
'''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

lowercase__						       =      logging.get_logger(__name__)




def    _UpperCamelCase						(       SCREAMING_SNAKE_CASE__ ,       SCREAMING_SNAKE_CASE__			)  ->							Any:


		'''simple docstring'''


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

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

				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    _UpperCamelCase						(       SCREAMING_SNAKE_CASE__ ,       SCREAMING_SNAKE_CASE__ ,       SCREAMING_SNAKE_CASE__			)  ->							["tf.Tensor"]:


		'''simple docstring'''


		snake_case  :       Dict     			=						random.Random()
		snake_case  :       Optional[Any]     			=						[rng.randint(0 ,       vocab_size - 1			) for i in range(batch_size * sequence_length			)]
		return tf.constant(SCREAMING_SNAKE_CASE__ ,       shape=(batch_size, sequence_length) ,       dtype=tf.intaa			)






class 	snake_case__       (      __SCREAMING_SNAKE_CASE					):







		"""simple docstring"""

		lowerCamelCase     							=						42
		lowerCamelCase     							=						42
		lowerCamelCase     							=						"TensorFlow"



		@property
		def       lowerCAmelCase				(      self				:			Tuple      )						->			Optional[int]:






				"""simple docstring"""







				return tf.__version__



		def       lowerCAmelCase				(      self				:			Any  ,	UpperCamelCase__				:			str  ,	UpperCamelCase__				:			int  ,	UpperCamelCase__				:			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  :       Tuple     			=						self._prepare_inference_func(UpperCamelCase__  ,	UpperCamelCase__  ,	UpperCamelCase__      )
				return self._measure_speed(_inference      )



		def       lowerCAmelCase				(      self				:			Any  ,	UpperCamelCase__				:			str  ,	UpperCamelCase__				:			int  ,	UpperCamelCase__				:			int      )						->			float:






				"""simple docstring"""







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



		def       lowerCAmelCase				(      self				:			List[Any]  ,	UpperCamelCase__				:			str  ,	UpperCamelCase__				:			int  ,	UpperCamelCase__				:			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]  ,	UpperCamelCase__      )
				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  :       int     			=						self._prepare_inference_func(UpperCamelCase__  ,	UpperCamelCase__  ,	UpperCamelCase__      )
				return self._measure_memory(_inference      )



		def       lowerCAmelCase				(      self				:			Dict  ,	UpperCamelCase__				:			str  ,	UpperCamelCase__				:			int  ,	UpperCamelCase__				:			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]  ,	UpperCamelCase__      )
				snake_case  :       Optional[int]     			=						self.args.strategy
				if strategy is None:
						raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''      )

				snake_case  :       Any     			=						self._prepare_train_func(UpperCamelCase__  ,	UpperCamelCase__  ,	UpperCamelCase__      )
				return self._measure_memory(_train      )



		def       lowerCAmelCase				(      self				:			Tuple  ,	UpperCamelCase__				:			str  ,	UpperCamelCase__				:			int  ,	UpperCamelCase__				:			int      )						->			Callable[[], None]:






				"""simple docstring"""







				snake_case  :       Optional[int]     			=						self.config_dict[model_name]

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

				snake_case  :       int     			=						(
				    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:
								snake_case  :       int     			=						'''TF''' + config.architectures[0]  # prepend 'TF' for tensorflow model
								snake_case  :       str     			=						__import__('''transformers'''  ,	fromlist=[model_class]      )
								snake_case  :       Optional[Any]     			=						getattr(UpperCamelCase__  ,	UpperCamelCase__      )
								snake_case  :       List[Any]     			=						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:
						snake_case  :       List[Any]     			=						TF_MODEL_MAPPING[config.__class__](UpperCamelCase__      )

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

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

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

				snake_case  :       str     			=						encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

				return _inference



		def       lowerCAmelCase				(      self				:			Union[str, Any]  ,	UpperCamelCase__				:			str  ,	UpperCamelCase__				:			int  ,	UpperCamelCase__				:			int      )						->			Callable[[], None]:






				"""simple docstring"""







				snake_case  :       Tuple     			=						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  :       Any     			=						(
				    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:
								snake_case  :       Tuple     			=						'''TF''' + config.architectures[0]  # prepend 'TF' for tensorflow model
								snake_case  :       str     			=						__import__('''transformers'''  ,	fromlist=[model_class]      )
								snake_case  :       int     			=						getattr(UpperCamelCase__  ,	UpperCamelCase__      )
								snake_case  :       Optional[int]     			=						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:
						snake_case  :       List[Any]     			=						TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](UpperCamelCase__      )

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

				@run_with_tf_optimizations(self.args.eager_mode  ,	self.args.use_xla      )
				def encoder_decoder_train():
						snake_case  :       Dict     			=						model(UpperCamelCase__  ,	decoder_input_ids=UpperCamelCase__  ,	labels=UpperCamelCase__  ,	training=UpperCamelCase__      )[0]
						snake_case  :       List[Any]     			=						tf.gradients(UpperCamelCase__  ,	model.trainable_variables      )
						return gradients

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

				snake_case  :       List[Any]     			=						encoder_decoder_train if config.is_encoder_decoder else encoder_train

				return _train



		def       lowerCAmelCase				(      self				:			Tuple  ,	UpperCamelCase__				:			int      )						->			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(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
								snake_case  :       Optional[int]     			=						timeit.repeat(
								    UpperCamelCase__  ,	repeat=self.args.repeat  ,	number=10  ,	)

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



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

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




'''simple docstring'''
from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow


if is_tf_available():
						import numpy as np
						import tensorflow as tf

						from transformers import TFCamembertModel






@require_tf
@require_sentencepiece
@require_tokenizers
class 	snake_case__       (      unittest.TestCase					):







		"""simple docstring"""

		@slow
		def       lowerCAmelCase				(      self				:			Optional[Any]      )						->			Optional[Any]:






				"""simple docstring"""







				snake_case  :       Any     			=						TFCamembertModel.from_pretrained('''jplu/tf-camembert-base'''      )

				snake_case  :       List[Any]     			=						tf.convert_to_tensor(
				    [[5, 121, 11, 660, 16, 730, 2_5543, 110, 83, 6]]  ,	dtype=tf.intaa  ,	)  # J'aime le camembert !"

				snake_case  :       Any     			=						model(UpperCamelCase__      )['''last_hidden_state''']
				snake_case  :       int     			=						tf.TensorShape((1, 10, 768)      )
				self.assertEqual(output.shape  ,	UpperCamelCase__      )
				# compare the actual values for a slice.
				snake_case  :       Dict     			=						tf.convert_to_tensor(
				    [[[-0.0_254, 0.0_235, 0.1_027], [0.0_606, -0.1_811, -0.0_418], [-0.1_561, -0.1_127, 0.2_687]]]  ,	dtype=tf.floataa  ,	)
				# camembert = torch.hub.load('pytorch/fairseq', 'camembert.v0')
				# camembert.eval()
				# expected_slice = roberta.model.forward(input_ids)[0][:, :3, :3].detach()

				self.assertTrue(np.allclose(output[:, :3, :3].numpy()  ,	expected_slice.numpy()  ,	atol=1e-4      )      )


