




import inspect
import unittest
import warnings
from math import ceil, floor

from transformers import LevitConfig
from transformers.file_utils import cached_property, is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import (
        MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
        MODEL_MAPPING,
        LevitForImageClassification,
        LevitForImageClassificationWithTeacher,
        LevitModel,
    )
    from transformers.models.levit.modeling_levit import LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
    from PIL import Image

    from transformers import LevitImageProcessor







class __snake_case			(			a		):







   def lowerCamelCase						(      self		:						Any):

        """simple docstring"""


        UpperCAmelCase_          =							self.config_class(**self.inputs_dict)
        self.parent.assertTrue(hasattr(_snake_case					,						'''hidden_sizes'''))
        self.parent.assertTrue(hasattr(_snake_case					,						'''num_attention_heads'''))







class __snake_case			:







   def __init__(      self		:						Dict					,						_snake_case		:						Union[str, Any]					,						_snake_case		:						str=13					,						_snake_case		:						int=64					,						_snake_case		:						List[Any]=3					,						_snake_case		:						Optional[int]=3					,						_snake_case		:						List[Any]=2					,						_snake_case		:						Dict=1					,						_snake_case		:						Union[str, Any]=16					,						_snake_case		:						Any=[128, 256, 384]					,						_snake_case		:						Any=[4, 6, 8]					,						_snake_case		:						Optional[int]=[2, 3, 4]					,						_snake_case		:						List[Any]=[16, 16, 16]					,						_snake_case		:						Union[str, Any]=0					,						_snake_case		:						Optional[int]=[2, 2, 2]					,						_snake_case		:						Any=[2, 2, 2]					,						_snake_case		:						List[Any]=0.0_2					,						_snake_case		:						List[str]=True					,						_snake_case		:						List[Any]=True					,						_snake_case		:						Any=2					,						):

        """simple docstring"""


        UpperCAmelCase_          =							parent
        UpperCAmelCase_          =							batch_size
        UpperCAmelCase_          =							image_size
        UpperCAmelCase_          =							num_channels
        UpperCAmelCase_          =							kernel_size
        UpperCAmelCase_          =							stride
        UpperCAmelCase_          =							padding
        UpperCAmelCase_          =							hidden_sizes
        UpperCAmelCase_          =							num_attention_heads
        UpperCAmelCase_          =							depths
        UpperCAmelCase_          =							key_dim
        UpperCAmelCase_          =							drop_path_rate
        UpperCAmelCase_          =							patch_size
        UpperCAmelCase_          =							attention_ratio
        UpperCAmelCase_          =							mlp_ratio
        UpperCAmelCase_          =							initializer_range
        UpperCAmelCase_          =							[
            ['''Subsample''', key_dim[0], hidden_sizes[0] // key_dim[0], 4, 2, 2],
            ['''Subsample''', key_dim[0], hidden_sizes[1] // key_dim[0], 4, 2, 2],
        ]
        UpperCAmelCase_          =							is_training
        UpperCAmelCase_          =							use_labels
        UpperCAmelCase_          =							num_labels
        UpperCAmelCase_          =							initializer_range







   def lowerCamelCase						(      self		:						List[str]):

        """simple docstring"""


        UpperCAmelCase_          =							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

        UpperCAmelCase_          =							None
        if self.use_labels:
             UpperCAmelCase_          =							ids_tensor([self.batch_size]					,						self.num_labels)

        UpperCAmelCase_          =							self.get_config()
        return config, pixel_values, labels







   def lowerCamelCase						(      self		:						List[str]):

        """simple docstring"""


        return LevitConfig(
            image_size=self.image_size					,						num_channels=self.num_channels					,						kernel_size=self.kernel_size					,						stride=self.stride					,						padding=self.padding					,						patch_size=self.patch_size					,						hidden_sizes=self.hidden_sizes					,						num_attention_heads=self.num_attention_heads					,						depths=self.depths					,						key_dim=self.key_dim					,						drop_path_rate=self.drop_path_rate					,						mlp_ratio=self.mlp_ratio					,						attention_ratio=self.attention_ratio					,						initializer_range=self.initializer_range					,						down_ops=self.down_ops					,						)







   def lowerCamelCase						(      self		:						str					,						_snake_case		:						Optional[int]					,						_snake_case		:						Tuple					,						_snake_case		:						Any):

        """simple docstring"""


        UpperCAmelCase_          =							LevitModel(config=_snake_case)
        model.to(_snake_case)
        model.eval()
        UpperCAmelCase_          =							model(_snake_case)
        UpperCAmelCase_          =							(self.image_size, self.image_size)
        UpperCAmelCase_     ,       UpperCAmelCase_          =							image_size[0], image_size[1]
        for _ in range(4):
             UpperCAmelCase_          =							floor(((height + 2 * self.padding - self.kernel_size) / self.stride) + 1)
             UpperCAmelCase_          =							floor(((width + 2 * self.padding - self.kernel_size) / self.stride) + 1)
        self.parent.assertEqual(
            result.last_hidden_state.shape					,						(self.batch_size, ceil(height / 4) * ceil(width / 4), self.hidden_sizes[-1])					,						)







   def lowerCamelCase						(      self		:						Optional[Any]					,						_snake_case		:						Any					,						_snake_case		:						Optional[Any]					,						_snake_case		:						Union[str, Any]):

        """simple docstring"""


        UpperCAmelCase_          =							self.num_labels
        UpperCAmelCase_          =							LevitForImageClassification(_snake_case)
        model.to(_snake_case)
        model.eval()
        UpperCAmelCase_          =							model(_snake_case					,						labels=_snake_case)
        self.parent.assertEqual(result.logits.shape					,						(self.batch_size, self.num_labels))







   def lowerCamelCase						(      self		:						str):

        """simple docstring"""


        UpperCAmelCase_          =							self.prepare_config_and_inputs()
        UpperCAmelCase_     ,       UpperCAmelCase_     ,       UpperCAmelCase_          =							config_and_inputs
        UpperCAmelCase_          =							{'''pixel_values''': pixel_values}
        return config, inputs_dict







@require_torch
class __snake_case			(			a ,  a ,  unittest.TestCase		):
   UpperCAmelCase__  :    Optional[int]          =  (
       (LevitModel, LevitForImageClassification, LevitForImageClassificationWithTeacher)
       if is_torch_available()
       else ()
   )
   UpperCAmelCase__  :    Optional[int]          =  (
       {
           '''feature-extraction''': LevitModel,
           '''image-classification''': (LevitForImageClassification, LevitForImageClassificationWithTeacher),
       }
       if is_torch_available()
       else {}
   )

   UpperCAmelCase__  :    Tuple          =  False
   UpperCAmelCase__  :    str          =  False
   UpperCAmelCase__  :    Dict          =  False
   UpperCAmelCase__  :    Optional[Any]          =  False
   UpperCAmelCase__  :    Any          =  False







   def lowerCamelCase						(      self		:						Tuple):

        """simple docstring"""


        UpperCAmelCase_          =							LevitModelTester(self)
        UpperCAmelCase_          =							ConfigTester(self					,						config_class=_snake_case					,						has_text_modality=_snake_case					,						hidden_size=37)







   def lowerCamelCase						(      self		:						Any):

        """simple docstring"""


        self.create_and_test_config_common_properties()
        self.config_tester.create_and_test_config_to_json_string()
        self.config_tester.create_and_test_config_to_json_file()
        self.config_tester.create_and_test_config_from_and_save_pretrained()
        self.config_tester.create_and_test_config_with_num_labels()
        self.config_tester.check_config_can_be_init_without_params()
        self.config_tester.check_config_arguments_init()







   def lowerCamelCase						(      self		:						Tuple):

        """simple docstring"""


        return







   @unittest.skip(reason='''Levit does not use inputs_embeds''')
   def lowerCamelCase						(      self		:						Optional[int]):

        """simple docstring"""


        pass







   @unittest.skip(reason='''Levit does not support input and output embeddings''')
   def lowerCamelCase						(      self		:						Optional[int]):

        """simple docstring"""


        pass







   @unittest.skip(reason='''Levit does not output attentions''')
   def lowerCamelCase						(      self		:						Optional[int]):

        """simple docstring"""


        pass







   def lowerCamelCase						(      self		:						Any):

        """simple docstring"""


        UpperCAmelCase_     ,       UpperCAmelCase_          =							self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
             UpperCAmelCase_          =							model_class(_snake_case)
             UpperCAmelCase_          =							inspect.signature(model.forward)
             # signature.parameters is an OrderedDict => so arg_names order is deterministic
             UpperCAmelCase_          =							[*signature.parameters.keys()]

             UpperCAmelCase_          =							['''pixel_values''']
             self.assertListEqual(arg_names[:1]					,						_snake_case)







   def lowerCamelCase						(      self		:						str):

        """simple docstring"""


        def check_hidden_states_output(_snake_case		:						Optional[Any]					,						_snake_case		:						Optional[Any]					,						_snake_case		:						Optional[Any]):
             UpperCAmelCase_          =							model_class(_snake_case)
             model.to(_snake_case)
             model.eval()

             with torch.no_grad():
                  UpperCAmelCase_          =							model(**self._prepare_for_class(_snake_case					,						_snake_case))

             UpperCAmelCase_          =							outputs.hidden_states

             UpperCAmelCase_          =							len(self.model_tester.depths) + 1
             self.assertEqual(len(_snake_case)					,						_snake_case)

             UpperCAmelCase_          =							(self.model_tester.image_size, self.model_tester.image_size)
             UpperCAmelCase_     ,       UpperCAmelCase_          =							image_size[0], image_size[1]
             for _ in range(4):
                  UpperCAmelCase_          =							floor(
                      (
                          (height + 2 * self.model_tester.padding - self.model_tester.kernel_size)
                          / self.model_tester.stride
                      )
                      + 1)
                  UpperCAmelCase_          =							floor(
                      (
                          (width + 2 * self.model_tester.padding - self.model_tester.kernel_size)
                          / self.model_tester.stride
                      )
                      + 1)
             # verify the first hidden states (first block)
             self.assertListEqual(
                 list(hidden_states[0].shape[-2:])					,						[
                     height * width,
                     self.model_tester.hidden_sizes[0],
                 ]					,						)

        UpperCAmelCase_     ,       UpperCAmelCase_          =							self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
             UpperCAmelCase_          =							True
             check_hidden_states_output(_snake_case					,						_snake_case					,						_snake_case)

             # check that output_hidden_states also work using config
             del inputs_dict["output_hidden_states"]
             UpperCAmelCase_          =							True

             check_hidden_states_output(_snake_case					,						_snake_case					,						_snake_case)







   @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.''')
   def lowerCamelCase						(      self		:						Optional[int]):

        """simple docstring"""


        pass







   def lowerCamelCase						(      self		:						Any					,						_snake_case		:						Optional[int]					,						_snake_case		:						str					,						_snake_case		:						int=False):

        """simple docstring"""


        UpperCAmelCase_          =							super()._prepare_for_class(_snake_case					,						_snake_case					,						return_labels=_snake_case)

        if return_labels:
             if model_class.__name__ == "LevitForImageClassificationWithTeacher":
                  del inputs_dict["labels"]

        return inputs_dict







   def lowerCamelCase						(      self		:						Dict):

        """simple docstring"""


        UpperCAmelCase_          =							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*_snake_case)







   def lowerCamelCase						(      self		:						Tuple):

        """simple docstring"""


        UpperCAmelCase_          =							self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*_snake_case)







   def lowerCamelCase						(      self		:						List[str]):

        """simple docstring"""


        if not self.model_tester.is_training:
             return

        UpperCAmelCase_     ,       UpperCAmelCase_          =							self.model_tester.prepare_config_and_inputs_for_common()
        UpperCAmelCase_          =							True

        for model_class in self.all_model_classes:
             # LevitForImageClassificationWithTeacher supports inference-only
             if (
                 model_class in get_values(_snake_case)
                 or model_class.__name__ == "LevitForImageClassificationWithTeacher"
             ):
                  continue
             UpperCAmelCase_          =							model_class(_snake_case)
             model.to(_snake_case)
             model.train()
             UpperCAmelCase_          =							self._prepare_for_class(_snake_case					,						_snake_case					,						return_labels=_snake_case)
             UpperCAmelCase_          =							model(**_snake_case).loss
             loss.backward()







   def lowerCamelCase						(      self		:						Dict):

        """simple docstring"""


        UpperCAmelCase_     ,       UpperCAmelCase_          =							self.model_tester.prepare_config_and_inputs_for_common()
        if not self.model_tester.is_training:
             return

        UpperCAmelCase_          =							False
        UpperCAmelCase_          =							True

        for model_class in self.all_model_classes:
             if model_class in get_values(_snake_case) or not model_class.supports_gradient_checkpointing:
                  continue
             # LevitForImageClassificationWithTeacher supports inference-only
             if model_class.__name__ == "LevitForImageClassificationWithTeacher":
                  continue
             UpperCAmelCase_          =							model_class(_snake_case)
             model.gradient_checkpointing_enable()
             model.to(_snake_case)
             model.train()
             UpperCAmelCase_          =							self._prepare_for_class(_snake_case					,						_snake_case					,						return_labels=_snake_case)
             UpperCAmelCase_          =							model(**_snake_case).loss
             loss.backward()







   def lowerCamelCase						(      self		:						Dict):

        """simple docstring"""


        UpperCAmelCase_     ,       UpperCAmelCase_          =							self.model_tester.prepare_config_and_inputs_for_common()

        UpperCAmelCase_          =							[
            {'''title''': '''multi_label_classification''', '''num_labels''': 2, '''dtype''': torch.float},
            {'''title''': '''single_label_classification''', '''num_labels''': 1, '''dtype''': torch.long},
            {'''title''': '''regression''', '''num_labels''': 1, '''dtype''': torch.float},
        ]

        for model_class in self.all_model_classes:
             if (
                 model_class
                 not in [
                     *get_values(_snake_case),
                 ]
                 or model_class.__name__ == "LevitForImageClassificationWithTeacher"
             ):
                  continue

             for problem_type in problem_types:
                  with self.subTest(msg=F"""Testing {model_class} with {problem_type["title"]}"""):
                       UpperCAmelCase_          =							problem_type['''title''']
                       UpperCAmelCase_          =							problem_type['''num_labels''']

                       UpperCAmelCase_          =							model_class(_snake_case)
                       model.to(_snake_case)
                       model.train()

                       UpperCAmelCase_          =							self._prepare_for_class(_snake_case					,						_snake_case					,						return_labels=_snake_case)

                       if problem_type["num_labels"] > 1:
                            UpperCAmelCase_          =							inputs['''labels'''].unsqueeze(1).repeat(1					,						problem_type['''num_labels'''])

                       UpperCAmelCase_          =							inputs['''labels'''].to(problem_type['''dtype'''])

                       # This tests that we do not trigger the warning form PyTorch "Using a target size that is different
                       # to the input size. This will likely lead to incorrect results due to broadcasting. Please ensure
                       # they have the same size." which is a symptom something in wrong for the regression problem.
                       # See https://github.com/huggingface/transformers/issues/11780
                       with warnings.catch_warnings(record=_snake_case) as warning_list:
                            UpperCAmelCase_          =							model(**_snake_case).loss
                       for w in warning_list:
                            if "Using a target size that is different to the input size" in str(w.message):
                                 raise ValueError(
                                     F"""Something is going wrong in the regression problem: intercepted {w.message}""")

                       loss.backward()







   @slow
   def lowerCamelCase						(      self		:						List[str]):

        """simple docstring"""


        for model_name in LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             UpperCAmelCase_          =							LevitModel.from_pretrained(_snake_case)
             self.assertIsNotNone(_snake_case)







def    A						()   ->      Any:







     """simple docstring"""


     UpperCAmelCase_          =							Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''       )
     return image







@require_torch
@require_vision
class __snake_case			(			unittest.TestCase		):







   @cached_property
   def lowerCamelCase						(      self		:						Union[str, Any]):

        """simple docstring"""


        return LevitImageProcessor.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0])







   @slow
   def lowerCamelCase						(      self		:						Union[str, Any]):

        """simple docstring"""


        UpperCAmelCase_          =							LevitForImageClassificationWithTeacher.from_pretrained(LEVIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]).to(
            _snake_case)

        UpperCAmelCase_          =							self.default_image_processor
        UpperCAmelCase_          =							prepare_img()
        UpperCAmelCase_          =							image_processor(images=_snake_case					,						return_tensors='''pt''').to(_snake_case)

        # forward pass
        with torch.no_grad():
             UpperCAmelCase_          =							model(**_snake_case)

        # verify the logits
        UpperCAmelCase_          =							torch.Size((1, 1000))
        self.assertEqual(outputs.logits.shape					,						_snake_case)

        UpperCAmelCase_          =							torch.tensor([1.0_4_4_8, -0.3_7_4_5, -1.8_3_1_7]).to(_snake_case)

        self.assertTrue(torch.allclose(outputs.logits[0, :3]					,						_snake_case					,						atol=1e-4))












import random
import timeit
from functools import wraps
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_tf_auto import TF_MODEL_MAPPING, TF_MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_tf_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_tf_available():
    import tensorflow as tf
    from tensorflow.python.framework.errors_impl import ResourceExhaustedError

    from .benchmark_args_tf import TensorFlowBenchmarkArguments

if is_pyanvml_available():
    import pyanvml.pyanvml as nvml

snake_case_ : Tuple    		=       logging.get_logger(__name__)


def    A						(__A   :    bool   ,				__A   :    bool       )   ->      Optional[Any]:







     """simple docstring"""


     def run_func(__A   :    Optional[Any]       ):
          @wraps(__A       )
          def run_in_eager_mode(*__A   :    Dict   ,				**__A   :    List[Any]       ):
               return func(*__A   ,				**__A       )

          @wraps(__A       )
          @tf.function(experimental_compile=__A       )
          def run_in_graph_mode(*__A   :    Optional[Any]   ,				**__A   :    Any       ):
               return func(*__A   ,				**__A       )

          if do_eager_mode is True:
               if use_xla is not False:
                    raise ValueError(
                        '''Cannot run model in XLA, if `args.eager_mode` is set to `True`. Please set `args.eager_mode=False`.'''       )
               return run_in_eager_mode
          else:
               return run_in_graph_mode

     return run_func


def    A						(__A   :    int   ,				__A   :    int   ,				__A   :    int       )   ->      ["tf.Tensor"]:







     """simple docstring"""


     UpperCAmelCase_          =							random.Random()
     UpperCAmelCase_          =							[rng.randint(0   ,				vocab_size - 1       ) for i in range(batch_size * sequence_length       )]
     return tf.constant(__A   ,				shape=(batch_size, sequence_length)   ,				dtype=tf.intaa       )







class __snake_case			(			a		):
   UpperCAmelCase__  :    TensorFlowBenchmarkArguments
   UpperCAmelCase__  :    PretrainedConfig
   UpperCAmelCase__  :    str          =  "TensorFlow"







   @property
   def lowerCamelCase						(      self		:						List[str]):

        """simple docstring"""


        return tf.__version__







   def lowerCamelCase						(      self		:						Dict					,						_snake_case		:						str					,						_snake_case		:						int					,						_snake_case		:						int):

        """simple docstring"""


        UpperCAmelCase_          =							self.args.strategy
        if strategy is None:
             raise ValueError('''A device strategy has to be initialized before using TensorFlow.''')
        UpperCAmelCase_          =							self._prepare_inference_func(_snake_case					,						_snake_case					,						_snake_case)
        return self._measure_speed(_inference)







   def lowerCamelCase						(      self		:						Any					,						_snake_case		:						str					,						_snake_case		:						int					,						_snake_case		:						int):

        """simple docstring"""


        UpperCAmelCase_          =							self.args.strategy
        if strategy is None:
             raise ValueError('''A device strategy has to be initialized before using TensorFlow.''')
        UpperCAmelCase_          =							self._prepare_train_func(_snake_case					,						_snake_case					,						_snake_case)
        return self._measure_speed(_train)







   def lowerCamelCase						(      self		:						Any					,						_snake_case		:						str					,						_snake_case		:						int					,						_snake_case		:						int):

        """simple docstring"""


        if self.args.is_gpu:
             tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]					,						_snake_case)
        UpperCAmelCase_          =							self.args.strategy
        if strategy is None:
             raise ValueError('''A device strategy has to be initialized before using TensorFlow.''')
        UpperCAmelCase_          =							self._prepare_inference_func(_snake_case					,						_snake_case					,						_snake_case)
        return self._measure_memory(_inference)







   def lowerCamelCase						(      self		:						Optional[Any]					,						_snake_case		:						str					,						_snake_case		:						int					,						_snake_case		:						int):

        """simple docstring"""


        if self.args.is_gpu:
             tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]					,						_snake_case)
        UpperCAmelCase_          =							self.args.strategy
        if strategy is None:
             raise ValueError('''A device strategy has to be initialized before using TensorFlow.''')

        UpperCAmelCase_          =							self._prepare_train_func(_snake_case					,						_snake_case					,						_snake_case)
        return self._measure_memory(_train)







   def lowerCamelCase						(      self		:						Optional[int]					,						_snake_case		:						str					,						_snake_case		:						int					,						_snake_case		:						int):

        """simple docstring"""


        UpperCAmelCase_          =							self.config_dict[model_name]

        if self.args.fpaa:
             raise NotImplementedError('''Mixed precision is currently not supported.''')

        UpperCAmelCase_          =							(
            hasattr(_snake_case					,						'''architectures''')
            and isinstance(config.architectures					,						_snake_case)
            and len(config.architectures) > 0
        )
        if not self.args.only_pretrain_model and has_model_class_in_config:
             try:
                  UpperCAmelCase_          =							'''TF''' + config.architectures[0]  # prepend 'TF' for tensorflow model
                  UpperCAmelCase_          =							__import__('''transformers'''					,						fromlist=[model_class])
                  UpperCAmelCase_          =							getattr(_snake_case					,						_snake_case)
                  UpperCAmelCase_          =							model_cls(_snake_case)
             except ImportError:
                  raise ImportError(
                      F"""{model_class} does not exist. If you just want to test the pretrained model, you might want to"""
                      ''' set `--only_pretrain_model` or `args.only_pretrain_model=True`.''')
        else:
             UpperCAmelCase_          =							TF_MODEL_MAPPING[config.__class__](_snake_case)

        # encoder-decoder has vocab size saved differently
        UpperCAmelCase_          =							config.vocab_size if hasattr(_snake_case					,						'''vocab_size''') else config.encoder.vocab_size
        UpperCAmelCase_          =							random_input_ids(_snake_case					,						_snake_case					,						_snake_case)

        @run_with_tf_optimizations(self.args.eager_mode					,						self.args.use_xla)
        def encoder_decoder_forward():
             return model(_snake_case					,						decoder_input_ids=_snake_case					,						training=_snake_case)

        @run_with_tf_optimizations(self.args.eager_mode					,						self.args.use_xla)
        def encoder_forward():
             return model(_snake_case					,						training=_snake_case)

        UpperCAmelCase_          =							encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

        return _inference







   def lowerCamelCase						(      self		:						Optional[Any]					,						_snake_case		:						str					,						_snake_case		:						int					,						_snake_case		:						int):

        """simple docstring"""


        UpperCAmelCase_          =							self.config_dict[model_name]

        if self.args.eager_mode is not False:
             raise ValueError('''Training cannot be done in eager mode. Please make sure that `args.eager_mode = False`.''')

        if self.args.fpaa:
             raise NotImplementedError('''Mixed precision is currently not supported.''')

        UpperCAmelCase_          =							(
            hasattr(_snake_case					,						'''architectures''')
            and isinstance(config.architectures					,						_snake_case)
            and len(config.architectures) > 0
        )
        if not self.args.only_pretrain_model and has_model_class_in_config:
             try:
                  UpperCAmelCase_          =							'''TF''' + config.architectures[0]  # prepend 'TF' for tensorflow model
                  UpperCAmelCase_          =							__import__('''transformers'''					,						fromlist=[model_class])
                  UpperCAmelCase_          =							getattr(_snake_case					,						_snake_case)
                  UpperCAmelCase_          =							model_cls(_snake_case)
             except ImportError:
                  raise ImportError(
                      F"""{model_class} does not exist. If you just want to test the pretrained model, you might want to"""
                      ''' set `--only_pretrain_model` or `args.only_pretrain_model=True`.''')
        else:
             UpperCAmelCase_          =							TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](_snake_case)

        # encoder-decoder has vocab size saved differently
        UpperCAmelCase_          =							config.vocab_size if hasattr(_snake_case					,						'''vocab_size''') else config.encoder.vocab_size
        UpperCAmelCase_          =							random_input_ids(_snake_case					,						_snake_case					,						_snake_case)

        @run_with_tf_optimizations(self.args.eager_mode					,						self.args.use_xla)
        def encoder_decoder_train():
             UpperCAmelCase_          =							model(_snake_case					,						decoder_input_ids=_snake_case					,						labels=_snake_case					,						training=_snake_case)[0]
             UpperCAmelCase_          =							tf.gradients(_snake_case					,						model.trainable_variables)
             return gradients

        @run_with_tf_optimizations(self.args.eager_mode					,						self.args.use_xla)
        def encoder_train():
             UpperCAmelCase_          =							model(_snake_case					,						labels=_snake_case					,						training=_snake_case)[0]
             UpperCAmelCase_          =							tf.gradients(_snake_case					,						model.trainable_variables)
             return gradients

        UpperCAmelCase_          =							encoder_decoder_train if config.is_encoder_decoder else encoder_train

        return _train







   def lowerCamelCase						(      self		:						Any					,						_snake_case		:						Optional[Any]):

        """simple docstring"""


        with self.args.strategy.scope():
             try:
                  if self.args.is_tpu or self.args.use_xla:
                       # run additional 10 times to stabilize compilation for tpu
                       logger.info('''Do inference on TPU. Running model 5 times to stabilize compilation''')
                       timeit.repeat(_snake_case					,						repeat=1					,						number=5)

                  # as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
                  UpperCAmelCase_          =							timeit.repeat(
                      _snake_case					,						repeat=self.args.repeat					,						number=10					,						)

                  return min(_snake_case) / 1_0.0
             except ResourceExhaustedError as e:
                  self.print_fn(F"""Doesn't fit on GPU. {e}""")







   def lowerCamelCase						(      self		:						Dict					,						_snake_case		:						Callable[[], None]):

        """simple docstring"""


        logger.info(
            '''Note that TensorFlow allocates more memory than '''
            '''it might need to speed up computation. '''
            '''The memory reported here corresponds to the memory '''
            '''reported by `nvidia-smi`, which can vary depending '''
            '''on total available memory on the GPU that is used.''')
        with self.args.strategy.scope():
             try:
                  if self.args.trace_memory_line_by_line:
                       if not self.args.eager_mode:
                            raise ValueError(
                                '''`args.eager_mode` is set to `False`. Make sure to run model in eager mode to measure memory'''
                                ''' consumption line by line.''')
                       UpperCAmelCase_          =							start_memory_tracing('''transformers''')

                  if self.args.is_tpu:
                       # tpu
                       raise NotImplementedError(
                           '''Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking'''
                           ''' with `args.memory=False`''')
                  elif self.args.is_gpu:
                       # gpu
                       if not is_pyanvml_available():
                            logger.warning(
                                '''py3nvml not installed, we won\'t log GPU memory usage. '''
                                '''Install py3nvml (pip install py3nvml) to log information about GPU.''')
                            UpperCAmelCase_          =							'''N/A'''
                       else:
                            logger.info(
                                '''Measuring total GPU usage on GPU device. Make sure to not have additional processes'''
                                ''' running on the same GPU.''')
                            # init nvml
                            nvml.nvmlInit()
                            func()
                            UpperCAmelCase_          =							nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx)
                            UpperCAmelCase_          =							nvml.nvmlDeviceGetMemoryInfo(_snake_case)
                            UpperCAmelCase_          =							meminfo.used
                            UpperCAmelCase_          =							Memory(_snake_case)
                            # shutdown nvml
                            nvml.nvmlShutdown()
                  else:
                       # cpu
                       if self.args.trace_memory_line_by_line:
                            logger.info(
                                '''When enabling line by line tracing, the max peak memory for CPU is inaccurate in'''
                                ''' TensorFlow.''')
                            UpperCAmelCase_          =							None
                       else:
                            UpperCAmelCase_          =							measure_peak_memory_cpu(_snake_case)
                            UpperCAmelCase_          =							Memory(_snake_case) if isinstance(_snake_case					,						_snake_case) else memory_bytes
                  if self.args.trace_memory_line_by_line:
                       UpperCAmelCase_          =							stop_memory_tracing(_snake_case)
                       if memory is None:
                            UpperCAmelCase_          =							summary.total
                  else:
                       UpperCAmelCase_          =							None

                  return memory, summary
             except ResourceExhaustedError as e:
                  self.print_fn(F"""Doesn't fit on GPU. {e}""")
                  return "N/A", None






