





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






def lowerCamelCase__    (							snake_case_	:  bool       ,      snake_case_	:  bool       )      ->				Optional[Any]:
   def run_func(snake_case_	:  Union[str, Any]       ):
      @wraps(snake_case_       )
      def run_in_eager_mode(*snake_case_	:  str       ,      **snake_case_	:  Any       ):
         return func(*snake_case_       ,      **snake_case_       )

      @wraps(snake_case_       )
      @tf.function(experimental_compile=snake_case_       )
      def run_in_graph_mode(*snake_case_	:  List[str]       ,      **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 lowerCamelCase__    (							snake_case_	:  int       ,      snake_case_	:  int       ,      snake_case_	:  int       )      ->				["tf.Tensor"]:
   __snake_case							       =   random.Random()
   __snake_case							       =   [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 		SCREAMING_SNAKE_CASE__				(       _UpperCAmelCase  ):
      A_						:				TensorFlowBenchmarkArguments
      A_						:				PretrainedConfig
      A_						:				str							=    "TensorFlow"


      @property
      def    a (self       :      str      ):





         """simple docstring"""



         return tf.__version__


      def    a (self       :      Optional[int]				,			a__       :      str				,			a__       :      int				,			a__       :      int      ):





         """simple docstring"""



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


      def    a (self       :      Dict				,			a__       :      str				,			a__       :      int				,			a__       :      int      ):





         """simple docstring"""



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


      def    a (self       :      List[str]				,			a__       :      str				,			a__       :      int				,			a__       :      int      ):





         """simple docstring"""



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


      def    a (self       :      Tuple				,			a__       :      str				,			a__       :      int				,			a__       :      int      ):





         """simple docstring"""



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

         __snake_case							       =   self._prepare_train_func(a__				,			a__				,			a__      )
         return self._measure_memory(_train      )


      def    a (self       :      Union[str, Any]				,			a__       :      str				,			a__       :      int				,			a__       :      int      ):





         """simple docstring"""



         __snake_case							       =   self.config_dict[model_name]

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

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

         # encoder-decoder has vocab size saved differently
         __snake_case							       =   config.vocab_size if hasattr(a__				,			'''vocab_size'''      ) else config.encoder.vocab_size
         __snake_case							       =   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__      )

         __snake_case							       =   encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

         return _inference


      def    a (self       :      Union[str, Any]				,			a__       :      str				,			a__       :      int				,			a__       :      int      ):





         """simple docstring"""



         __snake_case							       =   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.'''      )

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

         # encoder-decoder has vocab size saved differently
         __snake_case							       =   config.vocab_size if hasattr(a__				,			'''vocab_size'''      ) else config.encoder.vocab_size
         __snake_case							       =   random_input_ids(a__				,			a__				,			a__      )

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

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

         __snake_case							       =   encoder_decoder_train if config.is_encoder_decoder else encoder_train

         return _train


      def    a (self       :      List[Any]				,			a__       :      Dict      ):





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

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


      def    a (self       :      Dict				,			a__       :      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.'''      )
                  __snake_case							       =   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.'''      )
                     __snake_case							       =   '''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()
                     __snake_case							       =   nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx      )
                     __snake_case							       =   nvml.nvmlDeviceGetMemoryInfo(a__      )
                     __snake_case							       =   meminfo.used
                     __snake_case							       =   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.'''      )
                     __snake_case							       =   None
                  else:
                     __snake_case							       =   measure_peak_memory_cpu(a__      )
                     __snake_case							       =   Memory(a__      ) if isinstance(a__				,			a__      ) else memory_bytes
               if self.args.trace_memory_line_by_line:
                  __snake_case							       =   stop_memory_tracing(a__      )
                  if memory is None:
                     __snake_case							       =   summary.total
               else:
                  __snake_case							       =   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 accelerate import debug_launcher
from accelerate.test_utils import require_cpu, test_ops, test_script
@require_cpu
class        __A       (	unittest.TestCase    ):







		'''simple docstring'''





		def 					UpperCAmelCase    (	self						:  Optional[int]						)					->		str:






				"""simple docstring"""



				debug_launcher(test_script.main						)







		def 					UpperCAmelCase    (	self						:  Dict						)					->		Tuple:






				"""simple docstring"""



				debug_launcher(test_ops.main						)
