






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


lowercase     :							str	     =			logging.get_logger(__name__)




class 						UpperCAmelCase_  (		SCREAMING_SNAKE_CASE__    ):


						'''simple docstring'''




						A   :      PyTorchBenchmarkArguments
						A   :      PretrainedConfig
						A   :      str						 =    "PyTorch"






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






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )       -> float:
								snake_case_ :							Tuple         =				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]]:
								snake_case_ :							int         =				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:
								snake_case_ :							str         =				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]]:
								snake_case_ :							List[Any]         =				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]:
								snake_case_ :							str         =				self.config_dict[model_name]

								if self.args.torchscript:
										snake_case_ :							Union[str, Any]         =				True

								snake_case_ :							List[str]         =				(
								    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:
												snake_case_ :							List[str]         =				config.architectures[0]
												snake_case_ :							List[str]         =				__import__("transformers"						, fromlist=[model_class] )
												snake_case_ :							Union[str, Any]         =				getattr(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )
												snake_case_ :							int         =				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:
										snake_case_ :							Any         =				MODEL_MAPPING[config.__class__](_SCREAMING_SNAKE_CASE )

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

								# encoder-decoder has vocab size saved differently
								snake_case_ :							List[str]         =				config.vocab_size if hasattr(_SCREAMING_SNAKE_CASE						, "vocab_size" ) else config.encoder.vocab_size
								snake_case_ :							Optional[int]         =				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():
												snake_case_ :							Union[str, Any]         =				torch.jit.trace(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )
								else:
										snake_case_ :							List[Any]         =				model

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

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

								snake_case_ :							Union[str, Any]         =				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]:
								snake_case_ :							Tuple         =				self.config_dict[model_name]

								snake_case_ :							Any         =				(
								    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:
												snake_case_ :							Optional[Any]         =				config.architectures[0]
												snake_case_ :							Optional[Any]         =				__import__("transformers"						, fromlist=[model_class] )
												snake_case_ :							int         =				getattr(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )
												snake_case_ :							Optional[Any]         =				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:
										snake_case_ :							Optional[int]         =				MODEL_WITH_LM_HEAD_MAPPING[config.__class__](_SCREAMING_SNAKE_CASE )

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

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

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

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

								snake_case_ :							Tuple         =				(
								    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:
								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
										snake_case_ :							Optional[int]         =				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]:
								try:
										if self.args.trace_memory_line_by_line:
												snake_case_ :							Dict         =				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_ :							int         =				"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_ :							Tuple         =				nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx )
														snake_case_ :							Optional[Any]         =				nvml.nvmlDeviceGetMemoryInfo(_SCREAMING_SNAKE_CASE )
														snake_case_ :							Tuple         =				meminfo.used
														snake_case_ :							Any         =				Memory(_SCREAMING_SNAKE_CASE )
														# shutdown nvml
														nvml.nvmlShutdown()
										else:
												# cpu
												snake_case_ :							str         =				measure_peak_memory_cpu(_SCREAMING_SNAKE_CASE )
												snake_case_ :							Optional[int]         =				Memory(_SCREAMING_SNAKE_CASE ) if isinstance(_SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE ) else memory_bytes

										if self.args.trace_memory_line_by_line:
												snake_case_ :							str         =				stop_memory_tracing(_SCREAMING_SNAKE_CASE )
										else:
												snake_case_ :							int         =				None

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








import glob
import os

from get_test_info import get_tester_classes


if __name__ == "__main__":
							lowercase     :							List[Any]	     =			[]

							lowercase     :							Optional[int]	     =			os.path.join('''tests''', '''models''', '''**''', '''test_modeling_*.py''')
							lowercase     :							Optional[Any]	     =			glob.glob(pattern)
							# TODO: deal with TF/Flax too
							lowercase     :							Tuple	     =			[
							    x for x in test_files if not (x.startswith('''test_modeling_tf_''') or x.startswith('''test_modeling_flax_'''))
							]

							for test_file in test_files:
														lowercase     :							Union[str, Any]	     =			get_tester_classes(test_file)
														for tester_class in tester_classes:
																					# A few tester classes don't have `parent` parameter in `__init__`.
																					# TODO: deal this better
																					try:
																												lowercase     :							Union[str, Any]	     =			tester_class(parent=None)
																					except Exception:
																												continue
																					if hasattr(tester, '''get_config'''):
																												lowercase     :							Tuple	     =			tester.get_config()
																												for k, v in config.to_dict().items():
																																			if isinstance(v, int):
																																										lowercase     :							Tuple	     =			None
																																										if k in ["vocab_size"]:
																																																	lowercase     :							str	     =			1_00
																																										elif k in ["max_position_embeddings"]:
																																																	lowercase     :							List[str]	     =			1_28
																																										elif k in ["hidden_size", "d_model"]:
																																																	lowercase     :							Union[str, Any]	     =			40
																																										elif k == ["num_layers", "num_hidden_layers", "num_encoder_layers", "num_decoder_layers"]:
																																																	lowercase     :							List[str]	     =			5
																																										if target is not None and v > target:
																																																	failures.append(
																																																	    F"""{tester_class.__name__} will produce a `config` of type `{config.__class__.__name__}`"""
																																																	    F""" with config[\"{k}\"] = {v} which is too large for testing! Set its value to be smaller"""
																																																	    F""" than {target}."""
																																																	)

							if len(failures) > 0:
														raise Exception(F"""There were {len(failures)} failures:\n""" + '''\n'''.join(failures))
