




import unittest

from transformers import ErnieMConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

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


if is_torch_available():
       import torch

       from transformers import (
           ErnieMForInformationExtraction,
           ErnieMForMultipleChoice,
           ErnieMForQuestionAnswering,
           ErnieMForSequenceClassification,
           ErnieMForTokenClassification,
           ErnieMModel,
       )
       from transformers.models.ernie_m.modeling_ernie_m import ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST




class 	UpperCamelCase__			:

    '''simple docstring'''



    def __init__(					self  ,						UpperCamelCase__  ,						UpperCamelCase__=13  ,						UpperCamelCase__=7  ,						UpperCamelCase__=True  ,						UpperCamelCase__=True  ,						UpperCamelCase__=True  ,						UpperCamelCase__=99  ,						UpperCamelCase__=32  ,						UpperCamelCase__=5  ,						UpperCamelCase__=4  ,						UpperCamelCase__=37  ,						UpperCamelCase__="gelu"  ,						UpperCamelCase__=0.1  ,						UpperCamelCase__=0.1  ,						UpperCamelCase__=512  ,						UpperCamelCase__=16  ,						UpperCamelCase__=2  ,						UpperCamelCase__=0.02  ,						UpperCamelCase__=3  ,						UpperCamelCase__=4  ,						UpperCamelCase__=None  ,						)  ->   List[Any]:
           lowerCamelCase     :   Optional[Any]    				=      parent
           lowerCamelCase     :   Dict    				=      batch_size
           lowerCamelCase     :   Union[str, Any]    				=      seq_length
           lowerCamelCase     :   str    				=      is_training
           lowerCamelCase     :   Dict    				=      use_input_mask
           lowerCamelCase     :   List[str]    				=      use_labels
           lowerCamelCase     :   Optional[Any]    				=      vocab_size
           lowerCamelCase     :   List[str]    				=      hidden_size
           lowerCamelCase     :   int    				=      num_hidden_layers
           lowerCamelCase     :   Any    				=      num_attention_heads
           lowerCamelCase     :   Union[str, Any]    				=      intermediate_size
           lowerCamelCase     :   Any    				=      hidden_act
           lowerCamelCase     :   Any    				=      hidden_dropout_prob
           lowerCamelCase     :   Tuple    				=      attention_probs_dropout_prob
           lowerCamelCase     :   List[Any]    				=      max_position_embeddings
           lowerCamelCase     :   str    				=      type_vocab_size
           lowerCamelCase     :   Optional[int]    				=      type_sequence_label_size
           lowerCamelCase     :   Any    				=      initializer_range
           lowerCamelCase     :   Dict    				=      num_labels
           lowerCamelCase     :   Union[str, Any]    				=      num_choices
           lowerCamelCase     :   int    				=      scope







    def 					_lowercase  (					self  )  ->   str:
           lowerCamelCase     :   List[Any]    				=      ids_tensor([self.batch_size, self.seq_length]  ,						self.vocab_size  )

           lowerCamelCase     :   Optional[int]    				=      None
           if self.use_input_mask:
                  lowerCamelCase     :   str    				=      random_attention_mask([self.batch_size, self.seq_length]  )

           lowerCamelCase     :   Dict    				=      None
           lowerCamelCase     :   Any    				=      None
           lowerCamelCase     :   Optional[int]    				=      None
           if self.use_labels:
                  lowerCamelCase     :   Tuple    				=      ids_tensor([self.batch_size]  ,						self.type_sequence_label_size  )
                  lowerCamelCase     :   str    				=      ids_tensor([self.batch_size, self.seq_length]  ,						self.num_labels  )
                  lowerCamelCase     :   Any    				=      ids_tensor([self.batch_size]  ,						self.num_choices  )

           lowerCamelCase     :   int    				=      self.get_config()

           return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels







    def 					_lowercase  (					self  )  ->   int:
           lowerCamelCase     :   Tuple    				=      ids_tensor([self.batch_size, self.seq_length]  ,						self.vocab_size  )

           lowerCamelCase     :   List[Any]    				=      None
           if self.use_input_mask:
                  lowerCamelCase     :   List[str]    				=      random_attention_mask([self.batch_size, self.seq_length]  )
           lowerCamelCase     :   int    				=      self.get_config()

           return config, input_ids, input_mask







    def 					_lowercase  (					self  )  ->   int:
           return ErnieMConfig(
               vocab_size=self.vocab_size  ,						hidden_size=self.hidden_size  ,						num_hidden_layers=self.num_hidden_layers  ,						num_attention_heads=self.num_attention_heads  ,						intermediate_size=self.intermediate_size  ,						hidden_act=self.hidden_act  ,						hidden_dropout_prob=self.hidden_dropout_prob  ,						attention_probs_dropout_prob=self.attention_probs_dropout_prob  ,						max_position_embeddings=self.max_position_embeddings  ,						type_vocab_size=self.type_vocab_size  ,						initializer_range=self.initializer_range  ,						)







    def 					_lowercase  (					self  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  )  ->   Optional[int]:
           lowerCamelCase     :   List[Any]    				=      ErnieMModel(config=UpperCamelCase__  )
           model.to(UpperCamelCase__  )
           model.eval()
           lowerCamelCase     :   Optional[int]    				=      model(UpperCamelCase__  ,						return_dict=UpperCamelCase__  )
           self.parent.assertEqual(result.last_hidden_state.shape  ,						(self.batch_size, self.seq_length, self.hidden_size)  )







    def 					_lowercase  (					self  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  )  ->   Tuple:
           lowerCamelCase     :   List[str]    				=      ErnieMForQuestionAnswering(config=UpperCamelCase__  )
           model.to(UpperCamelCase__  )
           model.eval()
           lowerCamelCase     :   Optional[int]    				=      model(
               UpperCamelCase__  ,						attention_mask=UpperCamelCase__  ,						start_positions=UpperCamelCase__  ,						end_positions=UpperCamelCase__  ,						)
           self.parent.assertEqual(result.start_logits.shape  ,						(self.batch_size, self.seq_length)  )
           self.parent.assertEqual(result.end_logits.shape  ,						(self.batch_size, self.seq_length)  )







    def 					_lowercase  (					self  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  )  ->   List[str]:
           lowerCamelCase     :   Tuple    				=      ErnieMForInformationExtraction(config=UpperCamelCase__  )
           model.to(UpperCamelCase__  )
           model.eval()
           lowerCamelCase     :   int    				=      torch.ones_like(UpperCamelCase__  ,						dtype=torch.floataa  )
           lowerCamelCase     :   Dict    				=      model(
               UpperCamelCase__  ,						attention_mask=UpperCamelCase__  ,						start_positions=UpperCamelCase__  ,						end_positions=UpperCamelCase__  ,						)
           self.parent.assertEqual(result.start_logits.shape  ,						(self.batch_size, self.seq_length)  )
           self.parent.assertEqual(result.end_logits.shape  ,						(self.batch_size, self.seq_length)  )







    def 					_lowercase  (					self  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  )  ->   List[Any]:
           lowerCamelCase     :   List[Any]    				=      self.num_labels
           lowerCamelCase     :   Dict    				=      ErnieMForSequenceClassification(UpperCamelCase__  )
           model.to(UpperCamelCase__  )
           model.eval()
           lowerCamelCase     :   int    				=      model(UpperCamelCase__  ,						attention_mask=UpperCamelCase__  ,						labels=UpperCamelCase__  )
           self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.num_labels)  )







    def 					_lowercase  (					self  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  )  ->   Tuple:
           lowerCamelCase     :   int    				=      self.num_labels
           lowerCamelCase     :   Optional[Any]    				=      ErnieMForTokenClassification(config=UpperCamelCase__  )
           model.to(UpperCamelCase__  )
           model.eval()
           input_ids.to(UpperCamelCase__  )
           input_mask.to(UpperCamelCase__  )
           token_labels.to(UpperCamelCase__  )

           lowerCamelCase     :   List[Any]    				=      model(UpperCamelCase__  ,						attention_mask=UpperCamelCase__  ,						labels=UpperCamelCase__  )

           self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.seq_length, self.num_labels)  )







    def 					_lowercase  (					self  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  )  ->   List[Any]:
           lowerCamelCase     :   int    				=      self.num_choices
           lowerCamelCase     :   int    				=      ErnieMForMultipleChoice(config=UpperCamelCase__  )
           model.to(UpperCamelCase__  )
           model.eval()
           lowerCamelCase     :   Union[str, Any]    				=      input_ids.unsqueeze(1  ).expand(-1  ,						self.num_choices  ,						-1  ).contiguous()
           lowerCamelCase     :   List[Any]    				=      input_mask.unsqueeze(1  ).expand(-1  ,						self.num_choices  ,						-1  ).contiguous()
           lowerCamelCase     :   List[str]    				=      model(
               UpperCamelCase__  ,						attention_mask=UpperCamelCase__  ,						labels=UpperCamelCase__  ,						)
           self.parent.assertEqual(result.logits.shape  ,						(self.batch_size, self.num_choices)  )







    def 					_lowercase  (					self  )  ->   Optional[Any]:
           lowerCamelCase     :   Any    				=      self.prepare_config_and_inputs()
           (
               lowerCamelCase   
           )     :   int    				=      config_and_inputs
           lowerCamelCase     :   Dict    				=      {"input_ids": input_ids, "attention_mask": input_mask}
           return config, inputs_dict




@require_torch
class 	UpperCamelCase__			(lowerCAmelCase__					,       lowerCAmelCase__					,       unittest.TestCase  ):

    '''simple docstring'''



    lowerCamelCase_   :			str			  =   (
        (
            ErnieMModel,
            ErnieMForMultipleChoice,
            ErnieMForQuestionAnswering,
            ErnieMForSequenceClassification,
            ErnieMForTokenClassification,
        )
        if is_torch_available()
        else ()
    )
    lowerCamelCase_   :			int			  =   ()
    lowerCamelCase_   :			Tuple			  =   (
        {
            """feature-extraction""": ErnieMModel,
            """question-answering""": ErnieMForQuestionAnswering,
            """text-classification""": ErnieMForSequenceClassification,
            """token-classification""": ErnieMForTokenClassification,
            """zero-shot""": ErnieMForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    lowerCamelCase_   :			str			  =   False







    def 					_lowercase  (					self  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  )  ->   str:
           if pipeline_test_casse_name == "QAPipelineTests":
                  return True

           return False







    def 					_lowercase  (					self  )  ->   List[str]:
           lowerCamelCase     :   List[str]    				=      ErnieMModelTester(self  )
           lowerCamelCase     :   int    				=      ConfigTester(self  ,						config_class=UpperCamelCase__  ,						hidden_size=37  )







    def 					_lowercase  (					self  )  ->   Dict:
           self.config_tester.run_common_tests()







    def 					_lowercase  (					self  )  ->   List[str]:
           lowerCamelCase     :   List[str]    				=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*UpperCamelCase__  )







    def 					_lowercase  (					self  )  ->   Union[str, Any]:
           lowerCamelCase     :   Optional[Any]    				=      self.model_tester.prepare_config_and_inputs()
           for type in ["absolute", "relative_key", "relative_key_query"]:
                  lowerCamelCase     :   Union[str, Any]    				=      type
                  self.model_tester.create_and_check_model(*UpperCamelCase__  )







    def 					_lowercase  (					self  )  ->   str:
           lowerCamelCase     :   Any    				=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_multiple_choice(*UpperCamelCase__  )







    def 					_lowercase  (					self  )  ->   Union[str, Any]:
           lowerCamelCase     :   Tuple    				=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_question_answering(*UpperCamelCase__  )







    def 					_lowercase  (					self  )  ->   List[str]:
           lowerCamelCase     :   Dict    				=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_sequence_classification(*UpperCamelCase__  )







    def 					_lowercase  (					self  )  ->   int:
           lowerCamelCase     :   Any    				=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_information_extraction(*UpperCamelCase__  )







    def 					_lowercase  (					self  )  ->   List[str]:
           lowerCamelCase     :   str    				=      self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_token_classification(*UpperCamelCase__  )







    @slow
    def 					_lowercase  (					self  )  ->   str:
           for model_name in ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  lowerCamelCase     :   Tuple    				=      ErnieMModel.from_pretrained(UpperCamelCase__  )
                  self.assertIsNotNone(UpperCamelCase__  )




@require_torch
class 	UpperCamelCase__			(unittest.TestCase  ):

    '''simple docstring'''



    @slow
    def 					_lowercase  (					self  )  ->   Any:
           lowerCamelCase     :   int    				=      ErnieMModel.from_pretrained("susnato/ernie-m-base_pytorch"  )
           model.eval()
           lowerCamelCase     :   Union[str, Any]    				=      torch.tensor([[0, 1, 2, 3, 4, 5]]  )
           lowerCamelCase     :   List[Any]    				=      model(UpperCamelCase__  )[0]

           # TODO Replace vocab size
           lowerCamelCase     :   Union[str, Any]    				=      768

           lowerCamelCase     :   Dict    				=      torch.Size((1, 6, hidden_size)  )
           self.assertEqual(output.shape  ,						UpperCamelCase__  )

           lowerCamelCase     :   List[Any]    				=      torch.tensor(
               [[[-0.0012, 0.1245, -0.0214], [-0.0742, 0.0244, -0.0771], [-0.0333, 0.1164, -0.1554]]]  )

           self.assertTrue(torch.allclose(output[:, :3, :3]  ,						UpperCamelCase__  ,						atol=1e-3  )  )











import timeit
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_auto import MODEL_MAPPING, MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_torch_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_torch_available():
       import torch

       from .benchmark_args import PyTorchBenchmarkArguments


if is_pyanvml_available():
       import pyanvml.pyanvml as nvml


SCREAMING_SNAKE_CASE__ :   List[Any]						=     logging.get_logger(__name__)




class 	UpperCamelCase__			(lowerCAmelCase__  ):

    '''simple docstring'''



    lowerCamelCase_   :			PyTorchBenchmarkArguments
    lowerCamelCase_   :			PretrainedConfig
    lowerCamelCase_   :			str			  =   "PyTorch"







    @property
    def 					_lowercase  (					self  )  ->   Optional[int]:
           return torch.__version__







    def 					_lowercase  (					self  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  )  ->   float:
           lowerCamelCase     :   List[Any]    				=      self._prepare_inference_func(UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  )
           return self._measure_speed(_inference  )







    def 					_lowercase  (					self  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  )  ->   [Memory, Optional[MemorySummary]]:
           lowerCamelCase     :   List[Any]    				=      self._prepare_inference_func(UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  )
           return self._measure_memory(_inference  )







    def 					_lowercase  (					self  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  )  ->   float:
           lowerCamelCase     :   Dict    				=      self._prepare_train_func(UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  )
           return self._measure_speed(_train  )







    def 					_lowercase  (					self  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  )  ->   [Memory, Optional[MemorySummary]]:
           lowerCamelCase     :   Optional[Any]    				=      self._prepare_train_func(UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  )
           return self._measure_memory(_train  )







    def 					_lowercase  (					self  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  )  ->   Callable[[], None]:
           lowerCamelCase     :   Optional[Any]    				=      self.config_dict[model_name]

           if self.args.torchscript:
                  lowerCamelCase     :   Optional[Any]    				=      True

           lowerCamelCase     :   Dict    				=      (
               hasattr(UpperCamelCase__  ,						"architectures"  )
               and isinstance(config.architectures  ,						UpperCamelCase__  )
               and len(config.architectures  ) > 0
           )
           if not self.args.only_pretrain_model and has_model_class_in_config:
                  try:
                         lowerCamelCase     :   Optional[Any]    				=      config.architectures[0]
                         lowerCamelCase     :   Optional[int]    				=      __import__("transformers"  ,						fromlist=[model_class]  )
                         lowerCamelCase     :   Optional[int]    				=      getattr(UpperCamelCase__  ,						UpperCamelCase__  )
                         lowerCamelCase     :   List[str]    				=      model_cls(UpperCamelCase__  )
                  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:
                  lowerCamelCase     :   Tuple    				=      MODEL_MAPPING[config.__class__](UpperCamelCase__  )

           model.eval()
           model.to(self.args.device  )

           # encoder-decoder has vocab size saved differently
           lowerCamelCase     :   str    				=      config.vocab_size if hasattr(UpperCamelCase__  ,						"vocab_size"  ) else config.encoder.vocab_size
           lowerCamelCase     :   int    				=      torch.randint(UpperCamelCase__  ,						(batch_size, sequence_length)  ,						dtype=torch.long  ,						device=self.args.device  )

           if self.args.fpaa:
                  logger.info("Running training in Mixed Precision..."  )
                  if not self.args.is_gpu:
                         raise ValueError("Mixed precision is possible only for GPU."  )
                  # amp seems to have memory leaks so that memory usage
                  # is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
                  model.half()

           if self.args.torchscript:
                  with torch.no_grad():
                         lowerCamelCase     :   Optional[int]    				=      torch.jit.trace(UpperCamelCase__  ,						UpperCamelCase__  )
           else:
                  lowerCamelCase     :   Dict    				=      model

           def encoder_decoder_forward():
                  with torch.no_grad():
                         lowerCamelCase     :   Dict    				=      inference_model(UpperCamelCase__  ,						decoder_input_ids=UpperCamelCase__  )
                  return outputs

           def encoder_forward():
                  with torch.no_grad():
                         lowerCamelCase     :   Union[str, Any]    				=      inference_model(UpperCamelCase__  )
                  return outputs

           lowerCamelCase     :   Dict    				=      encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
           return _forward







    def 					_lowercase  (					self  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  )  ->   Callable[[], None]:
           lowerCamelCase     :   Optional[int]    				=      self.config_dict[model_name]

           lowerCamelCase     :   List[Any]    				=      (
               hasattr(UpperCamelCase__  ,						"architectures"  )
               and isinstance(config.architectures  ,						UpperCamelCase__  )
               and len(config.architectures  ) > 0
           )
           if not self.args.only_pretrain_model and has_model_class_in_config:
                  try:
                         lowerCamelCase     :   Union[str, Any]    				=      config.architectures[0]
                         lowerCamelCase     :   int    				=      __import__("transformers"  ,						fromlist=[model_class]  )
                         lowerCamelCase     :   Tuple    				=      getattr(UpperCamelCase__  ,						UpperCamelCase__  )
                         lowerCamelCase     :   Tuple    				=      model_cls(UpperCamelCase__  )
                  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:
                  lowerCamelCase     :   List[str]    				=      MODEL_WITH_LM_HEAD_MAPPING[config.__class__](UpperCamelCase__  )

           if self.args.torchscript:
                  raise NotImplementedError("Training for torchscript is currently not implemented"  )
           else:
                  lowerCamelCase     :   Dict    				=      model

           model.train()
           model.to(self.args.device  )

           # encoder-decoder has vocab size saved differently
           lowerCamelCase     :   Union[str, Any]    				=      config.vocab_size if hasattr(UpperCamelCase__  ,						"vocab_size"  ) else config.encoder.vocab_size
           lowerCamelCase     :   List[str]    				=      torch.randint(UpperCamelCase__  ,						(batch_size, sequence_length)  ,						dtype=torch.long  ,						device=self.args.device  )

           if self.args.fpaa:
                  logger.info("Running training in Mixed Precision..."  )
                  if not self.args.is_gpu:
                         raise ValueError("Mixed precision is possible only for GPU."  )

                  # amp seems to have memory leaks so that memory usage
                  # is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
                  model.half()

           def compute_loss_and_backprob_encoder():
                  lowerCamelCase     :   Tuple    				=      train_model(UpperCamelCase__  ,						labels=UpperCamelCase__  )[0]
                  loss.backward()
                  return loss

           def compute_loss_and_backprob_encoder_decoder():
                  lowerCamelCase     :   Optional[int]    				=      train_model(UpperCamelCase__  ,						decoder_input_ids=UpperCamelCase__  ,						labels=UpperCamelCase__  )[0]
                  loss.backward()
                  return loss

           lowerCamelCase     :   Union[str, Any]    				=      (
               compute_loss_and_backprob_encoder_decoder
               if config.is_encoder_decoder
               else compute_loss_and_backprob_encoder
           )
           return _train







    def 					_lowercase  (					self  ,						UpperCamelCase__  )  ->   float:
           try:
                  if self.args.is_tpu or self.args.torchscript:
                         # run additional 10 times to stabilize compilation for tpu and torchscript
                         logger.info("Do inference on TPU or torchscript. Running model 5 times to stabilize compilation"  )
                         timeit.repeat(
                             UpperCamelCase__  ,						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
                  lowerCamelCase     :   List[str]    				=      timeit.repeat(
                      UpperCamelCase__  ,						repeat=self.args.repeat  ,						number=10  ,						)

                  if self.args.is_tpu and self.args.torch_xla_tpu_print_metrics:
                         import torch_xla.debug.metrics as met

                         self.print_fn(met.metrics_report()  )

                  return min(UpperCamelCase__  ) / 10.0
           except RuntimeError as e:
                  self.print_fn(F'''Doesn\'t fit on GPU. {e}'''  )
                  return "N/A"







    def 					_lowercase  (					self  ,						UpperCamelCase__  )  ->   [Memory, MemorySummary]:
           try:
                  if self.args.trace_memory_line_by_line:
                         lowerCamelCase     :   str    				=      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"
                             " `--no-memory` or `args.memory=False`"  )
                  elif self.args.is_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."  )
                                lowerCamelCase     :   Optional[Any]    				=      "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()
                                lowerCamelCase     :   int    				=      nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx  )
                                lowerCamelCase     :   Dict    				=      nvml.nvmlDeviceGetMemoryInfo(UpperCamelCase__  )
                                lowerCamelCase     :   int    				=      meminfo.used
                                lowerCamelCase     :   Dict    				=      Memory(UpperCamelCase__  )
                                # shutdown nvml
                                nvml.nvmlShutdown()
                  else:
                         # cpu
                         lowerCamelCase     :   Union[str, Any]    				=      measure_peak_memory_cpu(UpperCamelCase__  )
                         lowerCamelCase     :   int    				=      Memory(UpperCamelCase__  ) if isinstance(UpperCamelCase__  ,						UpperCamelCase__  ) else memory_bytes

                  if self.args.trace_memory_line_by_line:
                         lowerCamelCase     :   List[str]    				=      stop_memory_tracing(UpperCamelCase__  )
                  else:
                         lowerCamelCase     :   int    				=      None

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






