






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_  (		lowerCamelCase_    ):


						'''simple docstring'''




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






						@property
						def 					_lowerCAmelCase				(		self )       -> str:
								return torch.__version__






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )       -> List[str]:
								snake_case_ :							Any         =				self._prepare_inference_func(snake_case_						, snake_case_						, snake_case_ )
								return self._measure_speed(_inference )






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )       -> Optional[Any]:
								snake_case_ :							str         =				self._prepare_inference_func(snake_case_						, snake_case_						, snake_case_ )
								return self._measure_memory(_inference )






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )       -> Any:
								snake_case_ :							Tuple         =				self._prepare_train_func(snake_case_						, snake_case_						, snake_case_ )
								return self._measure_speed(_train )






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )       -> List[str]:
								snake_case_ :							Union[str, Any]         =				self._prepare_train_func(snake_case_						, snake_case_						, snake_case_ )
								return self._measure_memory(_train )






						def 					_lowerCAmelCase				(		self						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE						, _SCREAMING_SNAKE_CASE )       -> Any:
								snake_case_ :							Optional[int]         =				self.config_dict[model_name]

								if self.args.torchscript:
										snake_case_ :							Dict         =				True

								snake_case_ :							Dict         =				(
								    hasattr(snake_case_						, "architectures" )
								    and isinstance(config.architectures						, 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_ :							int         =				__import__("transformers"						, fromlist=[model_class] )
												snake_case_ :							Dict         =				getattr(snake_case_						, snake_case_ )
												snake_case_ :							Union[str, Any]         =				model_cls(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[Any]         =				MODEL_MAPPING[config.__class__](snake_case_ )

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

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

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

								def encoder_forward():
										with torch.no_grad():
												snake_case_ :							List[str]         =				inference_model(snake_case_ )
										return outputs

								snake_case_ :							List[str]         =				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 )       -> List[Any]:
								snake_case_ :							List[Any]         =				self.config_dict[model_name]

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

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

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

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

								def compute_loss_and_backprob_encoder_decoder():
										snake_case_ :							Optional[int]         =				train_model(snake_case_						, decoder_input_ids=snake_case_						, labels=snake_case_ )[0]
										loss.backward()
										return loss

								snake_case_ :							List[str]         =				(
								    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 )       -> Optional[int]:

								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(
												    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_ :							Tuple         =				timeit.repeat(
										    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(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 )       -> int:

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

										if self.args.trace_memory_line_by_line:
												snake_case_ :							Dict         =				stop_memory_tracing(snake_case_ )
										else:
												snake_case_ :							Tuple         =				None

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

'''simple docstring'''


from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


snake_case__     :			Optional[Any]      			=  {'''configuration_nat''': ['''NAT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''NatConfig''']}


try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 snake_case__     :			Tuple      			=  [
     '''NAT_PRETRAINED_MODEL_ARCHIVE_LIST''',
     '''NatForImageClassification''',
     '''NatModel''',
     '''NatPreTrainedModel''',
     '''NatBackbone''',
 ]

if TYPE_CHECKING:
 from .configuration_nat import NAT_PRETRAINED_CONFIG_ARCHIVE_MAP, NatConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_nat import (
      NAT_PRETRAINED_MODEL_ARCHIVE_LIST,
      NatBackbone,
      NatForImageClassification,
      NatModel,
      NatPreTrainedModel,
  )

else:
 import sys

 snake_case__     :			int      			=  _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)






