

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





def      snake_case_				(__A					:						bool ,       __A					:						bool					)      ->     Union[str, Any]:
   def run_func(__A					:						Dict					):
      @wraps(__A					)
      def run_in_eager_mode(*__A					:						List[str] ,       **__A					:						Union[str, Any]					):
         return func(*__A ,       **__A					)

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

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



class SCREAMING_SNAKE_CASE     (  a_			):
  """simple docstring"""






  lowerCamelCase   :      TensorFlowBenchmarkArguments
  lowerCamelCase   :      PretrainedConfig
  lowerCamelCase   :      str										="TensorFlow"







  @property
  def 	SCREAMING_SNAKE_CASE   (      self			:   Optional[int]   )							->     str:






     """simple docstring"""







     return tf.__version__







  def 	SCREAMING_SNAKE_CASE   (      self			:   Optional[Any]	,							lowerCAmelCase			:   str	,							lowerCAmelCase			:   int	,							lowerCAmelCase			:   int   )							->     float:






     """simple docstring"""







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







  def 	SCREAMING_SNAKE_CASE   (      self			:   Optional[Any]	,							lowerCAmelCase			:   str	,							lowerCAmelCase			:   int	,							lowerCAmelCase			:   int   )							->     float:






     """simple docstring"""







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







  def 	SCREAMING_SNAKE_CASE   (      self			:   Any	,							lowerCAmelCase			:   str	,							lowerCAmelCase			:   int	,							lowerCAmelCase			:   int   )							->     [Memory, Optional[MemorySummary]]:






     """simple docstring"""







     if self.args.is_gpu:
        tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]	,							lowerCAmelCase   )
     __lowerCAmelCase  :      Tuple    =						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 	SCREAMING_SNAKE_CASE   (      self			:   Dict	,							lowerCAmelCase			:   str	,							lowerCAmelCase			:   int	,							lowerCAmelCase			:   int   )							->     [Memory, Optional[MemorySummary]]:






     """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  :      Any    =						self._prepare_train_func(lowerCAmelCase	,							lowerCAmelCase	,							lowerCAmelCase   )
     return self._measure_memory(_train   )







  def 	SCREAMING_SNAKE_CASE   (      self			:   List[Any]	,							lowerCAmelCase			:   str	,							lowerCAmelCase			:   int	,							lowerCAmelCase			:   int   )							->     Callable[[], None]:






     """simple docstring"""







     __lowerCAmelCase  :      int    =						self.config_dict[model_name]

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

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

     # encoder-decoder has vocab size saved differently
     __lowerCAmelCase  :      int    =						config.vocab_size if hasattr(lowerCAmelCase	,							"""vocab_size"""   ) else config.encoder.vocab_size
     __lowerCAmelCase  :      str    =						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  :      Dict    =						encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

     return _inference







  def 	SCREAMING_SNAKE_CASE   (      self			:   int	,							lowerCAmelCase			:   str	,							lowerCAmelCase			:   int	,							lowerCAmelCase			:   int   )							->     Callable[[], None]:






     """simple docstring"""







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

     # encoder-decoder has vocab size saved differently
     __lowerCAmelCase  :      List[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_train():
        __lowerCAmelCase  :      List[Any]    =						model(lowerCAmelCase	,							decoder_input_ids=lowerCAmelCase	,							labels=lowerCAmelCase	,							training=lowerCAmelCase   )[0]
        __lowerCAmelCase  :      Optional[Any]    =						tf.gradients(lowerCAmelCase	,							model.trainable_variables   )
        return gradients

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

     __lowerCAmelCase  :      int    =						encoder_decoder_train if config.is_encoder_decoder else encoder_train

     return _train







  def 	SCREAMING_SNAKE_CASE   (      self			:   Dict	,							lowerCAmelCase			:   Any   )							->     float:






     """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  :      Union[str, Any]    =						timeit.repeat(
               lowerCAmelCase	,							repeat=self.args.repeat	,							number=10	,							)

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







  def 	SCREAMING_SNAKE_CASE   (      self			:   Union[str, Any]	,							lowerCAmelCase			:   Callable[[], None]   )							->     [Memory, MemorySummary]:






     """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  :      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  :      int    =						nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx   )
                 __lowerCAmelCase  :      Union[str, Any]    =						nvml.nvmlDeviceGetMemoryInfo(lowerCAmelCase   )
                 __lowerCAmelCase  :      Optional[Any]    =						meminfo.used
                 __lowerCAmelCase  :      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  :      int    =						None
              else:
                 __lowerCAmelCase  :      str    =						measure_peak_memory_cpu(lowerCAmelCase   )
                 __lowerCAmelCase  :      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  :      List[str]    =						None

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





import unittest

import numpy as np

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
   from transformers.models.bert.modeling_flax_bert import (
       FlaxBertForMaskedLM,
       FlaxBertForMultipleChoice,
       FlaxBertForNextSentencePrediction,
       FlaxBertForPreTraining,
       FlaxBertForQuestionAnswering,
       FlaxBertForSequenceClassification,
       FlaxBertForTokenClassification,
       FlaxBertModel,
   )



class SCREAMING_SNAKE_CASE     (  unittest.TestCase			):
  """simple docstring"""






  def __init__(      self			:   Tuple	,							lowerCAmelCase			:   Tuple	,							lowerCAmelCase			:   str=13	,							lowerCAmelCase			:   Optional[Any]=7	,							lowerCAmelCase			:   Optional[Any]=True	,							lowerCAmelCase			:   Any=True	,							lowerCAmelCase			:   str=True	,							lowerCAmelCase			:   Optional[Any]=True	,							lowerCAmelCase			:   Dict=99	,							lowerCAmelCase			:   Any=32	,							lowerCAmelCase			:   int=5	,							lowerCAmelCase			:   Tuple=4	,							lowerCAmelCase			:   Optional[Any]=37	,							lowerCAmelCase			:   str="gelu"	,							lowerCAmelCase			:   str=0.1	,							lowerCAmelCase			:   str=0.1	,							lowerCAmelCase			:   Any=5_12	,							lowerCAmelCase			:   Optional[Any]=16	,							lowerCAmelCase			:   Dict=2	,							lowerCAmelCase			:   Tuple=0.02	,							lowerCAmelCase			:   Optional[int]=4	,							)							->     List[Any]:






     """simple docstring"""







     __lowerCAmelCase  :      int    =						parent
     __lowerCAmelCase  :      Dict    =						batch_size
     __lowerCAmelCase  :      Dict    =						seq_length
     __lowerCAmelCase  :      Union[str, Any]    =						is_training
     __lowerCAmelCase  :      List[Any]    =						use_attention_mask
     __lowerCAmelCase  :      List[Any]    =						use_token_type_ids
     __lowerCAmelCase  :      Optional[int]    =						use_labels
     __lowerCAmelCase  :      str    =						vocab_size
     __lowerCAmelCase  :      Any    =						hidden_size
     __lowerCAmelCase  :      Optional[int]    =						num_hidden_layers
     __lowerCAmelCase  :      Optional[int]    =						num_attention_heads
     __lowerCAmelCase  :      Dict    =						intermediate_size
     __lowerCAmelCase  :      Tuple    =						hidden_act
     __lowerCAmelCase  :      Dict    =						hidden_dropout_prob
     __lowerCAmelCase  :      Any    =						attention_probs_dropout_prob
     __lowerCAmelCase  :      Union[str, Any]    =						max_position_embeddings
     __lowerCAmelCase  :      int    =						type_vocab_size
     __lowerCAmelCase  :      Tuple    =						type_sequence_label_size
     __lowerCAmelCase  :      int    =						initializer_range
     __lowerCAmelCase  :      Optional[int]    =						num_choices







  def 	SCREAMING_SNAKE_CASE   (      self			:   Optional[int]   )							->     Dict:






     """simple docstring"""







     __lowerCAmelCase  :      Optional[Any]    =						ids_tensor([self.batch_size, self.seq_length]	,							self.vocab_size   )

     __lowerCAmelCase  :      Dict    =						None
     if self.use_attention_mask:
        __lowerCAmelCase  :      Optional[int]    =						random_attention_mask([self.batch_size, self.seq_length]   )

     __lowerCAmelCase  :      Union[str, Any]    =						None
     if self.use_token_type_ids:
        __lowerCAmelCase  :      str    =						ids_tensor([self.batch_size, self.seq_length]	,							self.type_vocab_size   )

     __lowerCAmelCase  :      List[str]    =						BertConfig(
         vocab_size=self.vocab_size	,							hidden_size=self.hidden_size	,							num_hidden_layers=self.num_hidden_layers	,							num_attention_heads=self.num_attention_heads	,							intermediate_size=self.intermediate_size	,							hidden_act=self.hidden_act	,							hidden_dropout_prob=self.hidden_dropout_prob	,							attention_probs_dropout_prob=self.attention_probs_dropout_prob	,							max_position_embeddings=self.max_position_embeddings	,							type_vocab_size=self.type_vocab_size	,							is_decoder=lowerCAmelCase	,							initializer_range=self.initializer_range	,							)

     return config, input_ids, token_type_ids, attention_mask







  def 	SCREAMING_SNAKE_CASE   (      self			:   List[Any]   )							->     Optional[Any]:






     """simple docstring"""







     __lowerCAmelCase  :      Any    =						self.prepare_config_and_inputs()
     __lowerCAmelCase   ,__lowerCAmelCase   ,__lowerCAmelCase   ,__lowerCAmelCase  :      List[str]    =						config_and_inputs
     __lowerCAmelCase  :      Any    =						{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": attention_mask}
     return config, inputs_dict







  def 	SCREAMING_SNAKE_CASE   (      self			:   Union[str, Any]   )							->     Tuple:






     """simple docstring"""







     __lowerCAmelCase  :      List[str]    =						self.prepare_config_and_inputs()
     __lowerCAmelCase   ,__lowerCAmelCase   ,__lowerCAmelCase   ,__lowerCAmelCase  :      Dict    =						config_and_inputs

     __lowerCAmelCase  :      Any    =						True
     __lowerCAmelCase  :      str    =						floats_tensor([self.batch_size, self.seq_length, self.hidden_size]   )
     __lowerCAmelCase  :      int    =						ids_tensor([self.batch_size, self.seq_length]	,							vocab_size=2   )

     return (
         config,
         input_ids,
         attention_mask,
         encoder_hidden_states,
         encoder_attention_mask,
     )






@require_flax
class SCREAMING_SNAKE_CASE     (  a_       ,							unittest.TestCase			):
  """simple docstring"""






  lowerCamelCase   :      int										=True

  lowerCamelCase   :      Any										=(
      (
          FlaxBertModel,
          FlaxBertForPreTraining,
          FlaxBertForMaskedLM,
          FlaxBertForMultipleChoice,
          FlaxBertForQuestionAnswering,
          FlaxBertForNextSentencePrediction,
          FlaxBertForSequenceClassification,
          FlaxBertForTokenClassification,
          FlaxBertForQuestionAnswering,
      )
      if is_flax_available()
      else ()
  )







  def 	SCREAMING_SNAKE_CASE   (      self			:   str   )							->     Union[str, Any]:






     """simple docstring"""







     __lowerCAmelCase  :      Union[str, Any]    =						FlaxBertModelTester(self   )







  @slow
  def 	SCREAMING_SNAKE_CASE   (      self			:   Dict   )							->     Dict:






     """simple docstring"""







     __lowerCAmelCase  :      int    =						FlaxBertModel.from_pretrained("""bert-base-cased"""   )
     __lowerCAmelCase  :      str    =						model(np.ones((1, 1)   )   )
     self.assertIsNotNone(lowerCAmelCase   )


