




'''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_						:      List[str]       =						logging.get_logger(__name__)





def      __a     (       _UpperCamelCase:   bool ,						_UpperCamelCase:   bool   )						->			List[Any]:

     """simple docstring"""




     def run_func(_UpperCamelCase:   Optional[int]   ):
          @wraps(_UpperCamelCase   )
          def run_in_eager_mode(*_UpperCamelCase:   Dict ,						**_UpperCamelCase:   Union[str, Any]   ):
               return func(*_UpperCamelCase ,						**_UpperCamelCase   )

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

     """simple docstring"""




     _snake_case							=				random.Random()
     _snake_case							=				[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        _a    (						__lowerCAmelCase							):
     SCREAMING_SNAKE_CASE_							:   TensorFlowBenchmarkArguments
     SCREAMING_SNAKE_CASE_							:   PretrainedConfig
     SCREAMING_SNAKE_CASE_							:   str       =						"TensorFlow"






     @property
     def     _lowercase     (			self   )							->      str:
          return tf.__version__






     def     _lowercase     (			self ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE   )							->      float:
          # initialize GPU on separate process
          _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(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE   )
          return self._measure_speed(_inference   )






     def     _lowercase     (			self ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE   )							->      float:
          _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(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE   )
          return self._measure_speed(_train   )






     def     _lowercase     (			self ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE   )							->      [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] ,_SCREAMING_SNAKE_CASE   )
          _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(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE   )
          return self._measure_memory(_inference   )






     def     _lowercase     (			self ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE   )							->      [Memory, Optional[MemorySummary]]:
          if self.args.is_gpu:
               tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx] ,_SCREAMING_SNAKE_CASE   )
          _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(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE   )
          return self._measure_memory(_train   )






     def     _lowercase     (			self ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE   )							->      Callable[[], None]:
          _snake_case							=				self.config_dict[model_name]

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

          _snake_case							=				(
              hasattr(_SCREAMING_SNAKE_CASE ,"architectures"   )
              and isinstance(config.architectures ,_SCREAMING_SNAKE_CASE   )
              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(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE   )
                    _snake_case							=				model_cls(_SCREAMING_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:
               _snake_case							=				TF_MODEL_MAPPING[config.__class__](_SCREAMING_SNAKE_CASE   )

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

          @run_with_tf_optimizations(self.args.eager_mode ,self.args.use_xla   )
          def encoder_decoder_forward():
               return model(_SCREAMING_SNAKE_CASE ,decoder_input_ids=_SCREAMING_SNAKE_CASE ,training=_SCREAMING_SNAKE_CASE   )

          @run_with_tf_optimizations(self.args.eager_mode ,self.args.use_xla   )
          def encoder_forward():
               return model(_SCREAMING_SNAKE_CASE ,training=_SCREAMING_SNAKE_CASE   )

          _snake_case							=				encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

          return _inference






     def     _lowercase     (			self ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE   )							->      Callable[[], None]:
          _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(_SCREAMING_SNAKE_CASE ,"architectures"   )
              and isinstance(config.architectures ,_SCREAMING_SNAKE_CASE   )
              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(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE   )
                    _snake_case							=				model_cls(_SCREAMING_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:
               _snake_case							=				TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](_SCREAMING_SNAKE_CASE   )

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

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

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

          _snake_case							=				encoder_decoder_train if config.is_encoder_decoder else encoder_train

          return _train






     def     _lowercase     (			self ,_SCREAMING_SNAKE_CASE   )							->      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(_SCREAMING_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
                    _snake_case							=				timeit.repeat(
                        _SCREAMING_SNAKE_CASE ,repeat=self.args.repeat ,number=10 ,)

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






     def     _lowercase     (			self ,_SCREAMING_SNAKE_CASE   )							->      [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."   )
                         _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(_SCREAMING_SNAKE_CASE   )
                              _snake_case							=				meminfo.used
                              _snake_case							=				Memory(_SCREAMING_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."   )
                              _snake_case							=				None
                         else:
                              _snake_case							=				measure_peak_memory_cpu(_SCREAMING_SNAKE_CASE   )
                              _snake_case							=				Memory(_SCREAMING_SNAKE_CASE   ) if isinstance(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE   ) else memory_bytes
                    if self.args.trace_memory_line_by_line:
                         _snake_case							=				stop_memory_tracing(_SCREAMING_SNAKE_CASE   )
                         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'''





def      __a     (       _UpperCamelCase:   int   )						->			str:

     """simple docstring"""




     if number > 0:
          raise ValueError("input must be a negative integer"   )
     _snake_case							=				len(bin(_UpperCamelCase   )[3:]   )
     _snake_case							=				bin(abs(_UpperCamelCase   ) - (1 << binary_number_length)   )[3:]
     _snake_case							=				(
         (
             "1"
             + "0" * (binary_number_length - len(_UpperCamelCase   ))
             + twos_complement_number
         )
         if number < 0
         else "0"
     )
     return "0b" + twos_complement_number


if __name__ == "__main__":
   import doctest

   doctest.testmod()


