




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

snake_case_ : Tuple    		=       logging.get_logger(__name__)


def    A						(__A   :    bool   ,				__A   :    bool       )   ->      Optional[Any]:







     """simple docstring"""


     def run_func(__A   :    Optional[Any]       ):
          @wraps(__A       )
          def run_in_eager_mode(*__A   :    Dict   ,				**__A   :    List[Any]       ):
               return func(*__A   ,				**__A       )

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







     """simple docstring"""


     UpperCAmelCase_          =							random.Random()
     UpperCAmelCase_          =							[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 __snake_case			(			a		):
   UpperCAmelCase__  :    TensorFlowBenchmarkArguments
   UpperCAmelCase__  :    PretrainedConfig
   UpperCAmelCase__  :    str          =  "TensorFlow"







   @property
   def lowerCamelCase						(      self		:						List[str]):

        """simple docstring"""


        return tf.__version__







   def lowerCamelCase						(      self		:						Dict					,						_snake_case		:						str					,						_snake_case		:						int					,						_snake_case		:						int):

        """simple docstring"""


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







   def lowerCamelCase						(      self		:						Any					,						_snake_case		:						str					,						_snake_case		:						int					,						_snake_case		:						int):

        """simple docstring"""


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







   def lowerCamelCase						(      self		:						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_          =							self.args.strategy
        if strategy is None:
             raise ValueError('''A device strategy has to be initialized before using TensorFlow.''')
        UpperCAmelCase_          =							self._prepare_inference_func(_snake_case					,						_snake_case					,						_snake_case)
        return self._measure_memory(_inference)







   def lowerCamelCase						(      self		:						Optional[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_          =							self.args.strategy
        if strategy is None:
             raise ValueError('''A device strategy has to be initialized before using TensorFlow.''')

        UpperCAmelCase_          =							self._prepare_train_func(_snake_case					,						_snake_case					,						_snake_case)
        return self._measure_memory(_train)







   def lowerCamelCase						(      self		:						Optional[int]					,						_snake_case		:						str					,						_snake_case		:						int					,						_snake_case		:						int):

        """simple docstring"""


        UpperCAmelCase_          =							self.config_dict[model_name]

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

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

        # encoder-decoder has vocab size saved differently
        UpperCAmelCase_          =							config.vocab_size if hasattr(_snake_case					,						'''vocab_size''') else config.encoder.vocab_size
        UpperCAmelCase_          =							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_          =							encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

        return _inference







   def lowerCamelCase						(      self		:						Optional[Any]					,						_snake_case		:						str					,						_snake_case		:						int					,						_snake_case		:						int):

        """simple docstring"""


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

        # encoder-decoder has vocab size saved differently
        UpperCAmelCase_          =							config.vocab_size if hasattr(_snake_case					,						'''vocab_size''') else config.encoder.vocab_size
        UpperCAmelCase_          =							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_          =							model(_snake_case					,						decoder_input_ids=_snake_case					,						labels=_snake_case					,						training=_snake_case)[0]
             UpperCAmelCase_          =							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_          =							model(_snake_case					,						labels=_snake_case					,						training=_snake_case)[0]
             UpperCAmelCase_          =							tf.gradients(_snake_case					,						model.trainable_variables)
             return gradients

        UpperCAmelCase_          =							encoder_decoder_train if config.is_encoder_decoder else encoder_train

        return _train







   def lowerCamelCase						(      self		:						Any					,						_snake_case		:						Optional[Any]):

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

                  return min(_snake_case) / 1_0.0
             except ResourceExhaustedError as e:
                  self.print_fn(F"""Doesn't fit on GPU. {e}""")







   def lowerCamelCase						(      self		:						Dict					,						_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_          =							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_          =							'''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_          =							nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx)
                            UpperCAmelCase_          =							nvml.nvmlDeviceGetMemoryInfo(_snake_case)
                            UpperCAmelCase_          =							meminfo.used
                            UpperCAmelCase_          =							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_          =							None
                       else:
                            UpperCAmelCase_          =							measure_peak_memory_cpu(_snake_case)
                            UpperCAmelCase_          =							Memory(_snake_case) if isinstance(_snake_case					,						_snake_case) else memory_bytes
                  if self.args.trace_memory_line_by_line:
                       UpperCAmelCase_          =							stop_memory_tracing(_snake_case)
                       if memory is None:
                            UpperCAmelCase_          =							summary.total
                  else:
                       UpperCAmelCase_          =							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 Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxSeqaSeqConfigWithPast
from ...utils import logging


a		:       Tuple			     = logging.get_logger(__name__)

a		:       Dict			     = {
    "t5-small": "https://huggingface.co/t5-small/resolve/main/config.json",
    "t5-base": "https://huggingface.co/t5-base/resolve/main/config.json",
    "t5-large": "https://huggingface.co/t5-large/resolve/main/config.json",
    "t5-3b": "https://huggingface.co/t5-3b/resolve/main/config.json",
    "t5-11b": "https://huggingface.co/t5-11b/resolve/main/config.json",
}
class     a    ( lowercase__   ):



      """simple docstring"""


      a  : int							       =	't5'
      a  : Dict							       =	['past_key_values']
      a  : Tuple							       =	{'hidden_size': 'd_model', 'num_attention_heads': 'num_heads', 'num_hidden_layers': 'num_layers'}
      def __init__(      self					:			str		,      __lowercase					:			Optional[int]=32128		,      __lowercase					:			Optional[int]=512		,      __lowercase					:			int=64		,      __lowercase					:			Any=2048		,      __lowercase					:			Tuple=6		,      __lowercase					:			Tuple=None		,      __lowercase					:			int=8		,      __lowercase					:			List[Any]=32		,      __lowercase					:			Dict=128		,      __lowercase					:			Optional[int]=0.1		,      __lowercase					:			int=1e-6		,      __lowercase					:			List[str]=1.0		,      __lowercase					:			List[str]="relu"		,      __lowercase					:			Dict=True		,      __lowercase					:			Optional[Any]=True		,      __lowercase					:			Tuple=0		,      __lowercase					:			List[str]=1		,      **__lowercase					:			Any		,      )							->      str:
        __UpperCAmelCase							:	Tuple           =  vocab_size
        __UpperCAmelCase							:	Optional[int]           =  d_model
        __UpperCAmelCase							:	Union[str, Any]           =  d_kv
        __UpperCAmelCase							:	Union[str, Any]           =  d_ff
        __UpperCAmelCase							:	int           =  num_layers
        __UpperCAmelCase							:	Any           =  (
            num_decoder_layers if num_decoder_layers is not None else self.num_layers
        )  # default = symmetry
        __UpperCAmelCase							:	Dict           =  num_heads
        __UpperCAmelCase							:	List[Any]           =  relative_attention_num_buckets
        __UpperCAmelCase							:	List[str]           =  relative_attention_max_distance
        __UpperCAmelCase							:	Union[str, Any]           =  dropout_rate
        __UpperCAmelCase							:	List[str]           =  layer_norm_epsilon
        __UpperCAmelCase							:	str           =  initializer_factor
        __UpperCAmelCase							:	Dict           =  feed_forward_proj
        __UpperCAmelCase							:	Optional[int]           =  use_cache

        __UpperCAmelCase							:	List[Any]           =  self.feed_forward_proj.split("""-"""					)
        __UpperCAmelCase							:	Tuple           =  act_info[-1]
        __UpperCAmelCase							:	int           =  act_info[0] == """gated"""

        if len(__lowercase					) > 1 and act_info[0] != "gated" or len(__lowercase					) > 2:
          raise ValueError(
              f"""`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer."""
              """Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. """
              """'gated-gelu' or 'relu'"""					)

        # for backwards compatibility
        if feed_forward_proj == "gated-gelu":
          __UpperCAmelCase							:	Dict           =  """gelu_new"""

        super().__init__(
            pad_token_id=__lowercase		,      eos_token_id=__lowercase		,      is_encoder_decoder=__lowercase		,      **__lowercase		,      )


class     a    ( lowercase__   ):



      """simple docstring"""


      @property
      def 	UpperCAmelCase      (      self					:			str					)							->      Mapping[str, Mapping[int, str]]:
        __UpperCAmelCase							:	Union[str, Any]           =  {
            """input_ids""": {0: """batch""", 1: """encoder_sequence"""},
            """attention_mask""": {0: """batch""", 1: """encoder_sequence"""},
        }
        if self.use_past:
          __UpperCAmelCase							:	List[Any]           =  """past_encoder_sequence + sequence"""
          __UpperCAmelCase							:	Optional[int]           =  {0: """batch"""}
          __UpperCAmelCase							:	Union[str, Any]           =  {0: """batch""", 1: """past_decoder_sequence + sequence"""}
        else:
          __UpperCAmelCase							:	str           =  {0: """batch""", 1: """decoder_sequence"""}
          __UpperCAmelCase							:	str           =  {0: """batch""", 1: """decoder_sequence"""}

        if self.use_past:
          self.fill_with_past_key_values_(__lowercase		,      direction="""inputs"""					)

        return common_inputs




      @property
      def 	UpperCAmelCase      (      self					:			int					)							->      int:
        return 13


