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

_A								=   logging.get_logger(__name__)

def _UpperCamelCase   (       SCREAMING_SNAKE_CASE_  ,       SCREAMING_SNAKE_CASE_							):
     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_							):
     lowercase_					:       str     				=					random.Random()
     lowercase_					:       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       UpperCAmelCase__  (       _snake_case			):






   """simple docstring"""
   A							:       TensorFlowBenchmarkArguments
   A							:       PretrainedConfig
   A							:       str     =    "TensorFlow"
   @property
   def 		_lowerCamelCase (self				)    ->     Union[str, Any]:
        return tf.__version__
   def 		_lowerCamelCase (self       ,				_a       ,				_a       ,				_a				)    ->     float:
        # initialize GPU on separate process
        lowercase_					:       Optional[int]     				=					self.args.strategy
        if strategy is None:
             raise ValueError('A device strategy has to be initialized before using TensorFlow.'				)
        lowercase_					:       str     				=					self._prepare_inference_func(_a       ,				_a       ,				_a				)
        return self._measure_speed(_inference				)
   def 		_lowerCamelCase (self       ,				_a       ,				_a       ,				_a				)    ->     float:
        lowercase_					:       Dict     				=					self.args.strategy
        if strategy is None:
             raise ValueError('A device strategy has to be initialized before using TensorFlow.'				)
        lowercase_					:       Any     				=					self._prepare_train_func(_a       ,				_a       ,				_a				)
        return self._measure_speed(_train				)
   def 		_lowerCamelCase (self       ,				_a       ,				_a       ,				_a				)    ->     [Memory, Optional[MemorySummary]]:
        # initialize GPU on separate process
        if self.args.is_gpu:
             tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]       ,				_a				)
        lowercase_					:       Optional[Any]     				=					self.args.strategy
        if strategy is None:
             raise ValueError('A device strategy has to be initialized before using TensorFlow.'				)
        lowercase_					:       Dict     				=					self._prepare_inference_func(_a       ,				_a       ,				_a				)
        return self._measure_memory(_inference				)
   def 		_lowerCamelCase (self       ,				_a       ,				_a       ,				_a				)    ->     [Memory, Optional[MemorySummary]]:
        if self.args.is_gpu:
             tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]       ,				_a				)
        lowercase_					:       Optional[Any]     				=					self.args.strategy
        if strategy is None:
             raise ValueError('A device strategy has to be initialized before using TensorFlow.'				)

        lowercase_					:       Any     				=					self._prepare_train_func(_a       ,				_a       ,				_a				)
        return self._measure_memory(_train				)
   def 		_lowerCamelCase (self       ,				_a       ,				_a       ,				_a				)    ->     Callable[[], None]:
        lowercase_					:       Optional[int]     				=					self.config_dict[model_name]

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

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

        # encoder-decoder has vocab size saved differently
        lowercase_					:       str     				=					config.vocab_size if hasattr(_a       ,				'vocab_size'				) else config.encoder.vocab_size
        lowercase_					:       str     				=					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				)

        lowercase_					:       Optional[Any]     				=					encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

        return _inference
   def 		_lowerCamelCase (self       ,				_a       ,				_a       ,				_a				)    ->     Callable[[], None]:
        lowercase_					:       int     				=					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.'				)

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

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

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

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

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

        return _train
   def 		_lowerCamelCase (self       ,				_a				)    ->     float:
        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
                  lowercase_					:       str     				=					timeit.repeat(
                      _a       ,				repeat=self.args.repeat       ,				number=10       ,				)

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




from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_A								=   {
    'configuration_bigbird_pegasus': [
        'BIGBIRD_PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'BigBirdPegasusConfig',
        'BigBirdPegasusOnnxConfig',
    ],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _A								=   [
     'BIGBIRD_PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST',
     'BigBirdPegasusForCausalLM',
     'BigBirdPegasusForConditionalGeneration',
     'BigBirdPegasusForQuestionAnswering',
     'BigBirdPegasusForSequenceClassification',
     'BigBirdPegasusModel',
     'BigBirdPegasusPreTrainedModel',
 ]


if TYPE_CHECKING:
 from .configuration_bigbird_pegasus import (
     BIGBIRD_PEGASUS_PRETRAINED_CONFIG_ARCHIVE_MAP,
     BigBirdPegasusConfig,
     BigBirdPegasusOnnxConfig,
 )

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_bigbird_pegasus import (
      BIGBIRD_PEGASUS_PRETRAINED_MODEL_ARCHIVE_LIST,
      BigBirdPegasusForCausalLM,
      BigBirdPegasusForConditionalGeneration,
      BigBirdPegasusForQuestionAnswering,
      BigBirdPegasusForSequenceClassification,
      BigBirdPegasusModel,
      BigBirdPegasusPreTrainedModel,
  )


else:
 import sys

 _A								=   _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




