






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

_UpperCamelCase    :   Union[str, Any]   			=						logging.get_logger(__name__)




def       snake_case  (A_				:bool     ,      A_				:bool				):







  '''simple docstring'''

  def run_func(A_				:List[Any]				):
    @wraps(_SCREAMING_SNAKE_CASE				)
    def run_in_eager_mode(*A_				:List[Any]     ,      **A_				:Union[str, Any]				):
      return func(*_SCREAMING_SNAKE_CASE     ,      **_SCREAMING_SNAKE_CASE				)

    @wraps(_SCREAMING_SNAKE_CASE				)
    @tf.function(experimental_compile=_SCREAMING_SNAKE_CASE				)
    def run_in_graph_mode(*A_				:int     ,      **A_				:List[Any]				):
      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       snake_case  (A_				:int     ,      A_				:int     ,      A_				:int				):







  '''simple docstring'''

  _A					:      str    =       random.Random()
  _A					:      List[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	):
    __magic_name__   					=					42
    __magic_name__   					=					42
    __magic_name__   					=					"TensorFlow"





    @property
    def     lowerCamelCase__				(							self     : Optional[int]      ):


      '''simple docstring'''




      return tf.__version__





    def     lowerCamelCase__				(							self     : Tuple    ,		A     : Union[str, Any]    ,		A     : Optional[int]    ,		A     : Optional[int]      ):


      '''simple docstring'''




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





    def     lowerCamelCase__				(							self     : Union[str, Any]    ,		A     : Any    ,		A     : int    ,		A     : Optional[Any]      ):


      '''simple docstring'''




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





    def     lowerCamelCase__				(							self     : Optional[int]    ,		A     : str    ,		A     : Any    ,		A     : List[str]      ):


      '''simple docstring'''




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





    def     lowerCamelCase__				(							self     : Union[str, Any]    ,		A     : Optional[int]    ,		A     : str    ,		A     : Optional[int]      ):


      '''simple docstring'''




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

      _A					:      Dict    =       self._prepare_train_func(_snake_case    ,		_snake_case    ,		_snake_case      )
      return self._measure_memory(_train      )





    def     lowerCamelCase__				(							self     : Any    ,		A     : Any    ,		A     : List[Any]    ,		A     : int      ):


      '''simple docstring'''




      _A					:      List[Any]    =       self.config_dict[model_name]

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

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

      # encoder-decoder has vocab size saved differently
      _A					:      List[Any]    =       config.vocab_size if hasattr(_snake_case    ,		'vocab_size'      ) else config.encoder.vocab_size
      _A					:      List[str]    =       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      )

      _A					:      Optional[int]    =       encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

      return _inference





    def     lowerCamelCase__				(							self     : int    ,		A     : Any    ,		A     : Optional[int]    ,		A     : Union[str, Any]      ):


      '''simple docstring'''




      _A					:      Dict    =       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.'      )

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

      # encoder-decoder has vocab size saved differently
      _A					:      Optional[Any]    =       config.vocab_size if hasattr(_snake_case    ,		'vocab_size'      ) else config.encoder.vocab_size
      _A					:      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_train():
        _A					:      List[str]    =       model(_snake_case    ,		decoder_input_ids=_snake_case    ,		labels=_snake_case    ,		training=_snake_case      )[0]
        _A					:      int    =       tf.gradients(_snake_case    ,		model.trainable_variables      )
        return gradients

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

      _A					:      Dict    =       encoder_decoder_train if config.is_encoder_decoder else encoder_train

      return _train





    def     lowerCamelCase__				(							self     : List[Any]    ,		A     : Tuple      ):


      '''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
          _A					:      Dict    =       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     lowerCamelCase__				(							self     : Optional[Any]    ,		A     : Dict      ):


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

_UpperCamelCase    :   str   			=						os.path.join(os.path.dirname(__file__), 'src')
sys.path.append(SRC_DIR)


from transformers import (
    AutoConfig,
    AutoModel,
    AutoModelForCausalLM,
    AutoModelForMaskedLM,
    AutoModelForQuestionAnswering,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    add_start_docstrings,
)


_UpperCamelCase    :   int   			=						[
    'torch',
    'numpy',
    'tokenizers',
    'filelock',
    'requests',
    'tqdm',
    'regex',
    'sentencepiece',
    'sacremoses',
    'importlib_metadata',
    'huggingface_hub',
]




@add_start_docstrings(AutoConfig.__doc__				)
def       snake_case  (*A_				:Optional[int]     ,      **A_				:int				):







  '''simple docstring'''

  return AutoConfig.from_pretrained(*A_     ,      **A_				)




@add_start_docstrings(AutoTokenizer.__doc__				)
def       snake_case  (*A_				:Optional[int]     ,      **A_				:List[Any]				):







  '''simple docstring'''

  return AutoTokenizer.from_pretrained(*A_     ,      **A_				)




@add_start_docstrings(AutoModel.__doc__				)
def       snake_case  (*A_				:Optional[Any]     ,      **A_				:Tuple				):







  '''simple docstring'''

  return AutoModel.from_pretrained(*A_     ,      **A_				)




@add_start_docstrings(AutoModelForCausalLM.__doc__				)
def       snake_case  (*A_				:str     ,      **A_				:Optional[int]				):







  '''simple docstring'''

  return AutoModelForCausalLM.from_pretrained(*A_     ,      **A_				)




@add_start_docstrings(AutoModelForMaskedLM.__doc__				)
def       snake_case  (*A_				:Optional[Any]     ,      **A_				:Dict				):







  '''simple docstring'''

  return AutoModelForMaskedLM.from_pretrained(*A_     ,      **A_				)




@add_start_docstrings(AutoModelForSequenceClassification.__doc__				)
def       snake_case  (*A_				:Dict     ,      **A_				:str				):







  '''simple docstring'''

  return AutoModelForSequenceClassification.from_pretrained(*A_     ,      **A_				)






@add_start_docstrings(AutoModelForQuestionAnswering.__doc__				)
def       snake_case  (*A_				:Dict     ,      **A_				:List[Any]				):







  '''simple docstring'''

  return AutoModelForQuestionAnswering.from_pretrained(*A_     ,      **A_				)




