




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    :		Tuple			   =       logging.get_logger(__name__)

def       _SCREAMING_SNAKE_CASE      (							lowercase     :	bool     ,							lowercase     :	bool					):


						'''simple docstring'''



						def run_func(lowercase     :	Optional[int]					):
												@wraps(_UpperCamelCase					)
												def run_in_eager_mode(*lowercase     :	List[Any]     ,							**lowercase     :	Optional[int]					):
																		return func(*_UpperCamelCase     ,							**_UpperCamelCase					)

												@wraps(_UpperCamelCase					)
												@tf.function(experimental_compile=_UpperCamelCase					)
												def run_in_graph_mode(*lowercase     :	Optional[Any]     ,							**lowercase     :	Optional[int]					):
																		return func(*_UpperCamelCase     ,							**_UpperCamelCase					)

												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      (							lowercase     :	int     ,							lowercase     :	int     ,							lowercase     :	int					):


						'''simple docstring'''



						lowerCamelCase_							=   random.Random()
						lowerCamelCase_							=   [rng.randint(0     ,							vocab_size - 1					) for i in range(batch_size * sequence_length					)]
						return tf.constant(_UpperCamelCase     ,							shape=(batch_size, sequence_length)     ,							dtype=tf.intaa					)







class   A(						_UpperCamelCase       ):



	'''simple docstring'''


	UpperCamelCase						       =						42
	UpperCamelCase						       =						42
	UpperCamelCase						       =						"TensorFlow"

	@property
	def     a__			(    self :  Optional[Any]  )       ->     str:

							"""simple docstring"""






							return tf.__version__

	def     a__			(    self :  Union[str, Any]						,							A_ :  int						,							A_ :  Any						,							A_ :  str  )       ->     float:

							"""simple docstring"""






							lowerCamelCase_							=   self.args.strategy
							if strategy is None:
													raise ValueError('A device strategy has to be initialized before using TensorFlow.'  )
							lowerCamelCase_							=   self._prepare_inference_func(_SCREAMING_SNAKE_CASE						,							_SCREAMING_SNAKE_CASE						,							_SCREAMING_SNAKE_CASE  )
							return self._measure_speed(_inference  )

	def     a__			(    self :  Any						,							A_ :  Union[str, Any]						,							A_ :  Union[str, Any]						,							A_ :  Any  )       ->     float:

							"""simple docstring"""






							lowerCamelCase_							=   self.args.strategy
							if strategy is None:
													raise ValueError('A device strategy has to be initialized before using TensorFlow.'  )
							lowerCamelCase_							=   self._prepare_train_func(_SCREAMING_SNAKE_CASE						,							_SCREAMING_SNAKE_CASE						,							_SCREAMING_SNAKE_CASE  )
							return self._measure_speed(_train  )

	def     a__			(    self :  Union[str, Any]						,							A_ :  str						,							A_ :  Tuple						,							A_ :  List[Any]  )       ->     [Memory, Optional[MemorySummary]]:

							"""simple docstring"""






							if self.args.is_gpu:
													tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]						,							_SCREAMING_SNAKE_CASE  )
							lowerCamelCase_							=   self.args.strategy
							if strategy is None:
													raise ValueError('A device strategy has to be initialized before using TensorFlow.'  )
							lowerCamelCase_							=   self._prepare_inference_func(_SCREAMING_SNAKE_CASE						,							_SCREAMING_SNAKE_CASE						,							_SCREAMING_SNAKE_CASE  )
							return self._measure_memory(_inference  )

	def     a__			(    self :  Optional[int]						,							A_ :  Dict						,							A_ :  Optional[Any]						,							A_ :  Dict  )       ->     [Memory, Optional[MemorySummary]]:

							"""simple docstring"""






							if self.args.is_gpu:
													tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]						,							_SCREAMING_SNAKE_CASE  )
							lowerCamelCase_							=   self.args.strategy
							if strategy is None:
													raise ValueError('A device strategy has to be initialized before using TensorFlow.'  )

							lowerCamelCase_							=   self._prepare_train_func(_SCREAMING_SNAKE_CASE						,							_SCREAMING_SNAKE_CASE						,							_SCREAMING_SNAKE_CASE  )
							return self._measure_memory(_train  )

	def     a__			(    self :  Optional[int]						,							A_ :  List[Any]						,							A_ :  str						,							A_ :  Dict  )       ->     Callable[[], None]:

							"""simple docstring"""






							lowerCamelCase_							=   self.config_dict[model_name]

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

							lowerCamelCase_							=   (
							    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:
																			lowerCamelCase_							=   'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
																			lowerCamelCase_							=   __import__('transformers'						,							fromlist=[model_class]  )
																			lowerCamelCase_							=   getattr(_SCREAMING_SNAKE_CASE						,							_SCREAMING_SNAKE_CASE  )
																			lowerCamelCase_							=   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:
													lowerCamelCase_							=   TF_MODEL_MAPPING[config.__class__](_SCREAMING_SNAKE_CASE  )

							# encoder-decoder has vocab size saved differently
							lowerCamelCase_							=   config.vocab_size if hasattr(_SCREAMING_SNAKE_CASE						,							'vocab_size'  ) else config.encoder.vocab_size
							lowerCamelCase_							=   random_input_ids(_SCREAMING_SNAKE_CASE						,							_SCREAMING_SNAKE_CASE						,							_SCREAMING_SNAKE_CASE  )

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

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

							lowerCamelCase_							=   encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

							return _inference

	def     a__			(    self :  Optional[Any]						,							A_ :  Optional[int]						,							A_ :  Optional[int]						,							A_ :  List[str]  )       ->     Callable[[], None]:

							"""simple docstring"""






							lowerCamelCase_							=   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.'  )

							lowerCamelCase_							=   (
							    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:
																			lowerCamelCase_							=   'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
																			lowerCamelCase_							=   __import__('transformers'						,							fromlist=[model_class]  )
																			lowerCamelCase_							=   getattr(_SCREAMING_SNAKE_CASE						,							_SCREAMING_SNAKE_CASE  )
																			lowerCamelCase_							=   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:
													lowerCamelCase_							=   TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](_SCREAMING_SNAKE_CASE  )

							# encoder-decoder has vocab size saved differently
							lowerCamelCase_							=   config.vocab_size if hasattr(_SCREAMING_SNAKE_CASE						,							'vocab_size'  ) else config.encoder.vocab_size
							lowerCamelCase_							=   random_input_ids(_SCREAMING_SNAKE_CASE						,							_SCREAMING_SNAKE_CASE						,							_SCREAMING_SNAKE_CASE  )

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

							@run_with_tf_optimizations(self.args.eager_mode						,							self.args.use_xla  )
							def encoder_train():
													lowerCamelCase_							=   model(_SCREAMING_SNAKE_CASE						,							labels=_SCREAMING_SNAKE_CASE						,							training=_SCREAMING_SNAKE_CASE  )[0]
													lowerCamelCase_							=   tf.gradients(_SCREAMING_SNAKE_CASE						,							model.trainable_variables  )
													return gradients

							lowerCamelCase_							=   encoder_decoder_train if config.is_encoder_decoder else encoder_train

							return _train

	def     a__			(    self :  List[Any]						,							A_ :  Optional[Any]  )       ->     float:

							"""simple docstring"""






							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(_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
																			lowerCamelCase_							=   timeit.repeat(
																			    _SCREAMING_SNAKE_CASE						,							repeat=self.args.repeat						,							number=10						,							)

																			return min(_SCREAMING_SNAKE_CASE  ) / 10.0
													except ResourceExhaustedError as e:
																			self.print_fn(f"""Doesn\'t fit on GPU. {e}"""  )





	def     a__			(    self :  Tuple						,							A_ :  Tuple  )       ->     [Memory, MemorySummary]:

							"""simple docstring"""






							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.'  )
																									lowerCamelCase_							=   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.'  )
																															lowerCamelCase_							=   '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()
																															lowerCamelCase_							=   nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx  )
																															lowerCamelCase_							=   nvml.nvmlDeviceGetMemoryInfo(_SCREAMING_SNAKE_CASE  )
																															lowerCamelCase_							=   meminfo.used
																															lowerCamelCase_							=   Memory(_SCREAMING_SNAKE_CASE  )
																															# 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.'  )
																															lowerCamelCase_							=   None
																									else:
																															lowerCamelCase_							=   measure_peak_memory_cpu(_SCREAMING_SNAKE_CASE  )
																															lowerCamelCase_							=   Memory(_SCREAMING_SNAKE_CASE  ) if isinstance(_SCREAMING_SNAKE_CASE						,							_SCREAMING_SNAKE_CASE  ) else memory_bytes
																			if self.args.trace_memory_line_by_line:
																									lowerCamelCase_							=   stop_memory_tracing(_SCREAMING_SNAKE_CASE  )
																									if memory is None:
																															lowerCamelCase_							=   summary.total
																			else:
																									lowerCamelCase_							=   None

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









from ..utils import DummyObject, requires_backends







class   A(						metaclass=UpperCamelCase       ):



 '''simple docstring'''


 UpperCamelCase						       =						['''keras_nlp''']

 def __init__(    self :  Optional[int]						,							*A_ :  Any						,							**A_ :  Dict  )       ->     Optional[int]:

       """simple docstring"""






       requires_backends(self						,							['keras_nlp']  )



