





"""simple docstring"""




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

_SCREAMING_SNAKE_CASE 		=						logging.get_logger(__name__)


def 							SCREAMING_SNAKE_CASE__  (      __a				,	__a				):
	def run_func(__a				):
		@wraps(__a				)
		def run_in_eager_mode(*__a				,	**__a				):
			return func(*__a				,	**__a				)

		@wraps(__a				)
		@tf.function(experimental_compile=__a				)
		def run_in_graph_mode(*__a				,	**__a				):
			return func(*__a				,	**__a				)

		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__  (      __a				,	__a				,	__a				):
	snake_case_						:     Dict     =		random.Random()
	snake_case_						:     int     =		[rng.randint(0				,	vocab_size - 1				) for i in range(batch_size * sequence_length				)]
	return tf.constant(__a				,	shape=(batch_size, sequence_length)				,	dtype=tf.intaa				)


class SCREAMING_SNAKE_CASE_  (   snake_case_   ):
						__magic_name__:     TensorFlowBenchmarkArguments
						__magic_name__:     PretrainedConfig
						__magic_name__:     str	 =    "TensorFlow"






						@property
						def UpperCAmelCase_			( self : List[Any]		)							->      int:







							"""simple docstring"""


							return tf.__version__






						def UpperCAmelCase_			( self : Optional[Any]						,				_A : str						,				_A : int						,				_A : int		)							->      float:







							"""simple docstring"""


							snake_case_						:     Dict     =		self.args.strategy
							if strategy is None:
								raise ValueError('A device strategy has to be initialized before using TensorFlow.'		)
							snake_case_						:     Union[str, Any]     =		self._prepare_inference_func(_A						,				_A						,				_A		)
							return self._measure_speed(_inference		)






						def UpperCAmelCase_			( self : Tuple						,				_A : str						,				_A : int						,				_A : int		)							->      float:







							"""simple docstring"""


							snake_case_						:     List[str]     =		self.args.strategy
							if strategy is None:
								raise ValueError('A device strategy has to be initialized before using TensorFlow.'		)
							snake_case_						:     Optional[int]     =		self._prepare_train_func(_A						,				_A						,				_A		)
							return self._measure_speed(_train		)






						def UpperCAmelCase_			( self : Optional[int]						,				_A : str						,				_A : int						,				_A : int		)							->      [Memory, Optional[MemorySummary]]:







							"""simple docstring"""


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






						def UpperCAmelCase_			( self : Dict						,				_A : str						,				_A : int						,				_A : int		)							->      [Memory, Optional[MemorySummary]]:







							"""simple docstring"""


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

							snake_case_						:     Union[str, Any]     =		self._prepare_train_func(_A						,				_A						,				_A		)
							return self._measure_memory(_train		)






						def UpperCAmelCase_			( self : str						,				_A : str						,				_A : int						,				_A : int		)							->      Callable[[], None]:







							"""simple docstring"""


							snake_case_						:     Tuple     =		self.config_dict[model_name]

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

							snake_case_						:     Tuple     =		(
							    hasattr(_A						,				'architectures'		)
							    and isinstance(config.architectures						,				_A		)
							    and len(config.architectures		) > 0
							)
							if not self.args.only_pretrain_model and has_model_class_in_config:
								try:
									snake_case_						:     Dict     =		'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
									snake_case_						:     Union[str, Any]     =		__import__('transformers'						,				fromlist=[model_class]		)
									snake_case_						:     Optional[Any]     =		getattr(_A						,				_A		)
									snake_case_						:     Dict     =		model_cls(_A		)
								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_						:     int     =		TF_MODEL_MAPPING[config.__class__](_A		)

							# encoder-decoder has vocab size saved differently
							snake_case_						:     Union[str, Any]     =		config.vocab_size if hasattr(_A						,				'vocab_size'		) else config.encoder.vocab_size
							snake_case_						:     List[str]     =		random_input_ids(_A						,				_A						,				_A		)

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

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

							snake_case_						:     List[Any]     =		encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

							return _inference






						def UpperCAmelCase_			( self : Tuple						,				_A : str						,				_A : int						,				_A : int		)							->      Callable[[], None]:







							"""simple docstring"""


							snake_case_						:     Any     =		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.'		)

							snake_case_						:     Union[str, Any]     =		(
							    hasattr(_A						,				'architectures'		)
							    and isinstance(config.architectures						,				_A		)
							    and len(config.architectures		) > 0
							)
							if not self.args.only_pretrain_model and has_model_class_in_config:
								try:
									snake_case_						:     Optional[Any]     =		'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
									snake_case_						:     Dict     =		__import__('transformers'						,				fromlist=[model_class]		)
									snake_case_						:     List[Any]     =		getattr(_A						,				_A		)
									snake_case_						:     Union[str, Any]     =		model_cls(_A		)
								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_						:     str     =		TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](_A		)

							# encoder-decoder has vocab size saved differently
							snake_case_						:     Tuple     =		config.vocab_size if hasattr(_A						,				'vocab_size'		) else config.encoder.vocab_size
							snake_case_						:     int     =		random_input_ids(_A						,				_A						,				_A		)

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

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

							snake_case_						:     int     =		encoder_decoder_train if config.is_encoder_decoder else encoder_train

							return _train






						def UpperCAmelCase_			( self : int						,				_A : Union[str, 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(_A						,				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(
									    _A						,				repeat=self.args.repeat						,				number=10						,				)

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






						def UpperCAmelCase_			( self : Dict						,				_A : Callable[[], None]		)							->      [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.'		)
										snake_case_						:     Any     =		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.'		)
											snake_case_						:     str     =		'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_						:     List[Any]     =		nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx		)
											snake_case_						:     Dict     =		nvml.nvmlDeviceGetMemoryInfo(_A		)
											snake_case_						:     Tuple     =		meminfo.used
											snake_case_						:     Optional[int]     =		Memory(_A		)
											# 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.'		)
											snake_case_						:     List[str]     =		None
										else:
											snake_case_						:     Tuple     =		measure_peak_memory_cpu(_A		)
											snake_case_						:     Optional[Any]     =		Memory(_A		) if isinstance(_A						,				_A		) else memory_bytes
									if self.args.trace_memory_line_by_line:
										snake_case_						:     List[Any]     =		stop_memory_tracing(_A		)
										if memory is None:
											snake_case_						:     int     =		summary.total
									else:
										snake_case_						:     List[str]     =		None

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








from typing import TYPE_CHECKING

from ....utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_SCREAMING_SNAKE_CASE 		=						{
    """configuration_mctct""": ["""MCTCT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """MCTCTConfig"""],
    """feature_extraction_mctct""": ["""MCTCTFeatureExtractor"""],
    """processing_mctct""": ["""MCTCTProcessor"""],
}


try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						_SCREAMING_SNAKE_CASE 		=						[
						    """MCTCT_PRETRAINED_MODEL_ARCHIVE_LIST""",
						    """MCTCTForCTC""",
						    """MCTCTModel""",
						    """MCTCTPreTrainedModel""",
						]


if TYPE_CHECKING:
						from .configuration_mctct import MCTCT_PRETRAINED_CONFIG_ARCHIVE_MAP, MCTCTConfig
						from .feature_extraction_mctct import MCTCTFeatureExtractor
						from .processing_mctct import MCTCTProcessor

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_mctct import MCTCT_PRETRAINED_MODEL_ARCHIVE_LIST, MCTCTForCTC, MCTCTModel, MCTCTPreTrainedModel

else:
						import sys

						_SCREAMING_SNAKE_CASE 		=						_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

