


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




      def run_func(UpperCAmelCase      :			List[str]      ):
            @wraps(UpperCAmelCase      )
            def run_in_eager_mode(*UpperCAmelCase      :			Dict						,	**UpperCAmelCase      :			int      ):
                  return func(*UpperCAmelCase						,	**UpperCAmelCase      )

            @wraps(UpperCAmelCase      )
            @tf.function(experimental_compile=UpperCAmelCase      )
            def run_in_graph_mode(*UpperCAmelCase      :			str						,	**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   _UpperCAmelCase       (     UpperCAmelCase      :			int						,	UpperCAmelCase      :			int						,	UpperCAmelCase      :			int      ):
      """simple docstring"""




      __lowerCamelCase     :      Optional[int]       						=						random.Random()
      __lowerCamelCase     :      Any       						=						[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 							_UpperCamelCase (						A						):

   '''simple docstring'''







   a_				:							TensorFlowBenchmarkArguments
   a_				:							PretrainedConfig
   a_				:							str    = "TensorFlow"

   @property
   def   _snake_case    (  self  : Optional[int]							):






         '''simple docstring'''



         return tf.__version__

   def   _snake_case    (  self  : int       ,  _lowerCamelCase  : str       ,  _lowerCamelCase  : int       ,  _lowerCamelCase  : int							):






         '''simple docstring'''



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

   def   _snake_case    (  self  : List[str]       ,  _lowerCamelCase  : str       ,  _lowerCamelCase  : int       ,  _lowerCamelCase  : int							):






         '''simple docstring'''



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

   def   _snake_case    (  self  : Optional[int]       ,  _lowerCamelCase  : str       ,  _lowerCamelCase  : int       ,  _lowerCamelCase  : int							):






         '''simple docstring'''



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

   def   _snake_case    (  self  : List[str]       ,  _lowerCamelCase  : str       ,  _lowerCamelCase  : int       ,  _lowerCamelCase  : int							):






         '''simple docstring'''



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

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

   def   _snake_case    (  self  : str       ,  _lowerCamelCase  : str       ,  _lowerCamelCase  : int       ,  _lowerCamelCase  : int							):






         '''simple docstring'''



         __lowerCamelCase     :      List[str]       						=						self.config_dict[model_name]

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

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

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

         __lowerCamelCase     :      str       						=						encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

         return _inference

   def   _snake_case    (  self  : Dict       ,  _lowerCamelCase  : str       ,  _lowerCamelCase  : int       ,  _lowerCamelCase  : int							):






         '''simple docstring'''



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

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

         # encoder-decoder has vocab size saved differently
         __lowerCamelCase     :      Any       						=						config.vocab_size if hasattr(_lowerCamelCase       ,  """vocab_size"""							) else config.encoder.vocab_size
         __lowerCamelCase     :      Any       						=						random_input_ids(_lowerCamelCase       ,  _lowerCamelCase       ,  _lowerCamelCase							)

         @run_with_tf_optimizations(self.args.eager_mode       ,  self.args.use_xla							)
         def encoder_decoder_train():
               __lowerCamelCase     :      List[Any]       						=						model(_lowerCamelCase       ,  decoder_input_ids=_lowerCamelCase       ,  labels=_lowerCamelCase       ,  training=_lowerCamelCase							)[0]
               __lowerCamelCase     :      Tuple       						=						tf.gradients(_lowerCamelCase       ,  model.trainable_variables							)
               return gradients

         @run_with_tf_optimizations(self.args.eager_mode       ,  self.args.use_xla							)
         def encoder_train():
               __lowerCamelCase     :      Union[str, Any]       						=						model(_lowerCamelCase       ,  labels=_lowerCamelCase       ,  training=_lowerCamelCase							)[0]
               __lowerCamelCase     :      Tuple       						=						tf.gradients(_lowerCamelCase       ,  model.trainable_variables							)
               return gradients

         __lowerCamelCase     :      Union[str, Any]       						=						encoder_decoder_train if config.is_encoder_decoder else encoder_train

         return _train

   def   _snake_case    (  self  : Union[str, Any]       ,  _lowerCamelCase  : Tuple							):






         '''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(_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
                     __lowerCamelCase     :      Tuple       						=						timeit.repeat(
                         _lowerCamelCase       ,  repeat=self.args.repeat       ,  number=1_0       ,  )

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


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

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







import re
import warnings
from contextlib import contextmanager

from ...processing_utils import ProcessorMixin



class 							_UpperCamelCase (						A						):

   '''simple docstring'''







   a_				:							List[Any]    = ["image_processor", "tokenizer"]
   a_				:							Tuple    = "AutoImageProcessor"
   a_				:							Tuple    = "AutoTokenizer"

   def __init__(  self  : Union[str, Any]       ,  _lowerCamelCase  : List[Any]=None       ,  _lowerCamelCase  : Any=None       ,  **_lowerCamelCase  : List[str]							):






         '''simple docstring'''



         __lowerCamelCase     :      Tuple       						=						None
         if "feature_extractor" in kwargs:
               warnings.warn(
                   """The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"""
                   """ instead."""       ,  _lowerCamelCase       ,  )
               __lowerCamelCase     :      List[Any]       						=						kwargs.pop("""feature_extractor"""							)

         __lowerCamelCase     :      Optional[Any]       						=						image_processor if image_processor is not None else feature_extractor
         if image_processor is None:
               raise ValueError("""You need to specify an `image_processor`."""							)
         if tokenizer is None:
               raise ValueError("""You need to specify a `tokenizer`."""							)

         super().__init__(_lowerCamelCase       ,  _lowerCamelCase							)
         __lowerCamelCase     :      List[Any]       						=						self.image_processor
         __lowerCamelCase     :      Any       						=						False

   def __call__(  self  : Union[str, Any]       ,  *_lowerCamelCase  : str       ,  **_lowerCamelCase  : Dict							):






         '''simple docstring'''



         if self._in_target_context_manager:
               return self.current_processor(*_lowerCamelCase       ,  **_lowerCamelCase							)

         __lowerCamelCase     :      List[str]       						=						kwargs.pop("""images"""       ,  _lowerCamelCase							)
         __lowerCamelCase     :      List[str]       						=						kwargs.pop("""text"""       ,  _lowerCamelCase							)
         if len(_lowerCamelCase							) > 0:
               __lowerCamelCase     :      Tuple       						=						args[0]
               __lowerCamelCase     :      List[Any]       						=						args[1:]

         if images is None and text is None:
               raise ValueError("""You need to specify either an `images` or `text` input to process."""							)

         if images is not None:
               __lowerCamelCase     :      Tuple       						=						self.image_processor(_lowerCamelCase       ,  *_lowerCamelCase       ,  **_lowerCamelCase							)
         if text is not None:
               __lowerCamelCase     :      int       						=						self.tokenizer(_lowerCamelCase       ,  **_lowerCamelCase							)

         if text is None:
               return inputs
         elif images is None:
               return encodings
         else:
               __lowerCamelCase     :      Optional[int]       						=						encodings["""input_ids"""]
               return inputs

   def   _snake_case    (  self  : List[Any]       ,  *_lowerCamelCase  : Tuple       ,  **_lowerCamelCase  : Union[str, Any]							):






         '''simple docstring'''



         return self.tokenizer.batch_decode(*_lowerCamelCase       ,  **_lowerCamelCase							)

   def   _snake_case    (  self  : List[str]       ,  *_lowerCamelCase  : Optional[int]       ,  **_lowerCamelCase  : int							):






         '''simple docstring'''



         return self.tokenizer.decode(*_lowerCamelCase       ,  **_lowerCamelCase							)

   @contextmanager
   def   _snake_case    (  self  : Optional[Any]							):






         '''simple docstring'''



         warnings.warn(
             """`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your """
             """labels by using the argument `text` of the regular `__call__` method (either in the same call as """
             """your images inputs, or in a separate call."""							)
         __lowerCamelCase     :      Any       						=						True
         __lowerCamelCase     :      int       						=						self.tokenizer
         yield
         __lowerCamelCase     :      Dict       						=						self.image_processor
         __lowerCamelCase     :      List[Any]       						=						False

   def   _snake_case    (  self  : List[str]       ,  _lowerCamelCase  : str       ,  _lowerCamelCase  : List[str]=False       ,  _lowerCamelCase  : List[Any]=None							):






         '''simple docstring'''



         if added_vocab is None:
               __lowerCamelCase     :      List[str]       						=						self.tokenizer.get_added_vocab()

         __lowerCamelCase     :      List[Any]       						=						{}

         while tokens:
               __lowerCamelCase     :      Dict       						=						re.search(R"""<s_(.*?)>"""       ,  _lowerCamelCase       ,  re.IGNORECASE							)
               if start_token is None:
                     break
               __lowerCamelCase     :      Dict       						=						start_token.group(1							)
               __lowerCamelCase     :      List[str]       						=						re.search(RF"""</s_{key}>"""       ,  _lowerCamelCase       ,  re.IGNORECASE							)
               __lowerCamelCase     :      Optional[int]       						=						start_token.group()
               if end_token is None:
                     __lowerCamelCase     :      int       						=						tokens.replace(_lowerCamelCase       ,  """"""							)
               else:
                     __lowerCamelCase     :      Dict       						=						end_token.group()
                     __lowerCamelCase     :      Optional[Any]       						=						re.escape(_lowerCamelCase							)
                     __lowerCamelCase     :      List[str]       						=						re.escape(_lowerCamelCase							)
                     __lowerCamelCase     :      Any       						=						re.search(F"""{start_token_escaped}(.*?){end_token_escaped}"""       ,  _lowerCamelCase       ,  re.IGNORECASE							)
                     if content is not None:
                           __lowerCamelCase     :      Union[str, Any]       						=						content.group(1							).strip()
                           if r"<s_" in content and r"</s_" in content:  # non-leaf node
                                 __lowerCamelCase     :      Union[str, Any]       						=						self.tokenajson(_lowerCamelCase       ,  is_inner_value=_lowerCamelCase       ,  added_vocab=_lowerCamelCase							)
                                 if value:
                                       if len(_lowerCamelCase							) == 1:
                                             __lowerCamelCase     :      Optional[int]       						=						value[0]
                                       __lowerCamelCase     :      Union[str, Any]       						=						value
                           else:  # leaf nodes
                                 __lowerCamelCase     :      List[Any]       						=						[]
                                 for leaf in content.split(R"""<sep/>"""							):
                                       __lowerCamelCase     :      Optional[Any]       						=						leaf.strip()
                                       if leaf in added_vocab and leaf[0] == "<" and leaf[-2:] == "/>":
                                             __lowerCamelCase     :      str       						=						leaf[1:-2]  # for categorical special tokens
                                       output[key].append(_lowerCamelCase							)
                                 if len(output[key]							) == 1:
                                       __lowerCamelCase     :      Tuple       						=						output[key][0]

                     __lowerCamelCase     :      Tuple       						=						tokens[tokens.find(_lowerCamelCase							) + len(_lowerCamelCase							) :].strip()
                     if tokens[:6] == r"<sep/>":  # non-leaf nodes
                           return [output] + self.tokenajson(tokens[6:]       ,  is_inner_value=_lowerCamelCase       ,  added_vocab=_lowerCamelCase							)

         if len(_lowerCamelCase							):
               return [output] if is_inner_value else output
         else:
               return [] if is_inner_value else {"text_sequence": tokens}

   @property
   def   _snake_case    (  self  : Tuple							):






         '''simple docstring'''



         warnings.warn(
             """`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."""       ,  _lowerCamelCase       ,  )
         return self.image_processor_class


   @property
   def   _snake_case    (  self  : List[Any]							):






         '''simple docstring'''



         warnings.warn(
             """`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."""       ,  _lowerCamelCase       ,  )
         return self.image_processor



