




"""simple docstring"""

import unittest

from transformers import FunnelConfig, FunnelTokenizer, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

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


if is_torch_available():
		import torch

		from transformers import (
		    MODEL_FOR_PRETRAINING_MAPPING,
		    FunnelBaseModel,
		    FunnelForMaskedLM,
		    FunnelForMultipleChoice,
		    FunnelForPreTraining,
		    FunnelForQuestionAnswering,
		    FunnelForSequenceClassification,
		    FunnelForTokenClassification,
		    FunnelModel,
		)
class     A__      :




	def __init__(   self  :	List[Any]      ,    lowercase  :	Optional[Any]      ,    lowercase  :	Optional[int]=13      ,    lowercase  :	Dict=7      ,    lowercase  :	Union[str, Any]=True      ,    lowercase  :	List[str]=True      ,    lowercase  :	int=True      ,    lowercase  :	Tuple=True      ,    lowercase  :	Dict=99      ,    lowercase  :	Any=[1, 1, 2]      ,    lowercase  :	Dict=1      ,    lowercase  :	List[Any]=32      ,    lowercase  :	Optional[int]=4      ,    lowercase  :	Union[str, Any]=8      ,    lowercase  :	Tuple=37      ,    lowercase  :	List[Any]="gelu_new"      ,    lowercase  :	Union[str, Any]=0.1      ,    lowercase  :	Optional[Any]=0.1      ,    lowercase  :	Dict=0.0      ,    lowercase  :	int=512      ,    lowercase  :	Tuple=3      ,    lowercase  :	Any=0.02      ,    lowercase  :	Optional[Any]=3      ,    lowercase  :	List[str]=4      ,    lowercase  :	str=None      ,    lowercase  :	Any=False      ,    )					-> List[str]:
						__snake_case          =	parent
						__snake_case          =	batch_size
						__snake_case          =	seq_length
						__snake_case          =	is_training
						__snake_case          =	use_input_mask
						__snake_case          =	use_token_type_ids
						__snake_case          =	use_labels
						__snake_case          =	vocab_size
						__snake_case          =	block_sizes
						__snake_case          =	num_decoder_layers
						__snake_case          =	d_model
						__snake_case          =	n_head
						__snake_case          =	d_head
						__snake_case          =	d_inner
						__snake_case          =	hidden_act
						__snake_case          =	hidden_dropout
						__snake_case          =	attention_dropout
						__snake_case          =	activation_dropout
						__snake_case          =	max_position_embeddings
						__snake_case          =	type_vocab_size
						__snake_case          =	2
						__snake_case          =	num_labels
						__snake_case          =	num_choices
						__snake_case          =	scope
						__snake_case          =	initializer_std

						# Used in the tests to check the size of the first attention layer
						__snake_case          =	n_head
						# Used in the tests to check the size of the first hidden state
						__snake_case          =	self.d_model
						# Used in the tests to check the number of output hidden states/attentions
						__snake_case          =	sum(self.block_sizes	) + (0 if base else self.num_decoder_layers)
						# FunnelModel adds two hidden layers: input embeddings and the sum of the upsampled encoder hidden state with
						# the last hidden state of the first block (which is the first hidden state of the decoder).
						if not base:
											__snake_case          =	self.num_hidden_layers + 2




	def 							a_		(   self  :	Optional[Any]	)					-> List[str]:
						__snake_case          =	ids_tensor([self.batch_size, self.seq_length]      ,    self.vocab_size	)

						__snake_case          =	None
						if self.use_input_mask:
											__snake_case          =	ids_tensor([self.batch_size, self.seq_length]      ,    vocab_size=2	)

						__snake_case          =	None
						if self.use_token_type_ids:
											__snake_case          =	ids_tensor([self.batch_size, self.seq_length]      ,    self.type_vocab_size	)

						__snake_case          =	None
						__snake_case          =	None
						__snake_case          =	None
						if self.use_labels:
											__snake_case          =	ids_tensor([self.batch_size]      ,    self.type_sequence_label_size	)
											__snake_case          =	ids_tensor([self.batch_size, self.seq_length]      ,    self.num_labels	)
											__snake_case          =	ids_tensor([self.batch_size]      ,    self.num_choices	)
											__snake_case          =	ids_tensor([self.batch_size, self.seq_length]      ,    1	)

						__snake_case          =	self.get_config()

						return (
						    config,
						    input_ids,
						    token_type_ids,
						    input_mask,
						    sequence_labels,
						    token_labels,
						    choice_labels,
						    fake_token_labels,
						)




	def 							a_		(   self  :	List[str]	)					-> Optional[Any]:
						return FunnelConfig(
						    vocab_size=self.vocab_size      ,    block_sizes=self.block_sizes      ,    num_decoder_layers=self.num_decoder_layers      ,    d_model=self.d_model      ,    n_head=self.n_head      ,    d_head=self.d_head      ,    d_inner=self.d_inner      ,    hidden_act=self.hidden_act      ,    hidden_dropout=self.hidden_dropout      ,    attention_dropout=self.attention_dropout      ,    activation_dropout=self.activation_dropout      ,    max_position_embeddings=self.max_position_embeddings      ,    type_vocab_size=self.type_vocab_size      ,    initializer_std=self.initializer_std      ,    )




	def 							a_		(   self  :	str      ,    lowercase  :	List[Any]      ,    lowercase  :	str      ,    lowercase  :	Dict      ,    lowercase  :	List[str]      ,    lowercase  :	Any      ,    lowercase  :	Any      ,    lowercase  :	Tuple      ,    lowercase  :	Optional[int]      ,    )					-> Union[str, Any]:
						__snake_case          =	FunnelModel(config=lowercase	)
						model.to(lowercase	)
						model.eval()
						__snake_case          =	model(lowercase      ,    attention_mask=lowercase      ,    token_type_ids=lowercase	)
						__snake_case          =	model(lowercase      ,    token_type_ids=lowercase	)
						__snake_case          =	model(lowercase	)
						self.parent.assertEqual(result.last_hidden_state.shape      ,    (self.batch_size, self.seq_length, self.d_model)	)

						__snake_case          =	False
						__snake_case          =	model(lowercase	)
						self.parent.assertEqual(result.last_hidden_state.shape      ,    (self.batch_size, self.seq_length, self.d_model)	)

						__snake_case          =	False
						__snake_case          =	model(lowercase	)
						self.parent.assertEqual(result.last_hidden_state.shape      ,    (self.batch_size, self.seq_length, self.d_model)	)




	def 							a_		(   self  :	Optional[int]      ,    lowercase  :	str      ,    lowercase  :	List[str]      ,    lowercase  :	Union[str, Any]      ,    lowercase  :	Optional[Any]      ,    lowercase  :	List[Any]      ,    lowercase  :	Tuple      ,    lowercase  :	List[Any]      ,    lowercase  :	Tuple      ,    )					-> Union[str, Any]:
						__snake_case          =	FunnelBaseModel(config=lowercase	)
						model.to(lowercase	)
						model.eval()
						__snake_case          =	model(lowercase      ,    attention_mask=lowercase      ,    token_type_ids=lowercase	)
						__snake_case          =	model(lowercase      ,    token_type_ids=lowercase	)
						__snake_case          =	model(lowercase	)
						self.parent.assertEqual(result.last_hidden_state.shape      ,    (self.batch_size, 2, self.d_model)	)

						__snake_case          =	False
						__snake_case          =	model(lowercase	)
						self.parent.assertEqual(result.last_hidden_state.shape      ,    (self.batch_size, 3, self.d_model)	)

						__snake_case          =	False
						__snake_case          =	model(lowercase	)
						self.parent.assertEqual(result.last_hidden_state.shape      ,    (self.batch_size, 2, self.d_model)	)




	def 							a_		(   self  :	Tuple      ,    lowercase  :	List[str]      ,    lowercase  :	List[str]      ,    lowercase  :	Optional[int]      ,    lowercase  :	Optional[int]      ,    lowercase  :	List[str]      ,    lowercase  :	Tuple      ,    lowercase  :	Any      ,    lowercase  :	Optional[Any]      ,    )					-> Tuple:
						__snake_case          =	self.num_labels
						__snake_case          =	FunnelForPreTraining(config=lowercase	)
						model.to(lowercase	)
						model.eval()
						__snake_case          =	model(lowercase      ,    attention_mask=lowercase      ,    token_type_ids=lowercase      ,    labels=lowercase	)
						self.parent.assertEqual(result.logits.shape      ,    (self.batch_size, self.seq_length)	)




	def 							a_		(   self  :	str      ,    lowercase  :	Any      ,    lowercase  :	Union[str, Any]      ,    lowercase  :	Union[str, Any]      ,    lowercase  :	str      ,    lowercase  :	List[str]      ,    lowercase  :	Optional[Any]      ,    lowercase  :	Optional[Any]      ,    lowercase  :	Optional[int]      ,    )					-> Tuple:
						__snake_case          =	FunnelForMaskedLM(config=lowercase	)
						model.to(lowercase	)
						model.eval()
						__snake_case          =	model(lowercase      ,    attention_mask=lowercase      ,    token_type_ids=lowercase      ,    labels=lowercase	)
						self.parent.assertEqual(result.logits.shape      ,    (self.batch_size, self.seq_length, self.vocab_size)	)




	def 							a_		(   self  :	Optional[int]      ,    lowercase  :	str      ,    lowercase  :	Any      ,    lowercase  :	Any      ,    lowercase  :	List[Any]      ,    lowercase  :	int      ,    lowercase  :	Union[str, Any]      ,    lowercase  :	Optional[int]      ,    lowercase  :	str      ,    )					-> int:
						__snake_case          =	self.num_labels
						__snake_case          =	FunnelForSequenceClassification(lowercase	)
						model.to(lowercase	)
						model.eval()
						__snake_case          =	model(lowercase      ,    attention_mask=lowercase      ,    token_type_ids=lowercase      ,    labels=lowercase	)
						self.parent.assertEqual(result.logits.shape      ,    (self.batch_size, self.num_labels)	)




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




	def 							a_		(   self  :	Tuple      ,    lowercase  :	Dict      ,    lowercase  :	Optional[int]      ,    lowercase  :	Tuple      ,    lowercase  :	Dict      ,    lowercase  :	Optional[Any]      ,    lowercase  :	Optional[Any]      ,    lowercase  :	Optional[int]      ,    lowercase  :	int      ,    )					-> int:
						__snake_case          =	self.num_labels
						__snake_case          =	FunnelForTokenClassification(config=lowercase	)
						model.to(lowercase	)
						model.eval()
						__snake_case          =	model(lowercase      ,    attention_mask=lowercase      ,    token_type_ids=lowercase      ,    labels=lowercase	)
						self.parent.assertEqual(result.logits.shape      ,    (self.batch_size, self.seq_length, self.num_labels)	)




	def 							a_		(   self  :	Optional[Any]      ,    lowercase  :	Tuple      ,    lowercase  :	List[str]      ,    lowercase  :	Union[str, Any]      ,    lowercase  :	str      ,    lowercase  :	List[Any]      ,    lowercase  :	Union[str, Any]      ,    lowercase  :	Union[str, Any]      ,    lowercase  :	str      ,    )					-> Union[str, Any]:
						__snake_case          =	FunnelForQuestionAnswering(config=lowercase	)
						model.to(lowercase	)
						model.eval()
						__snake_case          =	model(
						    lowercase      ,    attention_mask=lowercase      ,    token_type_ids=lowercase      ,    start_positions=lowercase      ,    end_positions=lowercase      ,    )
						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 							a_		(   self  :	Any	)					-> Optional[int]:
						__snake_case          =	self.prepare_config_and_inputs()
						(
						    (
						    __snake_case
						)      ,      (
						    __snake_case
						)      ,      (
						    __snake_case
						)      ,      (
						    __snake_case
						)      ,      (
						    __snake_case
						)      ,      (
						    __snake_case
						)      ,      (
						    __snake_case
						)      ,      (
						    __snake_case
						)      ,      
						)          =	config_and_inputs
						__snake_case          =	{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
						return config, inputs_dict

@require_torch
class     A__      (      lowercase_     ,	lowercase_     ,	unittest.TestCase     ):
	A_							:				Tuple						     =      False
	A_							:				Union[str, Any]						     =      False
	A_							:				List[Any]						     =      (
	    (
	        FunnelModel,
	        FunnelForMaskedLM,
	        FunnelForPreTraining,
	        FunnelForQuestionAnswering,
	        FunnelForTokenClassification,
	    )
	    if is_torch_available()
	    else ()
	)
	A_							:				Any						     =      (
	    {
	        """feature-extraction""": (FunnelBaseModel, FunnelModel),
	        """fill-mask""": FunnelForMaskedLM,
	        """question-answering""": FunnelForQuestionAnswering,
	        """text-classification""": FunnelForSequenceClassification,
	        """token-classification""": FunnelForTokenClassification,
	        """zero-shot""": FunnelForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)




	def 							a_		(   self  :	Optional[Any]      ,    lowercase  :	str      ,    lowercase  :	List[Any]      ,    lowercase  :	Optional[int]=False	)					-> Dict:
						__snake_case          =	super()._prepare_for_class(lowercase      ,    lowercase      ,    return_labels=lowercase	)

						if return_labels:
											if model_class in get_values(lowercase	):
																__snake_case          =	torch.zeros(
																    (self.model_tester.batch_size, self.model_tester.seq_length)      ,    dtype=torch.long      ,    device=lowercase	)
						return inputs_dict




	def 							a_		(   self  :	Tuple	)					-> Optional[int]:
						__snake_case          =	FunnelModelTester(self	)
						__snake_case          =	ConfigTester(self      ,    config_class=lowercase	)




	def 							a_		(   self  :	Dict	)					-> List[str]:
						self.config_tester.run_common_tests()




	def 							a_		(   self  :	List[str]	)					-> Tuple:
						__snake_case          =	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*lowercase	)




	def 							a_		(   self  :	str	)					-> Optional[int]:
						__snake_case          =	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_pretraining(*lowercase	)




	def 							a_		(   self  :	Any	)					-> int:
						__snake_case          =	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*lowercase	)




	def 							a_		(   self  :	List[Any]	)					-> Dict:
						__snake_case          =	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*lowercase	)




	def 							a_		(   self  :	Tuple	)					-> List[str]:
						__snake_case          =	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_question_answering(*lowercase	)




	def 							a_		(   self  :	Optional[int]      ,    lowercase  :	Optional[int]	)					-> Dict:
						if hasattr(lowercase      ,    """weight"""	) and module.weight is not None:
											module.weight.data.fill_(3	)
						if hasattr(lowercase      ,    """bias"""	) and module.bias is not None:
											module.bias.data.fill_(3	)

						for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
											if hasattr(lowercase      ,    lowercase	) and getattr(lowercase      ,    lowercase	) is not None:
																__snake_case          =	getattr(lowercase      ,    lowercase	)
																weight.data.fill_(3	)

@require_torch
class     A__      (      lowercase_     ,	unittest.TestCase     ):
	A_							:				str						     =      False
	A_							:				Tuple						     =      False
	A_							:				Optional[Any]						     =      (
	    (FunnelBaseModel, FunnelForMultipleChoice, FunnelForSequenceClassification) if is_torch_available() else ()
	)




	def 							a_		(   self  :	Any	)					-> Optional[int]:
						__snake_case          =	FunnelModelTester(self      ,    base=lowercase	)
						__snake_case          =	ConfigTester(self      ,    config_class=lowercase	)




	def 							a_		(   self  :	Optional[Any]	)					-> List[Any]:
						self.config_tester.run_common_tests()




	def 							a_		(   self  :	Optional[int]	)					-> Dict:
						__snake_case          =	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_base_model(*lowercase	)




	def 							a_		(   self  :	str	)					-> str:
						__snake_case          =	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_sequence_classification(*lowercase	)




	def 							a_		(   self  :	str	)					-> Optional[int]:
						__snake_case          =	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_multiple_choice(*lowercase	)




	def 							a_		(   self  :	str	)					-> Union[str, Any]:
						__snake_case      ,      __snake_case          =	self.model_tester.prepare_config_and_inputs_for_common()
						__snake_case          =	True

						for model_class in self.all_model_classes:
											if model_class.__name__ == "FunnelBaseModel":
																continue
											__snake_case          =	model_class(lowercase	)
											model.to(lowercase	)
											model.train()
											__snake_case          =	self._prepare_for_class(lowercase      ,    lowercase      ,    return_labels=lowercase	)
											__snake_case          =	model(**lowercase	).loss
											loss.backward()




	def 							a_		(   self  :	Any      ,    lowercase  :	int	)					-> Optional[int]:
						if hasattr(lowercase      ,    """weight"""	) and module.weight is not None:
											module.weight.data.fill_(3	)
						if hasattr(lowercase      ,    """bias"""	) and module.bias is not None:
											module.bias.data.fill_(3	)

						for param in ["r_w_bias", "r_r_bias", "r_kernel", "r_s_bias", "seg_embed"]:
											if hasattr(lowercase      ,    lowercase	) and getattr(lowercase      ,    lowercase	) is not None:
																__snake_case          =	getattr(lowercase      ,    lowercase	)
																weight.data.fill_(3	)

@require_torch
@require_sentencepiece
@require_tokenizers
class     A__      (      unittest.TestCase     ):




	def 							a_		(   self  :	List[str]	)					-> List[str]:
						__snake_case          =	13
						__snake_case          =	7
						__snake_case          =	torch.arange(0      ,    batch_size * sequence_length	).long().reshape(lowercase      ,    lowercase	)
						__snake_case          =	[0, 1, 2, 3, 4, 5, 6, 4, 1, 3, 5, 0, 1]
						__snake_case          =	torch.tensor([[2] + [0] * a + [1] * (sequence_length - a - 1) for a in lengths]	)

						__snake_case          =	FunnelModel.from_pretrained("""sgugger/funnel-random-tiny"""	)
						__snake_case          =	model(lowercase      ,    token_type_ids=lowercase	)[0].abs()

						__snake_case          =	torch.tensor(23_44.83_52	)
						__snake_case          =	torch.tensor(0.80_52	)
						self.assertTrue(torch.allclose(output.sum()      ,    lowercase      ,    atol=1e-4	)	)
						self.assertTrue(torch.allclose(output.mean()      ,    lowercase      ,    atol=1e-4	)	)

						__snake_case          =	torch.tensor([[1] * 7, [1] * 4 + [0] * 3] * 6 + [[0, 1, 1, 0, 0, 1, 1]]	)
						__snake_case          =	model(lowercase      ,    attention_mask=lowercase      ,    token_type_ids=lowercase	)[0].abs()

						__snake_case          =	torch.tensor(23_43.84_25	)
						__snake_case          =	torch.tensor(0.80_49	)
						self.assertTrue(torch.allclose(output.sum()      ,    lowercase      ,    atol=1e-4	)	)
						self.assertTrue(torch.allclose(output.mean()      ,    lowercase      ,    atol=1e-4	)	)




	@slow
	def 							a_		(   self  :	Optional[Any]	)					-> Optional[Any]:
						__snake_case          =	FunnelTokenizer.from_pretrained("""huggingface/funnel-small"""	)
						__snake_case          =	FunnelModel.from_pretrained("""huggingface/funnel-small"""	)
						__snake_case          =	tokenizer("""Hello! I am the Funnel Transformer model."""      ,    return_tensors="""pt"""	)
						__snake_case          =	model(**lowercase	)[0]

						__snake_case          =	torch.tensor(2_35.72_46	)
						__snake_case          =	torch.tensor(0.02_56	)
						self.assertTrue(torch.allclose(output.sum()      ,    lowercase      ,    atol=1e-4	)	)
						self.assertTrue(torch.allclose(output.mean()      ,    lowercase      ,    atol=1e-4	)	)








"""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


UpperCamelCase__ :    List[str] 					=							logging.get_logger(__name__)
class     A__      (      lowercase_     ):
	A_							:				PyTorchBenchmarkArguments
	A_							:				PretrainedConfig
	A_							:				str						     =      "PyTorch"




	@property
	def 							a_		(   self  :	Optional[Any]	)					-> Tuple:
						return torch.__version__




	def 							a_		(   self  :	int      ,    lowercase  :	str      ,    lowercase  :	int      ,    lowercase  :	int	)					-> float:
						__snake_case          =	self._prepare_inference_func(lowercase      ,    lowercase      ,    lowercase	)
						return self._measure_speed(_inference	)




	def 							a_		(   self  :	Dict      ,    lowercase  :	str      ,    lowercase  :	int      ,    lowercase  :	int	)					-> [Memory, Optional[MemorySummary]]:
						__snake_case          =	self._prepare_inference_func(lowercase      ,    lowercase      ,    lowercase	)
						return self._measure_memory(_inference	)




	def 							a_		(   self  :	str      ,    lowercase  :	str      ,    lowercase  :	int      ,    lowercase  :	int	)					-> float:
						__snake_case          =	self._prepare_train_func(lowercase      ,    lowercase      ,    lowercase	)
						return self._measure_speed(_train	)




	def 							a_		(   self  :	List[Any]      ,    lowercase  :	str      ,    lowercase  :	int      ,    lowercase  :	int	)					-> [Memory, Optional[MemorySummary]]:
						__snake_case          =	self._prepare_train_func(lowercase      ,    lowercase      ,    lowercase	)
						return self._measure_memory(_train	)




	def 							a_		(   self  :	Union[str, Any]      ,    lowercase  :	str      ,    lowercase  :	int      ,    lowercase  :	int	)					-> Callable[[], None]:
						__snake_case          =	self.config_dict[model_name]

						if self.args.torchscript:
											__snake_case          =	True

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

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

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

						def encoder_decoder_forward():
											with torch.no_grad():
																__snake_case          =	inference_model(lowercase      ,    decoder_input_ids=lowercase	)
											return outputs

						def encoder_forward():
											with torch.no_grad():
																__snake_case          =	inference_model(lowercase	)
											return outputs

						__snake_case          =	encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
						return _forward




	def 							a_		(   self  :	Optional[int]      ,    lowercase  :	str      ,    lowercase  :	int      ,    lowercase  :	int	)					-> Callable[[], None]:
						__snake_case          =	self.config_dict[model_name]

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

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

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

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

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

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




	def 							a_		(   self  :	Tuple      ,    lowercase  :	List[Any]	)					-> 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(
																    lowercase      ,    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          =	timeit.repeat(
											    lowercase      ,    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(lowercase	) / 10.0
						except RuntimeError as e:
											self.print_fn(f'''Doesn\'t fit on GPU. {e}'''	)
											return "N/A"




	def 							a_		(   self  :	Union[str, Any]      ,    lowercase  :	Callable[[], None]	)					-> [Memory, MemorySummary]:
						try:
											if self.args.trace_memory_line_by_line:
																__snake_case          =	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."""	)
																					__snake_case          =	"""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          =	nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx	)
																					__snake_case          =	nvml.nvmlDeviceGetMemoryInfo(lowercase	)
																					__snake_case          =	meminfo.used
																					__snake_case          =	Memory(lowercase	)
																					# shutdown nvml
																					nvml.nvmlShutdown()
											else:
																# cpu
																__snake_case          =	measure_peak_memory_cpu(lowercase	)
																__snake_case          =	Memory(lowercase	) if isinstance(lowercase      ,    lowercase	) else memory_bytes

											if self.args.trace_memory_line_by_line:
																__snake_case          =	stop_memory_tracing(lowercase	)
											else:
																__snake_case          =	None

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


