





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

__lowerCamelCase	:					List[Any]			  =					logging.get_logger(__name__)



def 		A_    (					_lowerCAmelCase							,       _lowerCAmelCase					)						->			Union[str, Any]:
 def run_func(_lowerCAmelCase					):
  @wraps(_lowerCAmelCase					)
  def run_in_eager_mode(*_lowerCAmelCase							,       **_lowerCAmelCase					):
   return func(*_lowerCAmelCase							,       **_lowerCAmelCase					)

  @wraps(_lowerCAmelCase					)
  @tf.function(experimental_compile=_lowerCAmelCase					)
  def run_in_graph_mode(*_lowerCAmelCase							,       **_lowerCAmelCase					):
   return func(*_lowerCAmelCase							,       **_lowerCAmelCase					)

  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_    (					_lowerCAmelCase							,       _lowerCAmelCase							,       _lowerCAmelCase					)						->			["tf.Tensor"]:
 UpperCamelCase     : Optional[int]	     =       random.Random()
 UpperCamelCase     : int	     =       [rng.randint(0							,       vocab_size - 1					) for i in range(batch_size * sequence_length					)]
 return tf.constant(_lowerCAmelCase							,       shape=(batch_size, sequence_length)							,       dtype=tf.intaa					)



class  A__       (	__snake_case      ):
      _UpperCAmelCase     :TensorFlowBenchmarkArguments
      _UpperCAmelCase     :PretrainedConfig
      _UpperCAmelCase     :str						   =       "TensorFlow"





      @property
      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       return tf.__version__





      def 	__UpperCamelCase(			self	,							A_	,							A_	,							A_       ):
       '''simple docstring'''




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





      def 	__UpperCamelCase(			self	,							A_	,							A_	,							A_       ):
       '''simple docstring'''




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





      def 	__UpperCamelCase(			self	,							A_	,							A_	,							A_       ):
       '''simple docstring'''




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





      def 	__UpperCamelCase(			self	,							A_	,							A_	,							A_       ):
       '''simple docstring'''




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

       UpperCamelCase     : Any	     =       self._prepare_train_func(A_	,							A_	,							A_       )
       return self._measure_memory(_train       )





      def 	__UpperCamelCase(			self	,							A_	,							A_	,							A_       ):
       '''simple docstring'''




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

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

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

       # encoder-decoder has vocab size saved differently
       UpperCamelCase     : Union[str, Any]	     =       config.vocab_size if hasattr(A_	,							"vocab_size"       ) else config.encoder.vocab_size
       UpperCamelCase     : Optional[Any]	     =       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_       )

       UpperCamelCase     : Tuple	     =       encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

       return _inference





      def 	__UpperCamelCase(			self	,							A_	,							A_	,							A_       ):
       '''simple docstring'''




       UpperCamelCase     : str	     =       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."       )

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

       # encoder-decoder has vocab size saved differently
       UpperCamelCase     : Optional[Any]	     =       config.vocab_size if hasattr(A_	,							"vocab_size"       ) else config.encoder.vocab_size
       UpperCamelCase     : str	     =       random_input_ids(A_	,							A_	,							A_       )

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

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

       UpperCamelCase     : Tuple	     =       encoder_decoder_train if config.is_encoder_decoder else encoder_train

       return _train





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

         return memory, summary
        except ResourceExhaustedError as e:
         self.print_fn(F"""Doesn't fit on GPU. {e}"""       )
         return "N/A", None







__lowerCamelCase	:					Tuple			  =					"""ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"""



def 		A_    (					_lowerCAmelCase					)						->			bytes:
 # Make sure the supplied data is a bytes-like object
 if not isinstance(_lowerCAmelCase							,       _lowerCAmelCase					):
  UpperCamelCase     : str	     =       F"""a bytes-like object is required, not '{data.__class__.__name__}'"""
  raise TypeError(_lowerCAmelCase					)

 UpperCamelCase     : Optional[int]	     =       "".join(bin(_lowerCAmelCase					)[2:].zfill(8					) for byte in data					)

 UpperCamelCase     : str	     =       len(_lowerCAmelCase					) % 6 != 0

 if padding_needed:
  # The padding that will be added later
  UpperCamelCase     : Optional[Any]	     =       b"=" * ((6 - len(_lowerCAmelCase					) % 6) // 2)

  # Append binary_stream with arbitrary binary digits (0's by default) to make its
  # length a multiple of 6.
  binary_stream += "0" * (6 - len(_lowerCAmelCase					) % 6)
 else:
  UpperCamelCase     : List[Any]	     =       b""

 # Encode every 6 binary digits to their corresponding Base64 character
 return (
     "".join(
         B64_CHARSET[int(binary_stream[index : index + 6]							,       2					)]
         for index in range(0							,       len(_lowerCAmelCase					)							,       6					)					).encode()
     + padding
 )



def 		A_    (					_lowerCAmelCase					)						->			bytes:
 # Make sure encoded_data is either a string or a bytes-like object
 if not isinstance(_lowerCAmelCase							,       _lowerCAmelCase					) and not isinstance(_lowerCAmelCase							,       _lowerCAmelCase					):
  UpperCamelCase     : List[Any]	     =       (
      "argument should be a bytes-like object or ASCII string, "
      F"""not '{encoded_data.__class__.__name__}'"""
  )
  raise TypeError(_lowerCAmelCase					)

 # In case encoded_data is a bytes-like object, make sure it contains only
 # ASCII characters so we convert it to a string object
 if isinstance(_lowerCAmelCase							,       _lowerCAmelCase					):
  try:
   UpperCamelCase     : Any	     =       encoded_data.decode("utf-8"					)
  except UnicodeDecodeError:
   raise ValueError("base64 encoded data should only contain ASCII characters"					)

 UpperCamelCase     : Union[str, Any]	     =       encoded_data.count("="					)

 # Check if the encoded string contains non base64 characters
 if padding:
  assert all(
      char in B64_CHARSET for char in encoded_data[:-padding]					), "Invalid base64 character(s) found."
 else:
  assert all(
      char in B64_CHARSET for char in encoded_data					), "Invalid base64 character(s) found."

 # Check the padding
 assert len(_lowerCAmelCase					) % 4 == 0 and padding < 3, "Incorrect padding"

 if padding:
  # Remove padding if there is one
  UpperCamelCase     : List[str]	     =       encoded_data[:-padding]

  UpperCamelCase     : List[Any]	     =       "".join(
      bin(B64_CHARSET.index(_lowerCAmelCase					)					)[2:].zfill(6					) for char in encoded_data					)[: -padding * 2]
 else:
  UpperCamelCase     : List[Any]	     =       "".join(
      bin(B64_CHARSET.index(_lowerCAmelCase					)					)[2:].zfill(6					) for char in encoded_data					)

 UpperCamelCase     : Any	     =       [
     int(binary_stream[index : index + 8]							,       2					)
     for index in range(0							,       len(_lowerCAmelCase					)							,       8					)
 ]

 return bytes(_lowerCAmelCase					)


if __name__ == "__main__":
    import doctest

    doctest.testmod()
