






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

_lowerCAmelCase	:	Optional[Any]				=      logging.get_logger(__name__)


def 	SCREAMING_SNAKE_CASE__       (  snake_case       : bool			,					snake_case       : bool				)->		Optional[int]:






     '''simple docstring'''



     def run_func(snake_case       : Dict				):
          @wraps(snake_case				)
          def run_in_eager_mode(*snake_case       : Optional[int]			,					**snake_case       : Union[str, Any]				):
               return func(*snake_case			,					**snake_case				)

          @wraps(snake_case				)
          @tf.function(experimental_compile=snake_case				)
          def run_in_graph_mode(*snake_case       : Optional[Any]			,					**snake_case       : Any				):
               return func(*snake_case			,					**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 	SCREAMING_SNAKE_CASE__       (  snake_case       : int			,					snake_case       : int			,					snake_case       : int				)->		["tf.Tensor"]:






     '''simple docstring'''



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




class 			lowerCAmelCase__			(    __magic_name__   ):
   SCREAMING_SNAKE_CASE_    =42
   SCREAMING_SNAKE_CASE_    =42
   SCREAMING_SNAKE_CASE_    ="TensorFlow"
   @property
   def       __a							(		self      :	List[str]							):


        '''simple docstring'''
        return tf.__version__
   def       __a							(		self      :	Any							,			snake_case__      :	str							,			snake_case__      :	int							,			snake_case__      :	int							):


        '''simple docstring'''
        # initialize GPU on separate process
        UpperCAmelCase__						:    str							=							self.args.strategy
        if strategy is None:
             raise ValueError("A device strategy has to be initialized before using TensorFlow."							)
        UpperCAmelCase__						:    Union[str, Any]							=							self._prepare_inference_func(snake_case__							,			snake_case__							,			snake_case__							)
        return self._measure_speed(_inference							)
   def       __a							(		self      :	int							,			snake_case__      :	str							,			snake_case__      :	int							,			snake_case__      :	int							):


        '''simple docstring'''
        UpperCAmelCase__						:    Tuple							=							self.args.strategy
        if strategy is None:
             raise ValueError("A device strategy has to be initialized before using TensorFlow."							)
        UpperCAmelCase__						:    List[str]							=							self._prepare_train_func(snake_case__							,			snake_case__							,			snake_case__							)
        return self._measure_speed(_train							)
   def       __a							(		self      :	Tuple							,			snake_case__      :	str							,			snake_case__      :	int							,			snake_case__      :	int							):


        '''simple docstring'''
        # initialize GPU on separate process
        if self.args.is_gpu:
             tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]							,			snake_case__							)
        UpperCAmelCase__						:    int							=							self.args.strategy
        if strategy is None:
             raise ValueError("A device strategy has to be initialized before using TensorFlow."							)
        UpperCAmelCase__						:    List[Any]							=							self._prepare_inference_func(snake_case__							,			snake_case__							,			snake_case__							)
        return self._measure_memory(_inference							)
   def       __a							(		self      :	List[Any]							,			snake_case__      :	str							,			snake_case__      :	int							,			snake_case__      :	int							):


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

        UpperCAmelCase__						:    Dict							=							self._prepare_train_func(snake_case__							,			snake_case__							,			snake_case__							)
        return self._measure_memory(_train							)
   def       __a							(		self      :	Optional[int]							,			snake_case__      :	str							,			snake_case__      :	int							,			snake_case__      :	int							):


        '''simple docstring'''
        UpperCAmelCase__						:    Tuple							=							self.config_dict[model_name]

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

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

        # encoder-decoder has vocab size saved differently
        UpperCAmelCase__						:    List[Any]							=							config.vocab_size if hasattr(snake_case__							,			"vocab_size"							) else config.encoder.vocab_size
        UpperCAmelCase__						:    Tuple							=							random_input_ids(snake_case__							,			snake_case__							,			snake_case__							)

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

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

        UpperCAmelCase__						:    Tuple							=							encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

        return _inference
   def       __a							(		self      :	Tuple							,			snake_case__      :	str							,			snake_case__      :	int							,			snake_case__      :	int							):


        '''simple docstring'''
        UpperCAmelCase__						:    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."							)

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

        # encoder-decoder has vocab size saved differently
        UpperCAmelCase__						:    Any							=							config.vocab_size if hasattr(snake_case__							,			"vocab_size"							) else config.encoder.vocab_size
        UpperCAmelCase__						:    Any							=							random_input_ids(snake_case__							,			snake_case__							,			snake_case__							)

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

        @run_with_tf_optimizations(self.args.eager_mode							,			self.args.use_xla							)
        def encoder_train():
             UpperCAmelCase__						:    Union[str, Any]							=							model(snake_case__							,			labels=snake_case__							,			training=snake_case__							)[0]
             UpperCAmelCase__						:    int							=							tf.gradients(snake_case__							,			model.trainable_variables							)
             return gradients

        UpperCAmelCase__						:    List[str]							=							encoder_decoder_train if config.is_encoder_decoder else encoder_train

        return _train
   def       __a							(		self      :	List[str]							,			snake_case__      :	int							):


        '''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(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
                  UpperCAmelCase__						:    List[str]							=							timeit.repeat(
                      snake_case__							,			repeat=self.args.repeat							,			number=1_0							,			)

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




   def       __a							(		self      :	Optional[Any]							,			snake_case__      :	Callable[[], None]							):


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

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










"""simple docstring"""



import unittest

from transformers import is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, require_torch, slow


if is_flax_available():
 import optax
 from flax.training.common_utils import onehot

 from transformers import AutoTokenizer, FlaxMTaForConditionalGeneration
 from transformers.models.ta.modeling_flax_ta import shift_tokens_right




@require_torch
@require_sentencepiece
@require_tokenizers
@require_flax
class 			lowerCAmelCase__			(    unittest.TestCase   ):
   @slow
   def       __a							(		self      :	List[str]							):


        '''simple docstring'''
        UpperCAmelCase__						:    Union[str, Any]							=							FlaxMTaForConditionalGeneration.from_pretrained("google/mt5-small"							)
        UpperCAmelCase__						:    str							=							AutoTokenizer.from_pretrained("google/mt5-small"							)

        UpperCAmelCase__						:    str							=							tokenizer("Hello there"							,			return_tensors="np"							).input_ids
        UpperCAmelCase__						:    Optional[Any]							=							tokenizer("Hi I am"							,			return_tensors="np"							).input_ids

        UpperCAmelCase__						:    Any							=							shift_tokens_right(snake_case__							,			model.config.pad_token_id							,			model.config.decoder_start_token_id							)

        UpperCAmelCase__						:    Any							=							model(snake_case__							,			decoder_input_ids=snake_case__							).logits
        UpperCAmelCase__						:    Any							=							optax.softmax_cross_entropy(snake_case__							,			onehot(snake_case__							,			logits.shape[-1]							)							).mean()

        UpperCAmelCase__						:    List[str]							=							-(labels.shape[-1] * loss.item())

        UpperCAmelCase__						:    Dict							=							-84.9127
        self.assertTrue(abs(mtf_score - EXPECTED_SCORE							) < 1e-4							)


