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

SCREAMING_SNAKE_CASE_    					=						logging.get_logger(__name__)




def     lowercase__						(					lowerCAmelCase				:					bool ,				lowerCAmelCase				:					bool ) ->      List[Any]:






    """simple docstring"""




    def run_func(lowerCAmelCase				:					List[Any] ):
        @wraps(lowerCAmelCase )
        def run_in_eager_mode(*lowerCAmelCase				:					Optional[Any] ,				**lowerCAmelCase				:					Dict ):
            return func(*lowerCAmelCase ,				**lowerCAmelCase )

        @wraps(lowerCAmelCase )
        @tf.function(experimental_compile=lowerCAmelCase )
        def run_in_graph_mode(*lowerCAmelCase				:					int ,				**lowerCAmelCase				:					str ):
            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     lowercase__						(					lowerCAmelCase				:					int ,				lowerCAmelCase				:					int ,				lowerCAmelCase				:					int ) ->      List[Any]:






    """simple docstring"""




    UpperCAmelCase									=					random.Random()
    UpperCAmelCase									=					[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      _UpperCAmelCase							(		__lowerCAmelCase							):
      __SCREAMING_SNAKE_CASE   :							TensorFlowBenchmarkArguments
      __SCREAMING_SNAKE_CASE   :							PretrainedConfig
      __SCREAMING_SNAKE_CASE   :							str							      =	"TensorFlow"




      @property
      def     a_      (							self							)						->					List[str]:
          return tf.__version__




      def     a_      (							self   ,       lowercase_   ,       lowercase_   ,       lowercase_							)						->					float:
          UpperCAmelCase									=					self.args.strategy
          if strategy is None:
              raise ValueError('A device strategy has to be initialized before using TensorFlow.'							)
          UpperCAmelCase									=					self._prepare_inference_func(lowerCamelCase__   ,       lowerCamelCase__   ,       lowerCamelCase__							)
          return self._measure_speed(_inference							)




      def     a_      (							self   ,       lowercase_   ,       lowercase_   ,       lowercase_							)						->					float:
          UpperCAmelCase									=					self.args.strategy
          if strategy is None:
              raise ValueError('A device strategy has to be initialized before using TensorFlow.'							)
          UpperCAmelCase									=					self._prepare_train_func(lowerCamelCase__   ,       lowerCamelCase__   ,       lowerCamelCase__							)
          return self._measure_speed(_train							)




      def     a_      (							self   ,       lowercase_   ,       lowercase_   ,       lowercase_							)						->					[Memory, Optional[MemorySummary]]:






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




      def     a_      (							self   ,       lowercase_   ,       lowercase_   ,       lowercase_							)						->					[Memory, Optional[MemorySummary]]:






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

          UpperCAmelCase									=					self._prepare_train_func(lowerCamelCase__   ,       lowerCamelCase__   ,       lowerCamelCase__							)
          return self._measure_memory(_train							)




      def     a_      (							self   ,       lowercase_   ,       lowercase_   ,       lowercase_							)						->					Callable[[], None]:
          UpperCAmelCase									=					self.config_dict[model_name]

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

          UpperCAmelCase									=					(
              hasattr(lowerCamelCase__   ,       'architectures'							)
              and isinstance(config.architectures   ,       lowerCamelCase__							)
              and len(config.architectures							) > 0
          )
          if not self.args.only_pretrain_model and has_model_class_in_config:
              try:
                  UpperCAmelCase									=					'''TF''' + config.architectures[0]  # prepend 'TF' for tensorflow model
                  UpperCAmelCase									=					__import__('transformers'   ,       fromlist=[model_class]							)
                  UpperCAmelCase									=					getattr(lowerCamelCase__   ,       lowerCamelCase__							)
                  UpperCAmelCase									=					model_cls(lowerCamelCase__							)
              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									=					TF_MODEL_MAPPING[config.__class__](lowerCamelCase__							)

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

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

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

          UpperCAmelCase									=					encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

          return _inference




      def     a_      (							self   ,       lowercase_   ,       lowercase_   ,       lowercase_							)						->					Callable[[], None]:
          UpperCAmelCase									=					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									=					(
              hasattr(lowerCamelCase__   ,       'architectures'							)
              and isinstance(config.architectures   ,       lowerCamelCase__							)
              and len(config.architectures							) > 0
          )
          if not self.args.only_pretrain_model and has_model_class_in_config:
              try:
                  UpperCAmelCase									=					'''TF''' + config.architectures[0]  # prepend 'TF' for tensorflow model
                  UpperCAmelCase									=					__import__('transformers'   ,       fromlist=[model_class]							)
                  UpperCAmelCase									=					getattr(lowerCamelCase__   ,       lowerCamelCase__							)
                  UpperCAmelCase									=					model_cls(lowerCamelCase__							)
              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									=					TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](lowerCamelCase__							)

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

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

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

          UpperCAmelCase									=					encoder_decoder_train if config.is_encoder_decoder else encoder_train

          return _train




      def     a_      (							self   ,       lowercase_							)						->					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(lowerCamelCase__   ,       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									=					timeit.repeat(
                      lowerCamelCase__   ,       repeat=self.args.repeat   ,       number=1_0   ,       )

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




      def     a_      (							self   ,       lowercase_							)						->					[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.'							)
                      UpperCAmelCase									=					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									=					'''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									=					nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx							)
                          UpperCAmelCase									=					nvml.nvmlDeviceGetMemoryInfo(lowerCamelCase__							)
                          UpperCAmelCase									=					meminfo.used
                          UpperCAmelCase									=					Memory(lowerCamelCase__							)
                          # 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									=					None
                      else:
                          UpperCAmelCase									=					measure_peak_memory_cpu(lowerCamelCase__							)
                          UpperCAmelCase									=					Memory(lowerCamelCase__							) if isinstance(lowerCamelCase__   ,       lowerCamelCase__							) else memory_bytes
                  if self.args.trace_memory_line_by_line:
                      UpperCAmelCase									=					stop_memory_tracing(lowerCamelCase__							)
                      if memory is None:
                          UpperCAmelCase									=					summary.total
                  else:
                      UpperCAmelCase									=					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 os
from typing import Optional

import fsspec
from fsspec.archive import AbstractArchiveFileSystem
from fsspec.utils import DEFAULT_BLOCK_SIZE
class      _UpperCAmelCase							(		SCREAMING_SNAKE_CASE_							):
      __SCREAMING_SNAKE_CASE   :							str							      =	""
      __SCREAMING_SNAKE_CASE   :							str							      =	(
          None  # protocol passed in prefix to the url. ex: "gzip", for gzip://file.txt::http://foo.bar/file.txt.gz
      )
      __SCREAMING_SNAKE_CASE   :							str							      =	None  # compression type in fsspec. ex: "gzip"
      __SCREAMING_SNAKE_CASE   :							str							      =	None  # extension of the filename to strip. ex: "".gz" to get file.txt from file.txt.gz




      def __init__(							self   ,       lowercase_ = ""   ,       lowercase_ = None   ,       lowercase_ = None   ,       **lowercase_							)						->					str:
          super().__init__(self   ,       **lowercase_							)
          # always open as "rb" since fsspec can then use the TextIOWrapper to make it work for "r" mode
          UpperCAmelCase									=					fsspec.open(
              lowercase_   ,       mode='rb'   ,       protocol=lowercase_   ,       compression=self.compression   ,       client_kwargs={
                  'requote_redirect_url': False,  # see https://github.com/huggingface/datasets/pull/5459
                  'trust_env': True,  # Enable reading proxy env variables.
                  **(target_options or {}).pop('client_kwargs'   ,       {}							),  # To avoid issues if it was already passed.
              }   ,       **(target_options or {})   ,       )
          UpperCAmelCase									=					os.path.basename(self.file.path.split('::'							)[0]							)
          UpperCAmelCase									=					(
              self.compressed_name[: self.compressed_name.rindex('.'							)]
              if '.' in self.compressed_name
              else self.compressed_name
          )
          UpperCAmelCase									=					None




      @classmethod
      def     a_      (							cls   ,       lowercase_							)						->					Union[str, Any]:
          # compressed file paths are always relative to the archive root
          return super()._strip_protocol(lowercase_							).lstrip('/'							)




      def     a_      (							self							)						->					int:
          if self.dir_cache is None:
              UpperCAmelCase									=					{**self.file.fs.info(self.file.path							), 'name': self.uncompressed_name}
              UpperCAmelCase									=					{f['name']: f}




      def     a_      (							self   ,       lowercase_							)						->					Any:
          return self.file.open().read()




      def     a_      (							self   ,       lowercase_   ,       lowercase_ = "rb"   ,       lowercase_=None   ,       lowercase_=True   ,       lowercase_=None   ,       **lowercase_   ,       )						->					Optional[Any]:
          UpperCAmelCase									=					self._strip_protocol(lowercase_							)
          if mode != "rb":
              raise ValueError(F"Tried to read with mode {mode} on file {self.file.path} opened with mode 'rb'"							)
          return self.file.open()

class      _UpperCAmelCase							(		SCREAMING_SNAKE_CASE_							):
      __SCREAMING_SNAKE_CASE   :							List[Any]							      =	"bz2"
      __SCREAMING_SNAKE_CASE   :							Union[str, Any]							      =	"bz2"
      __SCREAMING_SNAKE_CASE   :							Union[str, Any]							      =	".bz2"

class      _UpperCAmelCase							(		SCREAMING_SNAKE_CASE_							):
      __SCREAMING_SNAKE_CASE   :							List[str]							      =	"gzip"
      __SCREAMING_SNAKE_CASE   :							Any							      =	"gzip"
      __SCREAMING_SNAKE_CASE   :							List[str]							      =	".gz"

class      _UpperCAmelCase							(		SCREAMING_SNAKE_CASE_							):
      __SCREAMING_SNAKE_CASE   :							List[str]							      =	"lz4"
      __SCREAMING_SNAKE_CASE   :							List[Any]							      =	"lz4"
      __SCREAMING_SNAKE_CASE   :							Tuple							      =	".lz4"

class      _UpperCAmelCase							(		SCREAMING_SNAKE_CASE_							):
      __SCREAMING_SNAKE_CASE   :							List[str]							      =	"xz"
      __SCREAMING_SNAKE_CASE   :							Union[str, Any]							      =	"xz"
      __SCREAMING_SNAKE_CASE   :							List[Any]							      =	".xz"

class      _UpperCAmelCase							(		SCREAMING_SNAKE_CASE_							):
      __SCREAMING_SNAKE_CASE   :							List[str]							      =	"zstd"
      __SCREAMING_SNAKE_CASE   :							int							      =	"zstd"
      __SCREAMING_SNAKE_CASE   :							Tuple							      =	".zst"




      def __init__(							self   ,       lowercase_   ,       lowercase_ = "rb"   ,       lowercase_ = None   ,       lowercase_ = None   ,       lowercase_ = DEFAULT_BLOCK_SIZE   ,       **lowercase_   ,       )						->					str:
          super().__init__(
              fo=lowercase_   ,       mode=lowercase_   ,       target_protocol=lowercase_   ,       target_options=lowercase_   ,       block_size=lowercase_   ,       **lowercase_   ,       )
          # We need to wrap the zstd decompressor to avoid this error in fsspec==2021.7.0 and zstandard==0.15.2:
          #
          # File "/Users/user/.virtualenvs/hf-datasets/lib/python3.7/site-packages/fsspec/core.py", line 145, in open
          #     out.close = close
          # AttributeError: 'zstd.ZstdDecompressionReader' object attribute 'close' is read-only
          #
          # see https://github.com/intake/filesystem_spec/issues/725
          UpperCAmelCase									=					self.file.__enter__
          class      _UpperCAmelCase							:




                def __init__(							self   ,       lowercase_							)						->					List[Any]:
                    UpperCAmelCase									=					file_




                def __enter__(							self							)						->					List[str]:
                    self._file.__enter__()
                    return self




                def __exit__(							self   ,       *lowercase_   ,       **lowercase_							)						->					List[str]:
                    self._file.__exit__(*lowercase_   ,       **lowercase_							)




                def __iter__(							self							)						->					Optional[int]:
                    return iter(self._file							)




                def     a_      (							self							)						->					Tuple:
                    return next(self._file							)




                def __getattr__(							self   ,       lowercase_							)						->					Optional[int]:
                    return getattr(self._file   ,       lowercase_							)

          def fixed_enter(*lowercase_   ,       **lowercase_							):
              return WrappedFile(_enter(*lowercase_   ,       **lowercase_							)							)

          UpperCAmelCase									=					fixed_enter



