


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       :       Optional[Any]						       =				logging.get_logger(__name__)

def 			_a  (				SCREAMING_SNAKE_CASE	:				bool    , SCREAMING_SNAKE_CASE	:				bool     ):







 """simple docstring"""



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

  @wraps(SCREAMING_SNAKE_CASE     )
  @tf.function(experimental_compile=SCREAMING_SNAKE_CASE     )
  def run_in_graph_mode(*SCREAMING_SNAKE_CASE	:				int    , **SCREAMING_SNAKE_CASE	:				str     ):
   return func(*SCREAMING_SNAKE_CASE    , **SCREAMING_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 			_a  (				SCREAMING_SNAKE_CASE	:				int    , SCREAMING_SNAKE_CASE	:				int    , SCREAMING_SNAKE_CASE	:				int     ):







 """simple docstring"""



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






class       __magic_name__	(     __lowerCAmelCase):
     A:		TensorFlowBenchmarkArguments
     A:		PretrainedConfig
     A:		str			=      "TensorFlow"


     @property
     def      UpperCAmelCase__	(      self							:     Union[str, Any]			)  ->      List[str]:





      '''simple docstring'''






      return tf.__version__


     def      UpperCAmelCase__	(      self							:     Optional[Any]							,	lowerCamelCase__							:     str							,	lowerCamelCase__							:     int							,	lowerCamelCase__							:     int			)  ->      float:





      '''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(lowerCamelCase__							,	lowerCamelCase__							,	lowerCamelCase__			)
      return self._measure_speed(_inference			)


     def      UpperCAmelCase__	(      self							:     Optional[Any]							,	lowerCamelCase__							:     str							,	lowerCamelCase__							:     int							,	lowerCamelCase__							:     int			)  ->      float:





      '''simple docstring'''






      UpperCamelCase__					:    Dict          =   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(lowerCamelCase__							,	lowerCamelCase__							,	lowerCamelCase__			)
      return self._measure_speed(_train			)


     def      UpperCAmelCase__	(      self							:     Union[str, 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__			)
      UpperCamelCase__					:    int          =   self.args.strategy
      if strategy is None:
       raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''			)
      UpperCamelCase__					:    Optional[Any]          =   self._prepare_inference_func(lowerCamelCase__							,	lowerCamelCase__							,	lowerCamelCase__			)
      return self._measure_memory(_inference			)


     def      UpperCAmelCase__	(      self							:     Tuple							,	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__			)
      UpperCamelCase__					:    Tuple          =   self.args.strategy
      if strategy is None:
       raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''			)

      UpperCamelCase__					:    Optional[Any]          =   self._prepare_train_func(lowerCamelCase__							,	lowerCamelCase__							,	lowerCamelCase__			)
      return self._measure_memory(_train			)


     def      UpperCAmelCase__	(      self							:     Union[str, Any]							,	lowerCamelCase__							:     str							,	lowerCamelCase__							:     int							,	lowerCamelCase__							:     int			)  ->      Callable[[], None]:





      '''simple docstring'''






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

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

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

      # encoder-decoder has vocab size saved differently
      UpperCamelCase__					:    int          =   config.vocab_size if hasattr(lowerCamelCase__							,	'''vocab_size'''			) else config.encoder.vocab_size
      UpperCamelCase__					:    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__			)

      UpperCamelCase__					:    List[str]          =   encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

      return _inference


     def      UpperCAmelCase__	(      self							:     Tuple							,	lowerCamelCase__							:     str							,	lowerCamelCase__							:     int							,	lowerCamelCase__							:     int			)  ->      Callable[[], None]:





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

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

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

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

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

      return _train


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

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






import contextlib
import importlib
import io
import unittest

import transformers

# Try to import everything from transformers to ensure every object can be loaded.
from transformers import *  # noqa F406
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, require_tf, require_torch
from transformers.utils import ContextManagers, find_labels, is_flax_available, is_tf_available, is_torch_available


if is_torch_available():
       from transformers import BertForPreTraining, BertForQuestionAnswering, BertForSequenceClassification

if is_tf_available():
       from transformers import TFBertForPreTraining, TFBertForQuestionAnswering, TFBertForSequenceClassification

if is_flax_available():
       from transformers import FlaxBertForPreTraining, FlaxBertForQuestionAnswering, FlaxBertForSequenceClassification


__UpperCamelCase       :       Dict						       =				DUMMY_UNKNOWN_IDENTIFIER
# An actual model hosted on huggingface.co

__UpperCamelCase       :       Optional[int]						       =				"main"
# Default branch name
__UpperCamelCase       :       Optional[Any]						       =				"f2c752cfc5c0ab6f4bdec59acea69eefbee381c2"
# One particular commit (not the top of `main`)
__UpperCamelCase       :       str						       =				"aaaaaaa"
# This commit does not exist, so we should 404.

__UpperCamelCase       :       Optional[Any]						       =				"d9e9f15bc825e4b2c9249e9578f884bbcb5e3684"
# Sha-1 of config.json on the top of `main`, for checking purposes
__UpperCamelCase       :       Tuple						       =				"4b243c475af8d0a7754e87d7d096c92e5199ec2fe168a2ee7998e3b8e9bcb1d3"

@contextlib.contextmanager
def 			_a  (				):







 """simple docstring"""



 print('''Welcome!'''     )
 yield
 print('''Bye!'''     )

@contextlib.contextmanager
def 			_a  (				):







 """simple docstring"""



 print('''Bonjour!'''     )
 yield
 print('''Au revoir!'''     )

class       __magic_name__	(     unittest.TestCase):


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





      '''simple docstring'''






      assert transformers.__spec__ is not None
      assert importlib.util.find_spec('''transformers'''			) is not None



class       __magic_name__	(     unittest.TestCase):


     @unittest.mock.patch('''sys.stdout'''							,	new_callable=io.StringIO			)
     def      UpperCAmelCase__	(      self							:     Union[str, Any]							,	lowerCamelCase__							:     Union[str, Any]			)  ->      Optional[Any]:





      '''simple docstring'''






      with ContextManagers([]			):
       print('''Transformers are awesome!'''			)
      # The print statement adds a new line at the end of the output
      self.assertEqual(mock_stdout.getvalue()							,	'''Transformers are awesome!\n'''			)


     @unittest.mock.patch('''sys.stdout'''							,	new_callable=io.StringIO			)
     def      UpperCAmelCase__	(      self							:     Any							,	lowerCamelCase__							:     Union[str, Any]			)  ->      Dict:





      '''simple docstring'''






      with ContextManagers([context_en()]			):
       print('''Transformers are awesome!'''			)
      # The output should be wrapped with an English welcome and goodbye
      self.assertEqual(mock_stdout.getvalue()							,	'''Welcome!\nTransformers are awesome!\nBye!\n'''			)


     @unittest.mock.patch('''sys.stdout'''							,	new_callable=io.StringIO			)
     def      UpperCAmelCase__	(      self							:     List[Any]							,	lowerCamelCase__							:     List[str]			)  ->      Tuple:





      '''simple docstring'''






      with ContextManagers([context_fr(), context_en()]			):
       print('''Transformers are awesome!'''			)
      # The output should be wrapped with an English and French welcome and goodbye
      self.assertEqual(mock_stdout.getvalue()							,	'''Bonjour!\nWelcome!\nTransformers are awesome!\nBye!\nAu revoir!\n'''			)


     @require_torch
     def      UpperCAmelCase__	(      self							:     Tuple			)  ->      Optional[int]:





      '''simple docstring'''






      self.assertEqual(find_labels(lowerCamelCase__			)							,	['''labels''']			)
      self.assertEqual(find_labels(lowerCamelCase__			)							,	['''labels''', '''next_sentence_label''']			)
      self.assertEqual(find_labels(lowerCamelCase__			)							,	['''start_positions''', '''end_positions''']			)

      class       __magic_name__	(     __lowerCAmelCase):
           pass



      self.assertEqual(find_labels(lowerCamelCase__			)							,	['''labels''']			)


     @require_tf
     def      UpperCAmelCase__	(      self							:     Union[str, Any]			)  ->      Tuple:





      '''simple docstring'''






      self.assertEqual(find_labels(lowerCamelCase__			)							,	['''labels''']			)
      self.assertEqual(find_labels(lowerCamelCase__			)							,	['''labels''', '''next_sentence_label''']			)
      self.assertEqual(find_labels(lowerCamelCase__			)							,	['''start_positions''', '''end_positions''']			)

      class       __magic_name__	(     __lowerCAmelCase):
           pass



      self.assertEqual(find_labels(lowerCamelCase__			)							,	['''labels''']			)


     @require_flax
     def      UpperCAmelCase__	(      self							:     str			)  ->      Optional[Any]:





      '''simple docstring'''






      self.assertEqual(find_labels(lowerCamelCase__			)							,	[]			)
      self.assertEqual(find_labels(lowerCamelCase__			)							,	[]			)
      self.assertEqual(find_labels(lowerCamelCase__			)							,	[]			)

      class       __magic_name__	(     __lowerCAmelCase):
           pass



      self.assertEqual(find_labels(lowerCamelCase__			)							,	[]			)


