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


A_:      Tuple        =			logging.get_logger(__name__)


class _lowercase					(       _UpperCAmelCase						):





							"""simple docstring"""




							lowerCAmelCase__       =					42
							lowerCAmelCase__       =					42
							lowerCAmelCase__       =					'PyTorch'



							@property
							def     _UpperCAmelCase       (    self  ):







														'''simple docstring'''






														return torch.__version__



							def     _UpperCAmelCase       (    self     ,   UpperCAmelCase     ,   UpperCAmelCase     ,   UpperCAmelCase  ):







														'''simple docstring'''






														_lowercase				=	self._prepare_inference_func(UpperCAmelCase     ,   UpperCAmelCase     ,   UpperCAmelCase  )
														return self._measure_speed(_inference  )



							def     _UpperCAmelCase       (    self     ,   UpperCAmelCase     ,   UpperCAmelCase     ,   UpperCAmelCase  ):







														'''simple docstring'''






														_lowercase				=	self._prepare_inference_func(UpperCAmelCase     ,   UpperCAmelCase     ,   UpperCAmelCase  )
														return self._measure_memory(_inference  )



							def     _UpperCAmelCase       (    self     ,   UpperCAmelCase     ,   UpperCAmelCase     ,   UpperCAmelCase  ):







														'''simple docstring'''






														_lowercase				=	self._prepare_train_func(UpperCAmelCase     ,   UpperCAmelCase     ,   UpperCAmelCase  )
														return self._measure_speed(_train  )



							def     _UpperCAmelCase       (    self     ,   UpperCAmelCase     ,   UpperCAmelCase     ,   UpperCAmelCase  ):







														'''simple docstring'''






														_lowercase				=	self._prepare_train_func(UpperCAmelCase     ,   UpperCAmelCase     ,   UpperCAmelCase  )
														return self._measure_memory(_train  )



							def     _UpperCAmelCase       (    self     ,   UpperCAmelCase     ,   UpperCAmelCase     ,   UpperCAmelCase  ):







														'''simple docstring'''






														_lowercase				=	self.config_dict[model_name]

														if self.args.torchscript:
																					_lowercase				=	True

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

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

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

														def encoder_decoder_forward():
																					with torch.no_grad():
																												_lowercase				=	inference_model(UpperCAmelCase     ,   decoder_input_ids=UpperCAmelCase  )
																					return outputs

														def encoder_forward():
																					with torch.no_grad():
																												_lowercase				=	inference_model(UpperCAmelCase  )
																					return outputs

														_lowercase				=	encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
														return _forward



							def     _UpperCAmelCase       (    self     ,   UpperCAmelCase     ,   UpperCAmelCase     ,   UpperCAmelCase  ):







														'''simple docstring'''






														_lowercase				=	self.config_dict[model_name]

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

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

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

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

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

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



							def     _UpperCAmelCase       (    self     ,   UpperCAmelCase  ):







														'''simple docstring'''






														try:
																					if self.args.is_tpu or self.args.torchscript:
																												# run additional 10 times to stabilize compilation for tpu and torchscript
																												logger.info("""Do inference on TPU or torchscript. Running model 5 times to stabilize compilation"""  )
																												timeit.repeat(
																												    UpperCAmelCase     ,   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
																					_lowercase				=	timeit.repeat(
																					    UpperCAmelCase     ,   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(UpperCAmelCase  ) / 10.0
														except RuntimeError as e:
																					self.print_fn(F'''Doesn\'t fit on GPU. {e}'''  )
																					return "N/A"



							def     _UpperCAmelCase       (    self     ,   UpperCAmelCase  ):







														'''simple docstring'''






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

																					if self.args.trace_memory_line_by_line:
																												_lowercase				=	stop_memory_tracing(UpperCAmelCase  )
																					else:
																												_lowercase				=	None

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



from ctypes import c_float, sizeof
from enum import Enum
from typing import TYPE_CHECKING, Optional, Union


if TYPE_CHECKING:
				from .. import AutoFeatureExtractor, AutoProcessor, AutoTokenizer  # tests_ignore


class _lowercase					(       _UpperCAmelCase						):





							"""simple docstring"""




							lowerCAmelCase__       =					c_float



							@property
							def     _UpperCAmelCase       (    self  ):







														'''simple docstring'''






														return sizeof(self.value  )






def 					__lowerCAmelCase			(  _A ,_A ,_A = 0					):



							"""simple docstring"""


							if dimension <= 0:
														_lowercase				=	fixed_dimension

							dimension -= num_token_to_add
							return dimension


def 					__lowerCAmelCase			(  _A ,_A					):



							"""simple docstring"""


							return num_parameters * dtype.size







def 					__lowerCAmelCase			(  _A					):



							"""simple docstring"""


							from .. import AutoFeatureExtractor, AutoProcessor, AutoTokenizer  # tests_ignore

							try:
														return AutoProcessor.from_pretrained(_A					)
							except (ValueError, OSError, KeyError):
														_lowercase				,       _lowercase				=	None, None
														try:
																					_lowercase				=	AutoTokenizer.from_pretrained(_A					)
														except (OSError, KeyError):
																					pass
														try:
																					_lowercase				=	AutoFeatureExtractor.from_pretrained(_A					)
														except (OSError, KeyError):
																					pass

														if tokenizer is not None and feature_extractor is not None:
																					raise ValueError(
																					    f'''Couldn\'t auto-detect preprocessor for {model_name}. Found both a tokenizer and a feature extractor.'''					)
														elif tokenizer is None and feature_extractor is None:
																					return None
														elif tokenizer is not None:
																					return tokenizer
														else:
																					return feature_extractor


