



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         =			logging.get_logger(__name__)







def    _snake_case (       __snake_case					,  __snake_case    )     ->     Union[str, Any]:






      '''simple docstring'''







      def run_func(__snake_case    ):
            @wraps(__snake_case    )
            def run_in_eager_mode(*__snake_case					,  **__snake_case    ):
                  return func(*__snake_case					,  **__snake_case    )

            @wraps(__snake_case    )
            @tf.function(experimental_compile=__snake_case    )
            def run_in_graph_mode(*__snake_case					,  **__snake_case    ):
                  return func(*__snake_case					,  **__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    _snake_case (       __snake_case					,  __snake_case					,  __snake_case    )     ->     ["tf.Tensor"]:






      '''simple docstring'''







      UpperCAmelCase_  :	List[str]				      =		random.Random()
      UpperCAmelCase_  :	Any				      =		[rng.randint(0					,  vocab_size - 1    ) for i in range(batch_size * sequence_length    )]
      return tf.constant(__snake_case					,  shape=(batch_size, sequence_length)					,  dtype=tf.intaa    )





class   snake_case_ (lowercase__  ):




       """simple docstring"""






       _lowerCamelCase			    =				42
       _lowerCamelCase			    =				42
       _lowerCamelCase			    =				"TensorFlow"



       @property
       def  A_   (     self):

             """simple docstring"""





             return tf.__version__



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	int				      =		self.args.strategy
             if strategy is None:
                   raise ValueError("A device strategy has to be initialized before using TensorFlow.")
             UpperCAmelCase_  :	Tuple				      =		self._prepare_inference_func(lowercase     ,lowercase     ,lowercase)
             return self._measure_speed(_inference)



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	Union[str, Any]				      =		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(lowercase     ,lowercase     ,lowercase)
             return self._measure_speed(_train)



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase):

             """simple docstring"""





             if self.args.is_gpu:
                   tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]     ,lowercase)
             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(lowercase     ,lowercase     ,lowercase)
             return self._measure_memory(_inference)



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase):

             """simple docstring"""





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

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



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	Dict				      =		self.config_dict[model_name]

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

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

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

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

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

             UpperCAmelCase_  :	Tuple				      =		encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

             return _inference



       def  A_   (     self     ,lowercase     ,lowercase     ,lowercase):

             """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_  :	Dict				      =		(
                 hasattr(lowercase     ,"architectures")
                 and isinstance(config.architectures     ,lowercase)
                 and len(config.architectures) > 0
             )
             if not self.args.only_pretrain_model and has_model_class_in_config:
                   try:
                         UpperCAmelCase_  :	Dict				      =		"TF" + config.architectures[0]  # prepend 'TF' for tensorflow model
                         UpperCAmelCase_  :	str				      =		__import__("transformers"     ,fromlist=[model_class])
                         UpperCAmelCase_  :	Tuple				      =		getattr(lowercase     ,lowercase)
                         UpperCAmelCase_  :	List[Any]				      =		model_cls(lowercase)
                   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[str]				      =		TF_MODEL_WITH_LM_HEAD_MAPPING[config.__class__](lowercase)

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

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

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

             UpperCAmelCase_  :	Optional[int]				      =		encoder_decoder_train if config.is_encoder_decoder else encoder_train

             return _train



       def  A_   (     self     ,lowercase):

             """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(lowercase     ,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_  :	Optional[int]				      =		timeit.repeat(
                             lowercase     ,repeat=self.args.repeat     ,number=10     ,)

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



       def  A_   (     self     ,lowercase):

             """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_  :	Tuple				      =		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_  :	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()
                                     UpperCAmelCase_  :	Optional[int]				      =		nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx)
                                     UpperCAmelCase_  :	Optional[int]				      =		nvml.nvmlDeviceGetMemoryInfo(lowercase)
                                     UpperCAmelCase_  :	Optional[Any]				      =		meminfo.used
                                     UpperCAmelCase_  :	Optional[int]				      =		Memory(lowercase)
                                     # 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_  :	List[Any]				      =		None
                               else:
                                     UpperCAmelCase_  :	Optional[int]				      =		measure_peak_memory_cpu(lowercase)
                                     UpperCAmelCase_  :	List[str]				      =		Memory(lowercase) if isinstance(lowercase     ,lowercase) else memory_bytes
                         if self.args.trace_memory_line_by_line:
                               UpperCAmelCase_  :	Any				      =		stop_memory_tracing(lowercase)
                               if memory is None:
                                     UpperCAmelCase_  :	Optional[int]				      =		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





from ...configuration_utils import PretrainedConfig
from ...utils import logging


__lowerCamelCase         =			logging.get_logger(__name__)

__lowerCamelCase         =			{
    '''microsoft/trocr-base-handwritten''': (
        '''https://huggingface.co/microsoft/trocr-base-handwritten/resolve/main/config.json'''
    ),
    # See all TrOCR models at https://huggingface.co/models?filter=trocr
}


class   snake_case_ (lowercase__  ):




       """simple docstring"""






       _lowerCamelCase			    =				"""trocr"""
       _lowerCamelCase			    =				["""past_key_values"""]
       _lowerCamelCase			    =				{
           """num_attention_heads""": """decoder_attention_heads""",
           """hidden_size""": """d_model""",
           """num_hidden_layers""": """decoder_layers""",
       }



       def __init__(     self     ,lowercase=50265     ,lowercase=1024     ,lowercase=12     ,lowercase=16     ,lowercase=4096     ,lowercase="gelu"     ,lowercase=512     ,lowercase=0.1     ,lowercase=0.0     ,lowercase=0.0     ,lowercase=2     ,lowercase=0.02     ,lowercase=0.0     ,lowercase=True     ,lowercase=False     ,lowercase=True     ,lowercase=True     ,lowercase=1     ,lowercase=0     ,lowercase=2     ,**lowercase     ,):

             """simple docstring"""





             UpperCAmelCase_  :	int				      =		vocab_size
             UpperCAmelCase_  :	Tuple				      =		d_model
             UpperCAmelCase_  :	List[str]				      =		decoder_layers
             UpperCAmelCase_  :	Union[str, Any]				      =		decoder_attention_heads
             UpperCAmelCase_  :	Optional[Any]				      =		decoder_ffn_dim
             UpperCAmelCase_  :	Tuple				      =		activation_function
             UpperCAmelCase_  :	List[Any]				      =		max_position_embeddings
             UpperCAmelCase_  :	List[Any]				      =		dropout
             UpperCAmelCase_  :	List[str]				      =		attention_dropout
             UpperCAmelCase_  :	Union[str, Any]				      =		activation_dropout
             UpperCAmelCase_  :	Optional[int]				      =		init_std
             UpperCAmelCase_  :	Dict				      =		decoder_layerdrop
             UpperCAmelCase_  :	List[str]				      =		use_cache
             UpperCAmelCase_  :	Dict				      =		scale_embedding
             UpperCAmelCase_  :	Dict				      =		use_learned_position_embeddings
             UpperCAmelCase_  :	Optional[Any]				      =		layernorm_embedding

             super().__init__(
                 pad_token_id=lowercase     ,bos_token_id=lowercase     ,eos_token_id=lowercase     ,decoder_start_token_id=lowercase     ,**lowercase     ,)
