






"""simple docstring"""



from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
						import tensorflow as tf

						from transformers import (
						    TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
						    TFXLMForMultipleChoice,
						    TFXLMForQuestionAnsweringSimple,
						    TFXLMForSequenceClassification,
						    TFXLMForTokenClassification,
						    TFXLMModel,
						    TFXLMWithLMHeadModel,
						    XLMConfig,
						)




class 							__lowercase:






						def __init__(       self     ,    _SCREAMING_SNAKE_CASE     ,    )      ->	Any:




											"""simple docstring"""





											lowerCamelCase__					=parent
											lowerCamelCase__					=13
											lowerCamelCase__					=7
											lowerCamelCase__					=True
											lowerCamelCase__					=True
											lowerCamelCase__					=True
											lowerCamelCase__					=True
											lowerCamelCase__					=True
											lowerCamelCase__					=False
											lowerCamelCase__					=False
											lowerCamelCase__					=False
											lowerCamelCase__					=2
											lowerCamelCase__					=99
											lowerCamelCase__					=0
											lowerCamelCase__					=32
											lowerCamelCase__					=2
											lowerCamelCase__					=4
											lowerCamelCase__					=0.1
											lowerCamelCase__					=0.1
											lowerCamelCase__					=5_12
											lowerCamelCase__					=16
											lowerCamelCase__					=2
											lowerCamelCase__					=0.02
											lowerCamelCase__					=3
											lowerCamelCase__					=4
											lowerCamelCase__					='last'
											lowerCamelCase__					=True
											lowerCamelCase__					=None
											lowerCamelCase__					=0






						def 		lowerCAmelCase (       self)      ->	List[Any]:




											"""simple docstring"""





											lowerCamelCase__					=ids_tensor([self.batch_size, self.seq_length]     ,    self.vocab_size)
											lowerCamelCase__					=random_attention_mask([self.batch_size, self.seq_length]     ,    dtype=tf.floataa)

											lowerCamelCase__					=None
											if self.use_input_lengths:
																lowerCamelCase__					=(
																    ids_tensor([self.batch_size]     ,    vocab_size=2) + self.seq_length - 2
																)  # small variation of seq_length

											lowerCamelCase__					=None
											if self.use_token_type_ids:
																lowerCamelCase__					=ids_tensor([self.batch_size, self.seq_length]     ,    self.n_langs)

											lowerCamelCase__					=None
											lowerCamelCase__					=None
											lowerCamelCase__					=None
											if self.use_labels:
																lowerCamelCase__					=ids_tensor([self.batch_size]     ,    self.type_sequence_label_size)
																lowerCamelCase__					=ids_tensor([self.batch_size, self.seq_length]     ,    self.num_labels)
																lowerCamelCase__					=ids_tensor([self.batch_size]     ,    2     ,    dtype=tf.floataa)
																lowerCamelCase__					=ids_tensor([self.batch_size]     ,    self.num_choices)

											lowerCamelCase__					=XLMConfig(
											    vocab_size=self.vocab_size     ,    n_special=self.n_special     ,    emb_dim=self.hidden_size     ,    n_layers=self.num_hidden_layers     ,    n_heads=self.num_attention_heads     ,    dropout=self.hidden_dropout_prob     ,    attention_dropout=self.attention_probs_dropout_prob     ,    gelu_activation=self.gelu_activation     ,    sinusoidal_embeddings=self.sinusoidal_embeddings     ,    asm=self.asm     ,    causal=self.causal     ,    n_langs=self.n_langs     ,    max_position_embeddings=self.max_position_embeddings     ,    initializer_range=self.initializer_range     ,    summary_type=self.summary_type     ,    use_proj=self.use_proj     ,    bos_token_id=self.bos_token_id     ,    )

											return (
											    config,
											    input_ids,
											    token_type_ids,
											    input_lengths,
											    sequence_labels,
											    token_labels,
											    is_impossible_labels,
											    choice_labels,
											    input_mask,
											)






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    )      ->	Union[str, Any]:




											"""simple docstring"""





											lowerCamelCase__					=TFXLMModel(config=_SCREAMING_SNAKE_CASE)
											lowerCamelCase__					={'input_ids': input_ids, 'lengths': input_lengths, 'langs': token_type_ids}
											lowerCamelCase__					=model(_SCREAMING_SNAKE_CASE)

											lowerCamelCase__					=[input_ids, input_mask]
											lowerCamelCase__					=model(_SCREAMING_SNAKE_CASE)
											self.parent.assertEqual(result.last_hidden_state.shape     ,    (self.batch_size, self.seq_length, self.hidden_size))






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    )      ->	int:




											"""simple docstring"""





											lowerCamelCase__					=TFXLMWithLMHeadModel(_SCREAMING_SNAKE_CASE)

											lowerCamelCase__					={'input_ids': input_ids, 'lengths': input_lengths, 'langs': token_type_ids}
											lowerCamelCase__					=model(_SCREAMING_SNAKE_CASE)

											lowerCamelCase__					=outputs

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






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    )      ->	Optional[int]:




											"""simple docstring"""





											lowerCamelCase__					=TFXLMForQuestionAnsweringSimple(_SCREAMING_SNAKE_CASE)

											lowerCamelCase__					={'input_ids': input_ids, 'lengths': input_lengths}

											lowerCamelCase__					=model(_SCREAMING_SNAKE_CASE)

											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 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    )      ->	Any:




											"""simple docstring"""





											lowerCamelCase__					=TFXLMForSequenceClassification(_SCREAMING_SNAKE_CASE)

											lowerCamelCase__					={'input_ids': input_ids, 'lengths': input_lengths}

											lowerCamelCase__					=model(_SCREAMING_SNAKE_CASE)

											self.parent.assertEqual(result.logits.shape     ,    (self.batch_size, self.type_sequence_label_size))






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    )      ->	str:




											"""simple docstring"""





											lowerCamelCase__					=self.num_labels
											lowerCamelCase__					=TFXLMForTokenClassification(config=_SCREAMING_SNAKE_CASE)
											lowerCamelCase__					={'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
											lowerCamelCase__					=model(_SCREAMING_SNAKE_CASE)
											self.parent.assertEqual(result.logits.shape     ,    (self.batch_size, self.seq_length, self.num_labels))






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    )      ->	Optional[Any]:




											"""simple docstring"""





											lowerCamelCase__					=self.num_choices
											lowerCamelCase__					=TFXLMForMultipleChoice(config=_SCREAMING_SNAKE_CASE)
											lowerCamelCase__					=tf.tile(tf.expand_dims(_SCREAMING_SNAKE_CASE     ,    1)     ,    (1, self.num_choices, 1))
											lowerCamelCase__					=tf.tile(tf.expand_dims(_SCREAMING_SNAKE_CASE     ,    1)     ,    (1, self.num_choices, 1))
											lowerCamelCase__					=tf.tile(tf.expand_dims(_SCREAMING_SNAKE_CASE     ,    1)     ,    (1, self.num_choices, 1))
											lowerCamelCase__					={
											    'input_ids': multiple_choice_inputs_ids,
											    'attention_mask': multiple_choice_input_mask,
											    'token_type_ids': multiple_choice_token_type_ids,
											}
											lowerCamelCase__					=model(_SCREAMING_SNAKE_CASE)
											self.parent.assertEqual(result.logits.shape     ,    (self.batch_size, self.num_choices))






						def 		lowerCAmelCase (       self)      ->	int:




											"""simple docstring"""





											lowerCamelCase__					=self.prepare_config_and_inputs()
											(
											    (
											    lowerCamelCase__
											)					,  (
											    lowerCamelCase__
											)					,  (
											    lowerCamelCase__
											)					,  (
											    lowerCamelCase__
											)					,  (
											    lowerCamelCase__
											)					,  (
											    lowerCamelCase__
											)					,  (
											    lowerCamelCase__
											)					,  (
											    lowerCamelCase__
											)					,  (
											    lowerCamelCase__
											)					,  
											)					=config_and_inputs
											lowerCamelCase__					={
											    'input_ids': input_ids,
											    'token_type_ids': token_type_ids,
											    'langs': token_type_ids,
											    'lengths': input_lengths,
											}
											return config, inputs_dict




@require_tf
class 							__lowercase(							_UpperCamelCase    ,				_UpperCamelCase    ,				unittest.TestCase	):
						lowerCAmelCase_		:				Union[str, Any]=							(
						    (
						        TFXLMModel,
						        TFXLMWithLMHeadModel,
						        TFXLMForSequenceClassification,
						        TFXLMForQuestionAnsweringSimple,
						        TFXLMForTokenClassification,
						        TFXLMForMultipleChoice,
						    )
						    if is_tf_available()
						    else ()
						)
						lowerCAmelCase_		:				Union[str, Any]=							(
						    (TFXLMWithLMHeadModel,) if is_tf_available() else ()
						)  # TODO (PVP): Check other models whether language generation is also applicable
						lowerCAmelCase_		:				Optional[Any]=							(
						    {
						        """feature-extraction""": TFXLMModel,
						        """fill-mask""": TFXLMWithLMHeadModel,
						        """question-answering""": TFXLMForQuestionAnsweringSimple,
						        """text-classification""": TFXLMForSequenceClassification,
						        """text-generation""": TFXLMWithLMHeadModel,
						        """token-classification""": TFXLMForTokenClassification,
						        """zero-shot""": TFXLMForSequenceClassification,
						    }
						    if is_tf_available()
						    else {}
						)
						lowerCAmelCase_		:				Tuple=							False
						lowerCAmelCase_		:				Dict=							False






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)      ->	Union[str, Any]:




											"""simple docstring"""





											if (
											    pipeline_test_casse_name == "QAPipelineTests"
											    and tokenizer_name is not None
											    and not tokenizer_name.endswith('Fast')
											):
																# `QAPipelineTests` fails for a few models when the slower tokenizer are used.
																# (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
																# TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
																return True

											return False






						def 		lowerCAmelCase (       self)      ->	int:




											"""simple docstring"""





											lowerCamelCase__					=TFXLMModelTester(self)
											lowerCamelCase__					=ConfigTester(self     ,    config_class=_SCREAMING_SNAKE_CASE     ,    emb_dim=37)






						def 		lowerCAmelCase (       self)      ->	Union[str, Any]:




											"""simple docstring"""





											self.config_tester.run_common_tests()






						def 		lowerCAmelCase (       self)      ->	int:




											"""simple docstring"""





											lowerCamelCase__					=self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_xlm_model(*_SCREAMING_SNAKE_CASE)






						def 		lowerCAmelCase (       self)      ->	List[str]:




											"""simple docstring"""





											lowerCamelCase__					=self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_xlm_lm_head(*_SCREAMING_SNAKE_CASE)






						def 		lowerCAmelCase (       self)      ->	int:




											"""simple docstring"""





											lowerCamelCase__					=self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_xlm_qa(*_SCREAMING_SNAKE_CASE)






						def 		lowerCAmelCase (       self)      ->	int:




											"""simple docstring"""





											lowerCamelCase__					=self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_xlm_sequence_classif(*_SCREAMING_SNAKE_CASE)






						def 		lowerCAmelCase (       self)      ->	Dict:




											"""simple docstring"""





											lowerCamelCase__					=self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_xlm_for_token_classification(*_SCREAMING_SNAKE_CASE)






						def 		lowerCAmelCase (       self)      ->	Tuple:




											"""simple docstring"""





											lowerCamelCase__					=self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_xlm_for_multiple_choice(*_SCREAMING_SNAKE_CASE)






						@slow
						def 		lowerCAmelCase (       self)      ->	int:




											"""simple docstring"""





											for model_name in TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																lowerCamelCase__					=TFXLMModel.from_pretrained(_SCREAMING_SNAKE_CASE)
																self.assertIsNotNone(_SCREAMING_SNAKE_CASE)




@require_tf
class 							__lowercase(							unittest.TestCase	):






						@slow
						def 		lowerCAmelCase (       self)      ->	Optional[Any]:




											"""simple docstring"""





											lowerCamelCase__					=TFXLMWithLMHeadModel.from_pretrained('xlm-mlm-en-2048')
											lowerCamelCase__					=tf.convert_to_tensor([[14, 4_47]]     ,    dtype=tf.intaa)  # the president
											lowerCamelCase__					=[
											    14,
											    4_47,
											    14,
											    4_47,
											    14,
											    4_47,
											    14,
											    4_47,
											    14,
											    4_47,
											    14,
											    4_47,
											    14,
											    4_47,
											    14,
											    4_47,
											    14,
											    4_47,
											    14,
											    4_47,
											]  # the president the president the president the president the president the president the president the president the president the president
											# TODO(PVP): this and other input_ids I tried for generation give pretty bad results. Not sure why. Model might just not be made for auto-regressive inference
											lowerCamelCase__					=model.generate(_SCREAMING_SNAKE_CASE     ,    do_sample=_SCREAMING_SNAKE_CASE)
											self.assertListEqual(output_ids[0].numpy().tolist()     ,    _SCREAMING_SNAKE_CASE)







"""simple docstring"""



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


_snake_case							:   str	  =       logging.get_logger(__name__)




class 							__lowercase(							_UpperCamelCase	):
						lowerCAmelCase_		:				PyTorchBenchmarkArguments
						lowerCAmelCase_		:				PretrainedConfig
						lowerCAmelCase_		:				str=							"PyTorch"






						@property
						def 		lowerCAmelCase (       self)      ->	List[Any]:




											"""simple docstring"""





											return torch.__version__






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)      ->	float:




											"""simple docstring"""





											lowerCamelCase__					=self._prepare_inference_func(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)
											return self._measure_speed(_inference)






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)      ->	[Memory, Optional[MemorySummary]]:




											"""simple docstring"""





											lowerCamelCase__					=self._prepare_inference_func(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)
											return self._measure_memory(_inference)






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)      ->	float:




											"""simple docstring"""





											lowerCamelCase__					=self._prepare_train_func(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)
											return self._measure_speed(_train)






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)      ->	[Memory, Optional[MemorySummary]]:




											"""simple docstring"""





											lowerCamelCase__					=self._prepare_train_func(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)
											return self._measure_memory(_train)






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)      ->	Callable[[], None]:




											"""simple docstring"""





											lowerCamelCase__					=self.config_dict[model_name]

											if self.args.torchscript:
																lowerCamelCase__					=True

											lowerCamelCase__					=(
											    hasattr(_SCREAMING_SNAKE_CASE     ,    'architectures')
											    and isinstance(config.architectures     ,    _SCREAMING_SNAKE_CASE)
											    and len(config.architectures) > 0
											)
											if not self.args.only_pretrain_model and has_model_class_in_config:
																try:
																					lowerCamelCase__					=config.architectures[0]
																					lowerCamelCase__					=__import__('transformers'     ,    fromlist=[model_class])
																					lowerCamelCase__					=getattr(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)
																					lowerCamelCase__					=model_cls(_SCREAMING_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:
																lowerCamelCase__					=MODEL_MAPPING[config.__class__](_SCREAMING_SNAKE_CASE)

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

											# encoder-decoder has vocab size saved differently
											lowerCamelCase__					=config.vocab_size if hasattr(_SCREAMING_SNAKE_CASE     ,    'vocab_size') else config.encoder.vocab_size
											lowerCamelCase__					=torch.randint(_SCREAMING_SNAKE_CASE     ,    (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__					=torch.jit.trace(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)
											else:
																lowerCamelCase__					=model

											def encoder_decoder_forward():
																with torch.no_grad():
																					lowerCamelCase__					=inference_model(_SCREAMING_SNAKE_CASE     ,    decoder_input_ids=_SCREAMING_SNAKE_CASE)
																return outputs

											def encoder_forward():
																with torch.no_grad():
																					lowerCamelCase__					=inference_model(_SCREAMING_SNAKE_CASE)
																return outputs

											lowerCamelCase__					=encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
											return _forward






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)      ->	Callable[[], None]:




											"""simple docstring"""





											lowerCamelCase__					=self.config_dict[model_name]

											lowerCamelCase__					=(
											    hasattr(_SCREAMING_SNAKE_CASE     ,    'architectures')
											    and isinstance(config.architectures     ,    _SCREAMING_SNAKE_CASE)
											    and len(config.architectures) > 0
											)
											if not self.args.only_pretrain_model and has_model_class_in_config:
																try:
																					lowerCamelCase__					=config.architectures[0]
																					lowerCamelCase__					=__import__('transformers'     ,    fromlist=[model_class])
																					lowerCamelCase__					=getattr(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE)
																					lowerCamelCase__					=model_cls(_SCREAMING_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:
																lowerCamelCase__					=MODEL_WITH_LM_HEAD_MAPPING[config.__class__](_SCREAMING_SNAKE_CASE)

											if self.args.torchscript:
																raise NotImplementedError('Training for torchscript is currently not implemented')
											else:
																lowerCamelCase__					=model

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

											# encoder-decoder has vocab size saved differently
											lowerCamelCase__					=config.vocab_size if hasattr(_SCREAMING_SNAKE_CASE     ,    'vocab_size') else config.encoder.vocab_size
											lowerCamelCase__					=torch.randint(_SCREAMING_SNAKE_CASE     ,    (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__					=train_model(_SCREAMING_SNAKE_CASE     ,    labels=_SCREAMING_SNAKE_CASE)[0]
																loss.backward()
																return loss

											def compute_loss_and_backprob_encoder_decoder():
																lowerCamelCase__					=train_model(_SCREAMING_SNAKE_CASE     ,    decoder_input_ids=_SCREAMING_SNAKE_CASE     ,    labels=_SCREAMING_SNAKE_CASE)[0]
																loss.backward()
																return loss

											lowerCamelCase__					=(
											    compute_loss_and_backprob_encoder_decoder
											    if config.is_encoder_decoder
											    else compute_loss_and_backprob_encoder
											)
											return _train






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE)      ->	float:




											"""simple docstring"""





											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(
																					    _SCREAMING_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
																lowerCamelCase__					=timeit.repeat(
																    _SCREAMING_SNAKE_CASE     ,    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(_SCREAMING_SNAKE_CASE) / 10.0
											except RuntimeError as e:
																self.print_fn(F"Doesn't fit on GPU. {e}")
																return "N/A"






						def 		lowerCAmelCase (       self     ,    _SCREAMING_SNAKE_CASE)      ->	[Memory, MemorySummary]:




											"""simple docstring"""





											try:
																if self.args.trace_memory_line_by_line:
																					lowerCamelCase__					=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__					='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__					=nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx)
																										lowerCamelCase__					=nvml.nvmlDeviceGetMemoryInfo(_SCREAMING_SNAKE_CASE)
																										lowerCamelCase__					=meminfo.used
																										lowerCamelCase__					=Memory(_SCREAMING_SNAKE_CASE)
																										# shutdown nvml
																										nvml.nvmlShutdown()
																else:
																					# cpu
																					lowerCamelCase__					=measure_peak_memory_cpu(_SCREAMING_SNAKE_CASE)
																					lowerCamelCase__					=Memory(_SCREAMING_SNAKE_CASE) if isinstance(_SCREAMING_SNAKE_CASE     ,    _SCREAMING_SNAKE_CASE) else memory_bytes

																if self.args.trace_memory_line_by_line:
																					lowerCamelCase__					=stop_memory_tracing(_SCREAMING_SNAKE_CASE)
																else:
																					lowerCamelCase__					=None

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