






import unittest

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

from ...generation.test_utils import GenerationTesterMixin
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 (
	    XLMForMultipleChoice,
	    XLMForQuestionAnswering,
	    XLMForQuestionAnsweringSimple,
	    XLMForSequenceClassification,
	    XLMForTokenClassification,
	    XLMModel,
	    XLMWithLMHeadModel,
	)
	from transformers.models.xlm.modeling_xlm import XLM_PRETRAINED_MODEL_ARCHIVE_LIST



class 							SCREAMING_SNAKE_CASE__						:





					"""simple docstring"""






					def __init__( self					:  str							,		__A					:  List[Any]							,		__A					:  List[str]=1_3							,		__A					:  Dict=7							,		__A					:  Dict=True							,		__A					:  Dict=True							,		__A					:  int=True							,		__A					:  Any=True							,		__A					:  Optional[Any]=True							,		__A					:  int=False							,		__A					:  Dict=False							,		__A					:  Dict=False							,		__A					:  Dict=2							,		__A					:  Dict=9_9							,		__A					:  Any=0							,		__A					:  Optional[Any]=3_2							,		__A					:  List[str]=5							,		__A					:  str=4							,		__A					:  int=0.1							,		__A					:  int=0.1							,		__A					:  str=5_1_2							,		__A					:  Tuple=2							,		__A					:  Tuple=0.0_2							,		__A					:  Union[str, Any]=2							,		__A					:  Optional[int]=4							,		__A					:  List[Any]="last"							,		__A					:  List[Any]=True							,		__A					:  Dict=None							,		__A					:  str=0							,		):
						snake_case__			:				int       =  parent
						snake_case__			:				str       =  batch_size
						snake_case__			:				int       =  seq_length
						snake_case__			:				Tuple       =  is_training
						snake_case__			:				Any       =  use_input_lengths
						snake_case__			:				Optional[Any]       =  use_token_type_ids
						snake_case__			:				Optional[int]       =  use_labels
						snake_case__			:				int       =  gelu_activation
						snake_case__			:				Tuple       =  sinusoidal_embeddings
						snake_case__			:				Dict       =  causal
						snake_case__			:				Optional[Any]       =  asm
						snake_case__			:				Any       =  n_langs
						snake_case__			:				Any       =  vocab_size
						snake_case__			:				Union[str, Any]       =  n_special
						snake_case__			:				Any       =  hidden_size
						snake_case__			:				Optional[Any]       =  num_hidden_layers
						snake_case__			:				List[Any]       =  num_attention_heads
						snake_case__			:				Optional[int]       =  hidden_dropout_prob
						snake_case__			:				Any       =  attention_probs_dropout_prob
						snake_case__			:				List[str]       =  max_position_embeddings
						snake_case__			:				Union[str, Any]       =  type_sequence_label_size
						snake_case__			:				Optional[Any]       =  initializer_range
						snake_case__			:				Union[str, Any]       =  num_labels
						snake_case__			:				List[Any]       =  num_choices
						snake_case__			:				Optional[int]       =  summary_type
						snake_case__			:				Optional[int]       =  use_proj
						snake_case__			:				Optional[int]       =  scope
						snake_case__			:				int       =  bos_token_id



					def 	_lowercase					( self					:  Optional[int]	):
						snake_case__			:				Any       =  ids_tensor([self.batch_size, self.seq_length]							,		self.vocab_size	)
						snake_case__			:				str       =  random_attention_mask([self.batch_size, self.seq_length]	)

						snake_case__			:				Union[str, Any]       =  None
						if self.use_input_lengths:
							snake_case__			:				List[Any]       =  (
							    ids_tensor([self.batch_size]							,		vocab_size=2	) + self.seq_length - 2
							)  # small variation of seq_length

						snake_case__			:				List[Any]       =  None
						if self.use_token_type_ids:
							snake_case__			:				int       =  ids_tensor([self.batch_size, self.seq_length]							,		self.n_langs	)

						snake_case__			:				int       =  None
						snake_case__			:				Dict       =  None
						snake_case__			:				Union[str, Any]       =  None
						if self.use_labels:
							snake_case__			:				str       =  ids_tensor([self.batch_size]							,		self.type_sequence_label_size	)
							snake_case__			:				int       =  ids_tensor([self.batch_size, self.seq_length]							,		self.num_labels	)
							snake_case__			:				Union[str, Any]       =  ids_tensor([self.batch_size]							,		2	).float()
							snake_case__			:				str       =  ids_tensor([self.batch_size]							,		self.num_choices	)

						snake_case__			:				Optional[Any]       =  self.get_config()

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



					def 	_lowercase					( self					:  List[Any]	):
						return 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							,		num_labels=self.num_labels							,		bos_token_id=self.bos_token_id							,		)



					def 	_lowercase					( self					:  str							,		__A					:  str							,		__A					:  Optional[Any]							,		__A					:  Any							,		__A					:  Any							,		__A					:  int							,		__A					:  Tuple							,		__A					:  int							,		__A					:  Optional[Any]							,		__A					:  Tuple							,		):
						snake_case__			:				Optional[int]       =  XLMModel(config=__A	)
						model.to(__A	)
						model.eval()
						snake_case__			:				List[str]       =  model(__A							,		lengths=__A							,		langs=__A	)
						snake_case__			:				Optional[int]       =  model(__A							,		langs=__A	)
						snake_case__			:				int       =  model(__A	)
						self.parent.assertEqual(result.last_hidden_state.shape							,		(self.batch_size, self.seq_length, self.hidden_size)	)



					def 	_lowercase					( self					:  str							,		__A					:  Dict							,		__A					:  Optional[int]							,		__A					:  int							,		__A					:  Optional[int]							,		__A					:  Optional[Any]							,		__A					:  List[Any]							,		__A					:  str							,		__A					:  Any							,		__A					:  Tuple							,		):
						snake_case__			:				List[str]       =  XLMWithLMHeadModel(__A	)
						model.to(__A	)
						model.eval()

						snake_case__			:				List[str]       =  model(__A							,		token_type_ids=__A							,		labels=__A	)
						self.parent.assertEqual(result.loss.shape							,		()	)
						self.parent.assertEqual(result.logits.shape							,		(self.batch_size, self.seq_length, self.vocab_size)	)



					def 	_lowercase					( self					:  int							,		__A					:  Optional[Any]							,		__A					:  Any							,		__A					:  List[str]							,		__A					:  Union[str, Any]							,		__A					:  str							,		__A					:  str							,		__A					:  int							,		__A					:  Dict							,		__A					:  Union[str, Any]							,		):
						snake_case__			:				Union[str, Any]       =  XLMForQuestionAnsweringSimple(__A	)
						model.to(__A	)
						model.eval()

						snake_case__			:				Any       =  model(__A	)

						snake_case__			:				int       =  model(__A							,		start_positions=__A							,		end_positions=__A	)
						snake_case__			:				int       =  outputs
						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					:  Optional[int]							,		__A					:  List[Any]							,		__A					:  int							,		__A					:  str							,		__A					:  str							,		__A					:  List[str]							,		__A					:  Optional[Any]							,		__A					:  Tuple							,		__A					:  List[Any]							,		__A					:  Dict							,		):
						snake_case__			:				Optional[Any]       =  XLMForQuestionAnswering(__A	)
						model.to(__A	)
						model.eval()

						snake_case__			:				Dict       =  model(__A	)

						snake_case__			:				Tuple       =  model(
						    __A							,		start_positions=__A							,		end_positions=__A							,		cls_index=__A							,		is_impossible=__A							,		p_mask=__A							,		)

						snake_case__			:				List[str]       =  model(
						    __A							,		start_positions=__A							,		end_positions=__A							,		cls_index=__A							,		is_impossible=__A							,		)

						((snake_case__),						)			:				Union[str, Any]       =  result_with_labels.to_tuple()

						snake_case__			:				Union[str, Any]       =  model(__A							,		start_positions=__A							,		end_positions=__A	)

						((snake_case__),						)			:				List[Any]       =  result_with_labels.to_tuple()

						self.parent.assertEqual(result_with_labels.loss.shape							,		()	)
						self.parent.assertEqual(result.start_top_log_probs.shape							,		(self.batch_size, model.config.start_n_top)	)
						self.parent.assertEqual(result.start_top_index.shape							,		(self.batch_size, model.config.start_n_top)	)
						self.parent.assertEqual(
						    result.end_top_log_probs.shape							,		(self.batch_size, model.config.start_n_top * model.config.end_n_top)	)
						self.parent.assertEqual(
						    result.end_top_index.shape							,		(self.batch_size, model.config.start_n_top * model.config.end_n_top)	)
						self.parent.assertEqual(result.cls_logits.shape							,		(self.batch_size,)	)



					def 	_lowercase					( self					:  int							,		__A					:  List[Any]							,		__A					:  Optional[Any]							,		__A					:  List[Any]							,		__A					:  Dict							,		__A					:  Tuple							,		__A					:  List[str]							,		__A					:  Tuple							,		__A					:  Any							,		__A					:  Optional[int]							,		):
						snake_case__			:				Union[str, Any]       =  XLMForSequenceClassification(__A	)
						model.to(__A	)
						model.eval()

						snake_case__			:				List[Any]       =  model(__A	)
						snake_case__			:				Any       =  model(__A							,		labels=__A	)
						self.parent.assertEqual(result.loss.shape							,		()	)
						self.parent.assertEqual(result.logits.shape							,		(self.batch_size, self.type_sequence_label_size)	)



					def 	_lowercase					( self					:  Optional[Any]							,		__A					:  Any							,		__A					:  Optional[Any]							,		__A					:  Union[str, Any]							,		__A					:  Dict							,		__A					:  Any							,		__A					:  List[Any]							,		__A					:  Dict							,		__A					:  List[str]							,		__A					:  str							,		):
						snake_case__			:				int       =  self.num_labels
						snake_case__			:				Any       =  XLMForTokenClassification(__A	)
						model.to(__A	)
						model.eval()

						snake_case__			:				str       =  model(__A							,		attention_mask=__A							,		labels=__A	)
						self.parent.assertEqual(result.logits.shape							,		(self.batch_size, self.seq_length, self.num_labels)	)



					def 	_lowercase					( self					:  Dict							,		__A					:  Optional[Any]							,		__A					:  Optional[int]							,		__A					:  Dict							,		__A					:  List[str]							,		__A					:  Dict							,		__A					:  Union[str, Any]							,		__A					:  Any							,		__A					:  str							,		__A					:  List[str]							,		):
						snake_case__			:				str       =  self.num_choices
						snake_case__			:				Optional[Any]       =  XLMForMultipleChoice(config=__A	)
						model.to(__A	)
						model.eval()
						snake_case__			:				Any       =  input_ids.unsqueeze(1	).expand(-1							,		self.num_choices							,		-1	).contiguous()
						snake_case__			:				str       =  token_type_ids.unsqueeze(1	).expand(-1							,		self.num_choices							,		-1	).contiguous()
						snake_case__			:				Optional[Any]       =  input_mask.unsqueeze(1	).expand(-1							,		self.num_choices							,		-1	).contiguous()
						snake_case__			:				Tuple       =  model(
						    __A							,		attention_mask=__A							,		token_type_ids=__A							,		labels=__A							,		)
						self.parent.assertEqual(result.logits.shape							,		(self.batch_size, self.num_choices)	)






					def 	_lowercase					( self					:  Union[str, Any]	):
						snake_case__			:				Any       =  self.prepare_config_and_inputs()
						(
						    (
						    snake_case__
						),						(
						    snake_case__
						),						(
						    snake_case__
						),						(
						    snake_case__
						),						(
						    snake_case__
						),						(
						    snake_case__
						),						(
						    snake_case__
						),						(
						    snake_case__
						),						(
						    snake_case__
						),						
						)			:				List[Any]       =  config_and_inputs
						snake_case__			:				Optional[Any]       =  {"input_ids": input_ids, "token_type_ids": token_type_ids, "lengths": input_lengths}
						return config, inputs_dict



@require_torch
class 							SCREAMING_SNAKE_CASE__						(    UpperCamelCase_			,       UpperCamelCase_			,       UpperCamelCase_			,       unittest.TestCase						):





					"""simple docstring"""






					a_     =		(
					    (
					        XLMModel,
					        XLMWithLMHeadModel,
					        XLMForQuestionAnswering,
					        XLMForSequenceClassification,
					        XLMForQuestionAnsweringSimple,
					        XLMForTokenClassification,
					        XLMForMultipleChoice,
					    )
					    if is_torch_available()
					    else ()
					)
					a_     =		(
					    (XLMWithLMHeadModel,) if is_torch_available() else ()
					)  # TODO (PVP): Check other models whether language generation is also applicable
					a_     =		(
					    {
					        "feature-extraction": XLMModel,
					        "fill-mask": XLMWithLMHeadModel,
					        "question-answering": XLMForQuestionAnsweringSimple,
					        "text-classification": XLMForSequenceClassification,
					        "text-generation": XLMWithLMHeadModel,
					        "token-classification": XLMForTokenClassification,
					        "zero-shot": XLMForSequenceClassification,
					    }
					    if is_torch_available()
					    else {}
					)



					def 	_lowercase					( self					:  List[Any]							,		__A					:  str							,		__A					:  Optional[int]							,		__A					:  Dict							,		__A					:  str							,		__A					:  List[Any]	):
						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 	_lowercase					( self					:  Any							,		__A					:  Dict							,		__A					:  Any							,		__A					:  Tuple=False	):
						snake_case__			:				List[Any]       =  super()._prepare_for_class(__A							,		__A							,		return_labels=__A	)

						if return_labels:
							if model_class.__name__ == "XLMForQuestionAnswering":
								snake_case__			:				Optional[Any]       =  torch.zeros(
								    self.model_tester.batch_size							,		dtype=torch.long							,		device=__A	)
								snake_case__			:				Union[str, Any]       =  torch.zeros(
								    self.model_tester.batch_size							,		dtype=torch.long							,		device=__A	)

						return inputs_dict



					def 	_lowercase					( self					:  Any	):
						snake_case__			:				Optional[int]       =  XLMModelTester(self	)
						snake_case__			:				Optional[Any]       =  ConfigTester(self							,		config_class=__A							,		emb_dim=3_7	)



					def 	_lowercase					( self					:  List[str]	):
						self.config_tester.run_common_tests()



					def 	_lowercase					( self					:  Union[str, Any]	):
						snake_case__			:				str       =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_model(*__A	)



					def 	_lowercase					( self					:  Union[str, Any]	):
						snake_case__			:				Any       =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_lm_head(*__A	)



					def 	_lowercase					( self					:  Tuple	):
						snake_case__			:				List[Any]       =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_simple_qa(*__A	)



					def 	_lowercase					( self					:  Tuple	):
						snake_case__			:				Dict       =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_qa(*__A	)



					def 	_lowercase					( self					:  int	):
						snake_case__			:				Union[str, Any]       =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_sequence_classif(*__A	)



					def 	_lowercase					( self					:  Dict	):
						snake_case__			:				Optional[int]       =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_token_classif(*__A	)



					def 	_lowercase					( self					:  Any	):
						snake_case__			:				List[Any]       =  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_for_multiple_choice(*__A	)



					def 	_lowercase					( self					:  Union[str, Any]							,		__A					:  Any							,		__A					:  List[str]							,		__A					:  str							,		__A					:  str							,		__A					:  str							,		__A					:  Tuple=False							,		__A					:  Optional[int]=1	):
						self.assertIsInstance(__A							,		__A	)
						self.assertListEqual(
						    [isinstance(__A							,		__A	) for iter_attentions in attentions]							,		[True] * len(__A	)	)
						self.assertEqual(len(__A	)							,		(max_length - min_length) * num_beam_groups	)

						for idx, iter_attentions in enumerate(__A	):
							# adds PAD dummy token
							snake_case__			:				List[str]       =  min_length + idx + 1
							snake_case__			:				Optional[int]       =  min_length + idx + 1

							snake_case__			:				Tuple       =  (
							    batch_size * num_beam_groups,
							    config.num_attention_heads,
							    tgt_len,
							    src_len,
							)
							# check attn size
							self.assertListEqual(
							    [layer_attention.shape for layer_attention in iter_attentions]							,		[expected_shape] * len(__A	)	)



					def 	_lowercase					( self					:  Tuple							,		__A					:  Any							,		__A					:  List[str]							,		__A					:  Any							,		__A					:  Tuple							,		__A					:  Optional[Any]							,		__A					:  str=False							,		__A					:  int=1	):
						self.assertIsInstance(__A							,		__A	)
						self.assertListEqual(
						    [isinstance(__A							,		__A	) for iter_hidden_states in hidden_states]							,		[True] * len(__A	)							,		)
						self.assertEqual(len(__A	)							,		(max_length - min_length) * num_beam_groups	)

						for idx, iter_hidden_states in enumerate(__A	):
							# adds PAD dummy token
							snake_case__			:				Union[str, Any]       =  min_length + idx + 1
							snake_case__			:				Tuple       =  (batch_size * num_beam_groups, seq_len, config.hidden_size)
							# check hidden size
							self.assertListEqual(
							    [layer_hidden_states.shape for layer_hidden_states in iter_hidden_states]							,		[expected_shape] * len(__A	)							,		)
						pass






					@slow
					def 	_lowercase					( self					:  int	):
						for model_name in XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							snake_case__			:				List[str]       =  XLMModel.from_pretrained(__A	)
							self.assertIsNotNone(__A	)



@require_torch
class 							SCREAMING_SNAKE_CASE__						(    unittest.TestCase						):





					"""simple docstring"""






					@slow
					def 	_lowercase					( self					:  List[Any]	):
						snake_case__			:				Optional[int]       =  XLMWithLMHeadModel.from_pretrained("xlm-mlm-en-2048"	)
						model.to(__A	)
						snake_case__			:				Union[str, Any]       =  torch.tensor([[1_4, 4_4_7]]							,		dtype=torch.long							,		device=__A	)  # the president
						snake_case__			:				int       =  [
						    1_4,
						    4_4_7,
						    1_4,
						    4_4_7,
						    1_4,
						    4_4_7,
						    1_4,
						    4_4_7,
						    1_4,
						    4_4_7,
						    1_4,
						    4_4_7,
						    1_4,
						    4_4_7,
						    1_4,
						    4_4_7,
						    1_4,
						    4_4_7,
						    1_4,
						    4_4_7,
						]  # 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
						snake_case__			:				Tuple       =  model.generate(__A							,		do_sample=__A	)
						self.assertListEqual(output_ids[0].cpu().numpy().tolist()							,		__A	)











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

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







def 					SCREAMING_SNAKE_CASE    (			snake_case_ :							bool					,  snake_case_ :							bool       ):
	def run_func(snake_case_ :							str       ):
		@wraps(snake_case_       )
		def run_in_eager_mode(*snake_case_ :							str					,  **snake_case_ :							Union[str, Any]       ):
			return func(*snake_case_					,  **snake_case_       )

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

		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 					SCREAMING_SNAKE_CASE    (			snake_case_ :							int					,  snake_case_ :							int					,  snake_case_ :							int       ):
	snake_case__			:				Dict       =  random.Random()
	snake_case__			:				List[Any]       =  [rng.randint(0					,  vocab_size - 1       ) for i in range(batch_size * sequence_length       )]
	return tf.constant(snake_case_					,  shape=(batch_size, sequence_length)					,  dtype=tf.intaa       )





class 							SCREAMING_SNAKE_CASE__						(    UpperCamelCase_						):





					"""simple docstring"""






					a_     =		42
					a_     =		42
					a_     =		"TensorFlow"



					@property
					def 	_lowercase					( self					:  List[str]	):
						return tf.__version__



					def 	_lowercase					( self					:  List[str]							,		__A					:  str							,		__A					:  int							,		__A					:  int	):
						# initialize GPU on separate process
						snake_case__			:				str       =  self.args.strategy
						if strategy is None:
							raise ValueError("A device strategy has to be initialized before using TensorFlow."	)
						snake_case__			:				Dict       =  self._prepare_inference_func(__A							,		__A							,		__A	)
						return self._measure_speed(_inference	)



					def 	_lowercase					( self					:  Tuple							,		__A					:  str							,		__A					:  int							,		__A					:  int	):
						snake_case__			:				Optional[int]       =  self.args.strategy
						if strategy is None:
							raise ValueError("A device strategy has to be initialized before using TensorFlow."	)
						snake_case__			:				Any       =  self._prepare_train_func(__A							,		__A							,		__A	)
						return self._measure_speed(_train	)



					def 	_lowercase					( self					:  List[Any]							,		__A					:  str							,		__A					:  int							,		__A					:  int	):
						# initialize GPU on separate process
						if self.args.is_gpu:
							tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]							,		__A	)
						snake_case__			:				List[str]       =  self.args.strategy
						if strategy is None:
							raise ValueError("A device strategy has to be initialized before using TensorFlow."	)
						snake_case__			:				Optional[Any]       =  self._prepare_inference_func(__A							,		__A							,		__A	)
						return self._measure_memory(_inference	)



					def 	_lowercase					( self					:  str							,		__A					:  str							,		__A					:  int							,		__A					:  int	):
						if self.args.is_gpu:
							tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]							,		__A	)
						snake_case__			:				List[Any]       =  self.args.strategy
						if strategy is None:
							raise ValueError("A device strategy has to be initialized before using TensorFlow."	)

						snake_case__			:				int       =  self._prepare_train_func(__A							,		__A							,		__A	)
						return self._measure_memory(_train	)



					def 	_lowercase					( self					:  Union[str, Any]							,		__A					:  str							,		__A					:  int							,		__A					:  int	):
						snake_case__			:				int       =  self.config_dict[model_name]

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

						snake_case__			:				Tuple       =  (
						    hasattr(__A							,		"architectures"	)
						    and isinstance(config.architectures							,		__A	)
						    and len(config.architectures	) > 0
						)
						if not self.args.only_pretrain_model and has_model_class_in_config:
							try:
								snake_case__			:				Dict       =  "TF" + config.architectures[0]  # prepend 'TF' for tensorflow model
								snake_case__			:				Union[str, Any]       =  __import__("transformers"							,		fromlist=[model_class]	)
								snake_case__			:				Any       =  getattr(__A							,		__A	)
								snake_case__			:				Dict       =  model_cls(__A	)
							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:
							snake_case__			:				Dict       =  TF_MODEL_MAPPING[config.__class__](__A	)

						# encoder-decoder has vocab size saved differently
						snake_case__			:				Optional[int]       =  config.vocab_size if hasattr(__A							,		"vocab_size"	) else config.encoder.vocab_size
						snake_case__			:				List[Any]       =  random_input_ids(__A							,		__A							,		__A	)

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

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

						snake_case__			:				Optional[int]       =  encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

						return _inference



					def 	_lowercase					( self					:  List[str]							,		__A					:  str							,		__A					:  int							,		__A					:  int	):
						snake_case__			:				Optional[Any]       =  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."	)

						snake_case__			:				Any       =  (
						    hasattr(__A							,		"architectures"	)
						    and isinstance(config.architectures							,		__A	)
						    and len(config.architectures	) > 0
						)
						if not self.args.only_pretrain_model and has_model_class_in_config:
							try:
								snake_case__			:				Dict       =  "TF" + config.architectures[0]  # prepend 'TF' for tensorflow model
								snake_case__			:				List[Any]       =  __import__("transformers"							,		fromlist=[model_class]	)
								snake_case__			:				Optional[int]       =  getattr(__A							,		__A	)
								snake_case__			:				str       =  model_cls(__A	)
							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:
							snake_case__			:				Union[str, Any]       =  TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](__A	)

						# encoder-decoder has vocab size saved differently
						snake_case__			:				Union[str, Any]       =  config.vocab_size if hasattr(__A							,		"vocab_size"	) else config.encoder.vocab_size
						snake_case__			:				List[str]       =  random_input_ids(__A							,		__A							,		__A	)

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

						@run_with_tf_optimizations(self.args.eager_mode							,		self.args.use_xla	)
						def encoder_train():
							snake_case__			:				Optional[Any]       =  model(__A							,		labels=__A							,		training=__A	)[0]
							snake_case__			:				Dict       =  tf.gradients(__A							,		model.trainable_variables	)
							return gradients

						snake_case__			:				int       =  encoder_decoder_train if config.is_encoder_decoder else encoder_train

						return _train



					def 	_lowercase					( self					:  int							,		__A					:  List[Any]	):
						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(__A							,		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
								snake_case__			:				Optional[Any]       =  timeit.repeat(
								    __A							,		repeat=self.args.repeat							,		number=1_0							,		)

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






					def 	_lowercase					( self					:  str							,		__A					:  Callable[[], None]	):
						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."	)
									snake_case__			:				Optional[int]       =  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."	)
										snake_case__			:				List[str]       =  "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()
										snake_case__			:				Any       =  nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx	)
										snake_case__			:				Optional[Any]       =  nvml.nvmlDeviceGetMemoryInfo(__A	)
										snake_case__			:				Optional[int]       =  meminfo.used
										snake_case__			:				Any       =  Memory(__A	)
										# 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."	)
										snake_case__			:				int       =  None
									else:
										snake_case__			:				Any       =  measure_peak_memory_cpu(__A	)
										snake_case__			:				Tuple       =  Memory(__A	) if isinstance(__A							,		__A	) else memory_bytes
								if self.args.trace_memory_line_by_line:
									snake_case__			:				Optional[int]       =  stop_memory_tracing(__A	)
									if memory is None:
										snake_case__			:				Dict       =  summary.total
								else:
									snake_case__			:				List[str]       =  None

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



