

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


def        __lowercase    (  UpperCAmelCase__		,					UpperCAmelCase__      ):

					"""simple docstring"""






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

										@wraps(UpperCAmelCase__      )
										@tf.function(experimental_compile=UpperCAmelCase__      )
										def run_in_graph_mode(*UpperCAmelCase__		,					**UpperCAmelCase__      ):
															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        __lowercase    (  UpperCAmelCase__		,					UpperCAmelCase__		,					UpperCAmelCase__      ):

					"""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     snake_case_			(				_a							):


							"""simple docstring"""





							__UpperCAmelCase       =42
							__UpperCAmelCase       =42
							__UpperCAmelCase       ="TensorFlow"




							@property
							def        A__  ( self							):
												return tf.__version__




							def        A__  ( self					,						_A					,						_A					,						_A							):
												# initialize GPU on separate process
												__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(_A					,						_A					,						_A							)
												return self._measure_speed(_inference							)




							def        A__  ( self					,						_A					,						_A					,						_A							):
												__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(_A					,						_A					,						_A							)
												return self._measure_speed(_train							)




							def        A__  ( self					,						_A					,						_A					,						_A							):
												# initialize GPU on separate process
												if self.args.is_gpu:
																	tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]					,						_A							)
												__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(_A					,						_A					,						_A							)
												return self._measure_memory(_inference							)




							def        A__  ( self					,						_A					,						_A					,						_A							):
												if self.args.is_gpu:
																	tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]					,						_A							)
												__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(_A					,						_A					,						_A							)
												return self._measure_memory(_train							)




							def        A__  ( self					,						_A					,						_A					,						_A							):
												__lowerCAmelCase					=							self.config_dict[model_name]

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

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

												# encoder-decoder has vocab size saved differently
												__lowerCAmelCase					=							config.vocab_size if hasattr(_A					,						'vocab_size'							) else config.encoder.vocab_size
												__lowerCAmelCase					=							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							)

												__lowerCAmelCase					=							encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

												return _inference




							def        A__  ( self					,						_A					,						_A					,						_A							):
												__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(_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:
																						__lowerCAmelCase					=							'TF' + config.architectures[0]  # prepend 'TF' for tensorflow model
																						__lowerCAmelCase					=							__import__('transformers'					,						fromlist=[model_class]							)
																						__lowerCAmelCase					=							getattr(_A					,						_A							)
																						__lowerCAmelCase					=							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:
																	__lowerCAmelCase					=							TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](_A							)

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

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

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

												__lowerCAmelCase					=							encoder_decoder_train if config.is_encoder_decoder else encoder_train

												return _train




							def        A__  ( self					,						_A							):
												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
																						__lowerCAmelCase					=							timeit.repeat(
																						    _A					,						repeat=self.args.repeat					,						number=1_0					,						)

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




							def        A__  ( self					,						_A							):
												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(_A							)
																																__lowerCAmelCase					=							meminfo.used
																																__lowerCAmelCase					=							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.'							)
																																__lowerCAmelCase					=							None
																											else:
																																__lowerCAmelCase					=							measure_peak_memory_cpu(_A							)
																																__lowerCAmelCase					=							Memory(_A							) if isinstance(_A					,						_A							) else memory_bytes
																						if self.args.trace_memory_line_by_line:
																											__lowerCAmelCase					=							stop_memory_tracing(_A							)
																											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 typing import TYPE_CHECKING

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


lowerCamelCase      				=     {
    '''configuration_data2vec_audio''': ['''DATA2VEC_AUDIO_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''Data2VecAudioConfig'''],
    '''configuration_data2vec_text''': [
        '''DATA2VEC_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''Data2VecTextConfig''',
        '''Data2VecTextOnnxConfig''',
    ],
    '''configuration_data2vec_vision''': [
        '''DATA2VEC_VISION_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''Data2VecVisionConfig''',
        '''Data2VecVisionOnnxConfig''',
    ],
}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		lowerCamelCase      				=     [
		    '''DATA2VEC_AUDIO_PRETRAINED_MODEL_ARCHIVE_LIST''',
		    '''Data2VecAudioForAudioFrameClassification''',
		    '''Data2VecAudioForCTC''',
		    '''Data2VecAudioForSequenceClassification''',
		    '''Data2VecAudioForXVector''',
		    '''Data2VecAudioModel''',
		    '''Data2VecAudioPreTrainedModel''',
		]
		lowerCamelCase      				=     [
		    '''DATA2VEC_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST''',
		    '''Data2VecTextForCausalLM''',
		    '''Data2VecTextForMaskedLM''',
		    '''Data2VecTextForMultipleChoice''',
		    '''Data2VecTextForQuestionAnswering''',
		    '''Data2VecTextForSequenceClassification''',
		    '''Data2VecTextForTokenClassification''',
		    '''Data2VecTextModel''',
		    '''Data2VecTextPreTrainedModel''',
		]
		lowerCamelCase      				=     [
		    '''DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST''',
		    '''Data2VecVisionForImageClassification''',
		    '''Data2VecVisionForMaskedImageModeling''',
		    '''Data2VecVisionForSemanticSegmentation''',
		    '''Data2VecVisionModel''',
		    '''Data2VecVisionPreTrainedModel''',
		]

if is_tf_available():
		lowerCamelCase      				=     [
		    '''TFData2VecVisionForImageClassification''',
		    '''TFData2VecVisionForSemanticSegmentation''',
		    '''TFData2VecVisionModel''',
		    '''TFData2VecVisionPreTrainedModel''',
		]

if TYPE_CHECKING:
		from .configuration_dataavec_audio import DATA2VEC_AUDIO_PRETRAINED_CONFIG_ARCHIVE_MAP, DataaVecAudioConfig
		from .configuration_dataavec_text import (
		    DATA2VEC_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP,
		    DataaVecTextConfig,
		    DataaVecTextOnnxConfig,
		)
		from .configuration_dataavec_vision import (
		    DATA2VEC_VISION_PRETRAINED_CONFIG_ARCHIVE_MAP,
		    DataaVecVisionConfig,
		    DataaVecVisionOnnxConfig,
		)

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_dataavec_audio import (
				    DATA2VEC_AUDIO_PRETRAINED_MODEL_ARCHIVE_LIST,
				    DataaVecAudioForAudioFrameClassification,
				    DataaVecAudioForCTC,
				    DataaVecAudioForSequenceClassification,
				    DataaVecAudioForXVector,
				    DataaVecAudioModel,
				    DataaVecAudioPreTrainedModel,
				)
				from .modeling_dataavec_text import (
				    DATA2VEC_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    DataaVecTextForCausalLM,
				    DataaVecTextForMaskedLM,
				    DataaVecTextForMultipleChoice,
				    DataaVecTextForQuestionAnswering,
				    DataaVecTextForSequenceClassification,
				    DataaVecTextForTokenClassification,
				    DataaVecTextModel,
				    DataaVecTextPreTrainedModel,
				)
				from .modeling_dataavec_vision import (
				    DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST,
				    DataaVecVisionForImageClassification,
				    DataaVecVisionForMaskedImageModeling,
				    DataaVecVisionForSemanticSegmentation,
				    DataaVecVisionModel,
				    DataaVecVisionPreTrainedModel,
				)
		if is_tf_available():
				from .modeling_tf_dataavec_vision import (
				    TFDataaVecVisionForImageClassification,
				    TFDataaVecVisionForSemanticSegmentation,
				    TFDataaVecVisionModel,
				    TFDataaVecVisionPreTrainedModel,
				)

else:
		import sys

		lowerCamelCase      				=     _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)






