




"""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__	(snake_case__							: bool		,       snake_case__							: bool   ):
  """simple docstring"""

  def run_func(snake_case__							: List[Any]   ):
    @wraps(_SCREAMING_SNAKE_CASE   )
    def run_in_eager_mode(*snake_case__							: List[Any]		,       **snake_case__							: 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(*snake_case__							: int		,       **snake_case__							: 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  UpperCAmelCase__	(snake_case__							: int		,       snake_case__							: int		,       snake_case__							: int   ):
  """simple docstring"""

  _snake_case							:		str								=  random.Random()
  _snake_case							:		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  lowercase(					__SCREAMING_SNAKE_CASE			):







 '''simple docstring'''





 lowercase__			 = 42
 lowercase__			 = 42
 lowercase__			 = "TensorFlow"


 @property
 def 			UpperCamelCase_       (    self:		Optional[Any]		):




   '''simple docstring'''



   return tf.__version__


 def 			UpperCamelCase_       (    self:		str,				a_:		Any,				a_:		str,				a_:		Tuple		):




   '''simple docstring'''



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


 def 			UpperCamelCase_       (    self:		Tuple,				a_:		Any,				a_:		str,				a_:		List[Any]		):




   '''simple docstring'''



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


 def 			UpperCamelCase_       (    self:		Dict,				a_:		Dict,				a_:		List[Any],				a_:		str		):




   '''simple docstring'''



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


 def 			UpperCamelCase_       (    self:		Any,				a_:		Any,				a_:		Tuple,				a_:		Tuple		):




   '''simple docstring'''



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

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


 def 			UpperCamelCase_       (    self:		str,				a_:		Optional[Any],				a_:		Dict,				a_:		Any		):




   '''simple docstring'''



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

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

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

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

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

   return _inference


 def 			UpperCamelCase_       (    self:		Dict,				a_:		Tuple,				a_:		Union[str, Any],				a_:		Dict		):




   '''simple docstring'''



   _snake_case							:		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."""		)

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

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

   _snake_case							:		Dict								=  encoder_decoder_train if config.is_encoder_decoder else encoder_train

   return _train


 def 			UpperCamelCase_       (    self:		Any,				a_:		List[str]		):




   '''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
       _snake_case							:		Dict								=  timeit.repeat(
           _snake_case,				repeat=self.args.repeat,				number=10,				)

       return min(_snake_case		) / 10.0
     except ResourceExhaustedError as e:
       self.print_fn(f"Doesn\'t fit on GPU. {e}"		)


 def 			UpperCamelCase_       (    self:		Tuple,				a_:		Union[str, Any]		):




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


from collections.abc import Generator
from math import sin


def  UpperCAmelCase__	(snake_case__							: bytes   ):
  """simple docstring"""

  if len(snake_case__   ) != 32:
    raise ValueError("""Input must be of length 32"""   )

  _snake_case							:		Optional[int]								=  B""""""
  for i in [3, 2, 1, 0]:
    little_endian += string_aa[8 * i : 8 * i + 8]
  return little_endian


def  UpperCAmelCase__	(snake_case__							: int   ):
  """simple docstring"""

  if i < 0:
    raise ValueError("""Input must be non-negative"""   )

  _snake_case							:		Optional[Any]								=  format(snake_case__		,       """08x"""   )[-8:]
  _snake_case							:		Optional[Any]								=  B""""""
  for i in [3, 2, 1, 0]:
    little_endian_hex += hex_rep[2 * i : 2 * i + 2].encode("""utf-8"""   )
  return little_endian_hex


def  UpperCAmelCase__	(snake_case__							: bytes   ):
  """simple docstring"""

  _snake_case							:		Union[str, Any]								=  B""""""
  for char in message:
    bit_string += format(snake_case__		,       """08b"""   ).encode("""utf-8"""   )
  _snake_case							:		List[Any]								=  format(len(snake_case__   )		,       """064b"""   ).encode("""utf-8"""   )

  # Pad bit_string to a multiple of 512 chars
  bit_string += b"1"
  while len(snake_case__   ) % 5_12 != 4_48:
    bit_string += b"0"
  bit_string += to_little_endian(start_len[32:]   ) + to_little_endian(start_len[:32]   )

  return bit_string


def  UpperCAmelCase__	(snake_case__							: bytes   ):
  """simple docstring"""

  if len(snake_case__   ) % 5_12 != 0:
    raise ValueError("""Input must have length that's a multiple of 512"""   )

  for pos in range(0		,       len(snake_case__   )		,       5_12   ):
    _snake_case							:		List[str]								=  bit_string[pos : pos + 5_12]
    _snake_case							:		List[str]								=  []
    for i in range(0		,       5_12		,       32   ):
      block_words.append(int(to_little_endian(block[i : i + 32]   )		,       2   )   )
    yield block_words


def  UpperCAmelCase__	(snake_case__							: int   ):
  """simple docstring"""

  if i < 0:
    raise ValueError("""Input must be non-negative"""   )

  _snake_case							:		Optional[int]								=  format(snake_case__		,       """032b"""   )
  _snake_case							:		str								=  """"""
  for c in i_str:
    new_str += "1" if c == "0" else "0"
  return int(snake_case__		,       2   )


def  UpperCAmelCase__	(snake_case__							: int		,       snake_case__							: int   ):
  """simple docstring"""

  return (a + b) % 2**32


def  UpperCAmelCase__	(snake_case__							: int		,       snake_case__							: int   ):
  """simple docstring"""

  if i < 0:
    raise ValueError("""Input must be non-negative"""   )
  if shift < 0:
    raise ValueError("""Shift must be non-negative"""   )
  return ((i << shift) ^ (i >> (32 - shift))) % 2**32


def  UpperCAmelCase__	(snake_case__							: bytes   ):
  """simple docstring"""

  _snake_case							:		Any								=  preprocess(snake_case__   )

  _snake_case							:		Optional[Any]								=  [int(2**32 * abs(sin(i + 1   )   )   ) for i in range(64   )]

  # Starting states
  _snake_case							:		Union[str, Any]								=  0x6745_2301
  _snake_case							:		List[Any]								=  0xEFCD_AB89
  _snake_case							:		Optional[Any]								=  0x98BA_DCFE
  _snake_case							:		Optional[int]								=  0x1032_5476

  _snake_case							:		Tuple								=  [
      7,
      12,
      17,
      22,
      7,
      12,
      17,
      22,
      7,
      12,
      17,
      22,
      7,
      12,
      17,
      22,
      5,
      9,
      14,
      20,
      5,
      9,
      14,
      20,
      5,
      9,
      14,
      20,
      5,
      9,
      14,
      20,
      4,
      11,
      16,
      23,
      4,
      11,
      16,
      23,
      4,
      11,
      16,
      23,
      4,
      11,
      16,
      23,
      6,
      10,
      15,
      21,
      6,
      10,
      15,
      21,
      6,
      10,
      15,
      21,
      6,
      10,
      15,
      21,
  ]

  # Process bit string in chunks, each with 16 32-char words
  for block_words in get_block_words(snake_case__   ):
    _snake_case							:		Tuple								=  aa
    _snake_case							:		str								=  ba
    _snake_case							:		int								=  ca
    _snake_case							:		Dict								=  da

    # Hash current chunk
    for i in range(64   ):
      if i <= 15:
        # f = (b & c) | (not_32(b) & d)     # Alternate definition for f
        _snake_case							:		int								=  d ^ (b & (c ^ d))
        _snake_case							:		Dict								=  i
      elif i <= 31:
        # f = (d & b) | (not_32(d) & c)     # Alternate definition for f
        _snake_case							:		Tuple								=  c ^ (d & (b ^ c))
        _snake_case							:		int								=  (5 * i + 1) % 16
      elif i <= 47:
        _snake_case							:		List[Any]								=  b ^ c ^ d
        _snake_case							:		Union[str, Any]								=  (3 * i + 5) % 16
      else:
        _snake_case							:		Tuple								=  c ^ (b | not_aa(snake_case__   ))
        _snake_case							:		Optional[int]								=  (7 * i) % 16
      _snake_case							:		Dict								=  (f + a + added_consts[i] + block_words[g]) % 2**32
      _snake_case							:		List[str]								=  d
      _snake_case							:		List[Any]								=  c
      _snake_case							:		str								=  b
      _snake_case							:		List[str]								=  sum_aa(snake_case__		,       left_rotate_aa(snake_case__		,       shift_amounts[i]   )   )

    # Add hashed chunk to running total
    _snake_case							:		Union[str, Any]								=  sum_aa(snake_case__		,       snake_case__   )
    _snake_case							:		str								=  sum_aa(snake_case__		,       snake_case__   )
    _snake_case							:		Any								=  sum_aa(snake_case__		,       snake_case__   )
    _snake_case							:		List[str]								=  sum_aa(snake_case__		,       snake_case__   )

  _snake_case							:		Any								=  reformat_hex(snake_case__   ) + reformat_hex(snake_case__   ) + reformat_hex(snake_case__   ) + reformat_hex(snake_case__   )
  return digest


if __name__ == "__main__":
     import doctest

     doctest.testmod()


