






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






def        lowerCAmelCase (      UpperCamelCase_:							bool     , UpperCamelCase_:							bool					)       ->							Optional[int]:







  '''simple docstring'''







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

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







  '''simple docstring'''







  _a							=	random.Random()
  _a							=	[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    lowercase_    (_UpperCAmelCase     ):
      A__	:		TensorFlowBenchmarkArguments
      A__	:		PretrainedConfig
      A__	:		str       			="TensorFlow"






      @property
      def 			lowerCamelCase__    (    self	)	->Tuple:







        '''simple docstring'''


        return tf.__version__






      def 			lowerCamelCase__    (    self     ,     a_     ,     a_     ,     a_	)	->float:







        '''simple docstring'''


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






      def 			lowerCamelCase__    (    self     ,     a_     ,     a_     ,     a_	)	->float:







        '''simple docstring'''


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






      def 			lowerCamelCase__    (    self     ,     a_     ,     a_     ,     a_	)	->[Memory, Optional[MemorySummary]]:







        '''simple docstring'''


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






      def 			lowerCamelCase__    (    self     ,     a_     ,     a_     ,     a_	)	->[Memory, Optional[MemorySummary]]:







        '''simple docstring'''


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

        _a							=	self._prepare_train_func(a_     ,     a_     ,     a_	)
        return self._measure_memory(_train	)






      def 			lowerCamelCase__    (    self     ,     a_     ,     a_     ,     a_	)	->Callable[[], None]:







        '''simple docstring'''


        _a							=	self.config_dict[model_name]

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

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

        # encoder-decoder has vocab size saved differently
        _a							=	config.vocab_size if hasattr(a_     ,     "vocab_size"	) else config.encoder.vocab_size
        _a							=	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_	)

        _a							=	encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

        return _inference






      def 			lowerCamelCase__    (    self     ,     a_     ,     a_     ,     a_	)	->Callable[[], None]:







        '''simple docstring'''


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

        # encoder-decoder has vocab size saved differently
        _a							=	config.vocab_size if hasattr(a_     ,     "vocab_size"	) else config.encoder.vocab_size
        _a							=	random_input_ids(a_     ,     a_     ,     a_	)

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

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

        _a							=	encoder_decoder_train if config.is_encoder_decoder else encoder_train

        return _train






      def 			lowerCamelCase__    (    self     ,     a_	)	->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(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
            _a							=	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 			lowerCamelCase__    (    self     ,     a_	)	->[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."	)
              _a							=	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							=	"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							=	nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx	)
                _a							=	nvml.nvmlDeviceGetMemoryInfo(a_	)
                _a							=	meminfo.used
                _a							=	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."	)
                _a							=	None
              else:
                _a							=	measure_peak_memory_cpu(a_	)
                _a							=	Memory(a_	) if isinstance(a_     ,     a_	) else memory_bytes
            if self.args.trace_memory_line_by_line:
              _a							=	stop_memory_tracing(a_	)
              if memory is None:
                _a							=	summary.total
            else:
              _a							=	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 typing import TYPE_CHECKING

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


UpperCamelCase							       =    {
    """configuration_transfo_xl""": ["""TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP""", """TransfoXLConfig"""],
    """tokenization_transfo_xl""": ["""TransfoXLCorpus""", """TransfoXLTokenizer"""],
}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase							       =    [
         """TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST""",
         """AdaptiveEmbedding""",
         """TransfoXLForSequenceClassification""",
         """TransfoXLLMHeadModel""",
         """TransfoXLModel""",
         """TransfoXLPreTrainedModel""",
         """load_tf_weights_in_transfo_xl""",
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCamelCase							       =    [
         """TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST""",
         """TFAdaptiveEmbedding""",
         """TFTransfoXLForSequenceClassification""",
         """TFTransfoXLLMHeadModel""",
         """TFTransfoXLMainLayer""",
         """TFTransfoXLModel""",
         """TFTransfoXLPreTrainedModel""",
     ]


if TYPE_CHECKING:
     from .configuration_transfo_xl import TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP, TransfoXLConfig
     from .tokenization_transfo_xl import TransfoXLCorpus, TransfoXLTokenizer

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_transfo_xl import (
              TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
              AdaptiveEmbedding,
              TransfoXLForSequenceClassification,
              TransfoXLLMHeadModel,
              TransfoXLModel,
              TransfoXLPreTrainedModel,
              load_tf_weights_in_transfo_xl,
          )

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_transfo_xl import (
              TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
              TFAdaptiveEmbedding,
              TFTransfoXLForSequenceClassification,
              TFTransfoXLLMHeadModel,
              TFTransfoXLMainLayer,
              TFTransfoXLModel,
              TFTransfoXLPreTrainedModel,
          )

else:
     import sys

     UpperCamelCase							       =    _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


