




import os
import pickle
import tempfile
import unittest
from typing import Callable, Optional

import numpy as np

from transformers import (
    BatchEncoding,
    BertTokenizer,
    BertTokenizerFast,
    PreTrainedTokenizer,
    PreTrainedTokenizerFast,
    TensorType,
    TokenSpan,
    is_tokenizers_available,
)
from transformers.models.gpta.tokenization_gpta import GPTaTokenizer
from transformers.testing_utils import CaptureStderr, require_flax, require_tf, require_tokenizers, require_torch, slow


if is_tokenizers_available():
       from tokenizers import Tokenizer
       from tokenizers.models import WordPiece




class 	UpperCamelCase__			(unittest.TestCase  ):

    '''simple docstring'''



    def 					_lowercase  (					self  ,						UpperCamelCase__  )  ->   Tuple:
           lowerCamelCase     :   str    				=      list(tokenizer_class.max_model_input_sizes.keys()  )
           for model_name in sa_models[:1]:
                  lowerCamelCase     :   Optional[int]    				=      tokenizer_class.from_pretrained(_lowerCamelCase  )
                  self.assertIsNotNone(_lowerCamelCase  )
                  self.assertIsInstance(_lowerCamelCase  ,						_lowerCamelCase  )
                  self.assertIsInstance(_lowerCamelCase  ,						_lowerCamelCase  )

                  for special_tok in tokenizer.all_special_tokens:
                         self.assertIsInstance(_lowerCamelCase  ,						_lowerCamelCase  )
                         lowerCamelCase     :   Any    				=      tokenizer.convert_tokens_to_ids(_lowerCamelCase  )
                         self.assertIsInstance(_lowerCamelCase  ,						_lowerCamelCase  )







    def 					_lowercase  (					self  ,						UpperCamelCase__  ,						UpperCamelCase__ = None  )  ->   Tuple:
           lowerCamelCase     :   Optional[Any]    				=      pickle.dumps(_lowerCamelCase  )
           self.assertIsNotNone(_lowerCamelCase  )

           lowerCamelCase     :   str    				=      pickle.loads(_lowerCamelCase  )

           # Ensure is_fast is correctly restored
           self.assertEqual(be_restored.is_fast  ,						be_original.is_fast  )

           # Ensure encodings are potentially correctly restored
           if be_original.is_fast:
                  self.assertIsNotNone(be_restored.encodings  )
           else:
                  self.assertIsNone(be_restored.encodings  )

           # Ensure the keys are the same
           for original_v, restored_v in zip(be_original.values()  ,						be_restored.values()  ):
                  if equal_op:
                         self.assertTrue(equal_op(_lowerCamelCase  ,						_lowerCamelCase  )  )
                  else:
                         self.assertEqual(_lowerCamelCase  ,						_lowerCamelCase  )







    @slow
    def 					_lowercase  (					self  )  ->   Optional[int]:
           self.check_tokenizer_from_pretrained(_lowerCamelCase  )







    def 					_lowercase  (					self  )  ->   List[str]:
           self.assertEqual(TensorType("tf"  )  ,						TensorType.TENSORFLOW  )
           self.assertEqual(TensorType("pt"  )  ,						TensorType.PYTORCH  )
           self.assertEqual(TensorType("np"  )  ,						TensorType.NUMPY  )







    @require_tokenizers
    def 					_lowercase  (					self  )  ->   str:
           import numpy as np

           lowerCamelCase     :   List[str]    				=      BertTokenizer.from_pretrained("bert-base-cased"  )
           lowerCamelCase     :   List[str]    				=      BertTokenizerFast.from_pretrained("bert-base-cased"  )

           # Python no tensor
           with self.subTest("BatchEncoding (Python, return_tensors=None)"  ):
                  self.assert_dump_and_restore(tokenizer_p("Small example to encode"  )  )

           with self.subTest("BatchEncoding (Python, return_tensors=NUMPY)"  ):
                  self.assert_dump_and_restore(
                      tokenizer_p("Small example to encode"  ,						return_tensors=TensorType.NUMPY  )  ,						np.array_equal  )

           with self.subTest("BatchEncoding (Rust, return_tensors=None)"  ):
                  self.assert_dump_and_restore(tokenizer_r("Small example to encode"  )  )

           with self.subTest("BatchEncoding (Rust, return_tensors=NUMPY)"  ):
                  self.assert_dump_and_restore(
                      tokenizer_r("Small example to encode"  ,						return_tensors=TensorType.NUMPY  )  ,						np.array_equal  )







    @require_tf
    @require_tokenizers
    def 					_lowercase  (					self  )  ->   List[str]:
           import tensorflow as tf

           def tf_array_equals(UpperCamelCase__  ,						UpperCamelCase__  ):
                  return tf.reduce_all(tf.equal(_lowerCamelCase  ,						_lowerCamelCase  )  )

           lowerCamelCase     :   Union[str, Any]    				=      BertTokenizer.from_pretrained("bert-base-cased"  )
           lowerCamelCase     :   str    				=      BertTokenizerFast.from_pretrained("bert-base-cased"  )

           with self.subTest("BatchEncoding (Python, return_tensors=TENSORFLOW)"  ):
                  self.assert_dump_and_restore(
                      tokenizer_p("Small example to encode"  ,						return_tensors=TensorType.TENSORFLOW  )  ,						_lowerCamelCase  )

           with self.subTest("BatchEncoding (Rust, return_tensors=TENSORFLOW)"  ):
                  self.assert_dump_and_restore(
                      tokenizer_r("Small example to encode"  ,						return_tensors=TensorType.TENSORFLOW  )  ,						_lowerCamelCase  )







    @require_torch
    @require_tokenizers
    def 					_lowercase  (					self  )  ->   Any:
           import torch

           lowerCamelCase     :   Any    				=      BertTokenizer.from_pretrained("bert-base-cased"  )
           lowerCamelCase     :   Optional[int]    				=      BertTokenizerFast.from_pretrained("bert-base-cased"  )

           with self.subTest("BatchEncoding (Python, return_tensors=PYTORCH)"  ):
                  self.assert_dump_and_restore(
                      tokenizer_p("Small example to encode"  ,						return_tensors=TensorType.PYTORCH  )  ,						torch.equal  )

           with self.subTest("BatchEncoding (Rust, return_tensors=PYTORCH)"  ):
                  self.assert_dump_and_restore(
                      tokenizer_r("Small example to encode"  ,						return_tensors=TensorType.PYTORCH  )  ,						torch.equal  )







    @require_tokenizers
    def 					_lowercase  (					self  )  ->   Tuple:
           lowerCamelCase     :   List[str]    				=      BertTokenizer.from_pretrained("bert-base-cased"  )
           lowerCamelCase     :   str    				=      BertTokenizerFast.from_pretrained("bert-base-cased"  )

           with self.subTest("Python Tokenizer"  ):
                  self.assertFalse(tokenizer_p("Small example to_encode"  ).is_fast  )

           with self.subTest("Rust Tokenizer"  ):
                  self.assertTrue(tokenizer_r("Small example to_encode"  ).is_fast  )







    @require_tokenizers
    def 					_lowercase  (					self  )  ->   Dict:
           lowerCamelCase     :   Optional[int]    				=      BertTokenizerFast.from_pretrained("bert-base-cased"  )
           lowerCamelCase     :   Union[str, Any]    				=      tokenizer_r(["Test", "\xad", "test"]  ,						is_split_into_words=_lowerCamelCase  )

           self.assertEqual(encoded.word_to_tokens(0  )  ,						TokenSpan(start=1  ,						end=2  )  )
           self.assertEqual(encoded.word_to_tokens(1  )  ,						_lowerCamelCase  )
           self.assertEqual(encoded.word_to_tokens(2  )  ,						TokenSpan(start=2  ,						end=3  )  )







    def 					_lowercase  (					self  )  ->   int:
           lowerCamelCase     :   Tuple    				=      BatchEncoding({"inputs": [[1, 2, 3], [4, 5, 6]], "labels": [0, 1]}  )
           lowerCamelCase     :   Any    				=      batch.convert_to_tensors(tensor_type="np"  )
           self.assertEqual(tensor_batch["inputs"].shape  ,						(2, 3)  )
           self.assertEqual(tensor_batch["labels"].shape  ,						(2,)  )
           # test converting the converted
           with CaptureStderr() as cs:
                  lowerCamelCase     :   Any    				=      batch.convert_to_tensors(tensor_type="np"  )
           self.assertFalse(len(cs.err  )  ,						msg=F'''should have no warning, but got {cs.err}'''  )

           lowerCamelCase     :   Any    				=      BatchEncoding({"inputs": [1, 2, 3], "labels": 0}  )
           lowerCamelCase     :   List[Any]    				=      batch.convert_to_tensors(tensor_type="np"  ,						prepend_batch_axis=_lowerCamelCase  )
           self.assertEqual(tensor_batch["inputs"].shape  ,						(1, 3)  )
           self.assertEqual(tensor_batch["labels"].shape  ,						(1,)  )







    @require_torch
    def 					_lowercase  (					self  )  ->   Optional[Any]:
           lowerCamelCase     :   Union[str, Any]    				=      BatchEncoding({"inputs": [[1, 2, 3], [4, 5, 6]], "labels": [0, 1]}  )
           lowerCamelCase     :   Optional[Any]    				=      batch.convert_to_tensors(tensor_type="pt"  )
           self.assertEqual(tensor_batch["inputs"].shape  ,						(2, 3)  )
           self.assertEqual(tensor_batch["labels"].shape  ,						(2,)  )
           # test converting the converted
           with CaptureStderr() as cs:
                  lowerCamelCase     :   Union[str, Any]    				=      batch.convert_to_tensors(tensor_type="pt"  )
           self.assertFalse(len(cs.err  )  ,						msg=F'''should have no warning, but got {cs.err}'''  )

           lowerCamelCase     :   int    				=      BatchEncoding({"inputs": [1, 2, 3], "labels": 0}  )
           lowerCamelCase     :   Union[str, Any]    				=      batch.convert_to_tensors(tensor_type="pt"  ,						prepend_batch_axis=_lowerCamelCase  )
           self.assertEqual(tensor_batch["inputs"].shape  ,						(1, 3)  )
           self.assertEqual(tensor_batch["labels"].shape  ,						(1,)  )







    @require_tf
    def 					_lowercase  (					self  )  ->   Optional[int]:
           lowerCamelCase     :   Any    				=      BatchEncoding({"inputs": [[1, 2, 3], [4, 5, 6]], "labels": [0, 1]}  )
           lowerCamelCase     :   Optional[Any]    				=      batch.convert_to_tensors(tensor_type="tf"  )
           self.assertEqual(tensor_batch["inputs"].shape  ,						(2, 3)  )
           self.assertEqual(tensor_batch["labels"].shape  ,						(2,)  )
           # test converting the converted
           with CaptureStderr() as cs:
                  lowerCamelCase     :   List[Any]    				=      batch.convert_to_tensors(tensor_type="tf"  )
           self.assertFalse(len(cs.err  )  ,						msg=F'''should have no warning, but got {cs.err}'''  )

           lowerCamelCase     :   List[str]    				=      BatchEncoding({"inputs": [1, 2, 3], "labels": 0}  )
           lowerCamelCase     :   Optional[int]    				=      batch.convert_to_tensors(tensor_type="tf"  ,						prepend_batch_axis=_lowerCamelCase  )
           self.assertEqual(tensor_batch["inputs"].shape  ,						(1, 3)  )
           self.assertEqual(tensor_batch["labels"].shape  ,						(1,)  )







    @require_flax
    def 					_lowercase  (					self  )  ->   Any:
           lowerCamelCase     :   str    				=      BatchEncoding({"inputs": [[1, 2, 3], [4, 5, 6]], "labels": [0, 1]}  )
           lowerCamelCase     :   int    				=      batch.convert_to_tensors(tensor_type="jax"  )
           self.assertEqual(tensor_batch["inputs"].shape  ,						(2, 3)  )
           self.assertEqual(tensor_batch["labels"].shape  ,						(2,)  )
           # test converting the converted
           with CaptureStderr() as cs:
                  lowerCamelCase     :   Any    				=      batch.convert_to_tensors(tensor_type="jax"  )
           self.assertFalse(len(cs.err  )  ,						msg=F'''should have no warning, but got {cs.err}'''  )

           lowerCamelCase     :   Optional[Any]    				=      BatchEncoding({"inputs": [1, 2, 3], "labels": 0}  )
           lowerCamelCase     :   List[str]    				=      batch.convert_to_tensors(tensor_type="jax"  ,						prepend_batch_axis=_lowerCamelCase  )
           self.assertEqual(tensor_batch["inputs"].shape  ,						(1, 3)  )
           self.assertEqual(tensor_batch["labels"].shape  ,						(1,)  )







    def 					_lowercase  (					self  )  ->   Optional[int]:
           lowerCamelCase     :   Optional[int]    				=      [{"""input_ids""": np.array([0, 1, 2]  )}, {"""input_ids""": np.array([0, 1, 2, 3]  )}]
           lowerCamelCase     :   Optional[Any]    				=      BertTokenizer.from_pretrained("bert-base-cased"  )

           lowerCamelCase     :   int    				=      tokenizer.pad(_lowerCamelCase  ,						padding=_lowerCamelCase  )
           self.assertTrue(isinstance(batch["input_ids"]  ,						np.ndarray  )  )
           self.assertEqual(batch["input_ids"].tolist()  ,						[[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]]  )
           lowerCamelCase     :   Any    				=      tokenizer.pad(_lowerCamelCase  ,						padding=_lowerCamelCase  ,						return_tensors="np"  )
           self.assertTrue(isinstance(batch["input_ids"]  ,						np.ndarray  )  )
           self.assertEqual(batch["input_ids"].tolist()  ,						[[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]]  )







    @require_torch
    def 					_lowercase  (					self  )  ->   Union[str, Any]:
           import torch

           lowerCamelCase     :   int    				=      [{"""input_ids""": torch.tensor([0, 1, 2]  )}, {"""input_ids""": torch.tensor([0, 1, 2, 3]  )}]
           lowerCamelCase     :   Any    				=      BertTokenizer.from_pretrained("bert-base-cased"  )

           lowerCamelCase     :   Any    				=      tokenizer.pad(_lowerCamelCase  ,						padding=_lowerCamelCase  )
           self.assertTrue(isinstance(batch["input_ids"]  ,						torch.Tensor  )  )
           self.assertEqual(batch["input_ids"].tolist()  ,						[[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]]  )
           lowerCamelCase     :   Any    				=      tokenizer.pad(_lowerCamelCase  ,						padding=_lowerCamelCase  ,						return_tensors="pt"  )
           self.assertTrue(isinstance(batch["input_ids"]  ,						torch.Tensor  )  )
           self.assertEqual(batch["input_ids"].tolist()  ,						[[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]]  )







    @require_tf
    def 					_lowercase  (					self  )  ->   Any:
           import tensorflow as tf

           lowerCamelCase     :   Dict    				=      [{"""input_ids""": tf.constant([0, 1, 2]  )}, {"""input_ids""": tf.constant([0, 1, 2, 3]  )}]
           lowerCamelCase     :   List[Any]    				=      BertTokenizer.from_pretrained("bert-base-cased"  )

           lowerCamelCase     :   Optional[Any]    				=      tokenizer.pad(_lowerCamelCase  ,						padding=_lowerCamelCase  )
           self.assertTrue(isinstance(batch["input_ids"]  ,						tf.Tensor  )  )
           self.assertEqual(batch["input_ids"].numpy().tolist()  ,						[[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]]  )
           lowerCamelCase     :   Union[str, Any]    				=      tokenizer.pad(_lowerCamelCase  ,						padding=_lowerCamelCase  ,						return_tensors="tf"  )
           self.assertTrue(isinstance(batch["input_ids"]  ,						tf.Tensor  )  )
           self.assertEqual(batch["input_ids"].numpy().tolist()  ,						[[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]]  )







    @require_tokenizers
    def 					_lowercase  (					self  )  ->   Any:
           lowerCamelCase     :   Optional[Any]    				=      Tokenizer(WordPiece(unk_token="[UNK]"  )  )
           PreTrainedTokenizerFast(tokenizer_object=_lowerCamelCase  )







    @require_tokenizers
    def 					_lowercase  (					self  )  ->   Optional[Any]:
           lowerCamelCase     :   Dict    				=      Tokenizer(WordPiece(unk_token="[UNK]"  )  )
           with tempfile.TemporaryDirectory() as tmpdirname:
                  bert_tokenizer.save(os.path.join(_lowerCamelCase  ,						"tokenizer.json"  )  )
                  PreTrainedTokenizerFast(tokenizer_file=os.path.join(_lowerCamelCase  ,						"tokenizer.json"  )  )











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






