






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

__a       :	Dict             = logging.get_logger(__name__)





def        SCREAMING_SNAKE_CASE       (      lowerCamelCase_ ,     lowerCamelCase_):





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

          @wraps(__snake_case)
          @tf.function(experimental_compile=__snake_case)
          def run_in_graph_mode(*lowerCamelCase_ ,     **lowerCamelCase_):
               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        SCREAMING_SNAKE_CASE       (      lowerCamelCase_ ,     lowerCamelCase_ ,     lowerCamelCase_):
     a__          =  random.Random()
     a__          =  [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        _SCREAMING_SNAKE_CASE						(  __snake_case							):





       """simple docstring"""





       _SCREAMING_SNAKE_CASE    =42
       _SCREAMING_SNAKE_CASE    =42
       _SCREAMING_SNAKE_CASE    ='TensorFlow'




       @property
       def  lowercase		(   self:     Dict    ):

            '''simple docstring'''

            return tf.__version__




       def  lowercase		(   self:     Union[str, Any]				,		__A:     int				,		__A:     str				,		__A:     List[str]    ):

            '''simple docstring'''

            a__          =  self.args.strategy
            if strategy is None:
                 raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''    )
            a__          =  self._prepare_inference_func(__a				,		__a				,		__a    )
            return self._measure_speed(_inference    )




       def  lowercase		(   self:     List[str]				,		__A:     Optional[int]				,		__A:     Union[str, Any]				,		__A:     str    ):

            '''simple docstring'''

            a__          =  self.args.strategy
            if strategy is None:
                 raise ValueError('''A device strategy has to be initialized before using TensorFlow.'''    )
            a__          =  self._prepare_train_func(__a				,		__a				,		__a    )
            return self._measure_speed(_train    )




       def  lowercase		(   self:     Optional[Any]				,		__A:     Tuple				,		__A:     Optional[int]				,		__A:     Optional[Any]    ):

            '''simple docstring'''

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




       def  lowercase		(   self:     Union[str, Any]				,		__A:     List[Any]				,		__A:     Optional[Any]				,		__A:     int    ):

            '''simple docstring'''

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

            a__          =  self._prepare_train_func(__a				,		__a				,		__a    )
            return self._measure_memory(_train    )




       def  lowercase		(   self:     Union[str, Any]				,		__A:     Optional[int]				,		__A:     Dict				,		__A:     Tuple    ):

            '''simple docstring'''

            a__          =  self.config_dict[model_name]

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

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

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

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

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

            a__          =  encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

            return _inference




       def  lowercase		(   self:     Optional[int]				,		__A:     Union[str, Any]				,		__A:     Union[str, Any]				,		__A:     Tuple    ):

            '''simple docstring'''

            a__          =  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.'''    )

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

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

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

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

            a__          =  encoder_decoder_train if config.is_encoder_decoder else encoder_train

            return _train




       def  lowercase		(   self:     Dict				,		__A:     int    ):

            '''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(__a				,		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
                      a__          =  timeit.repeat(
                          __a				,		repeat=self.args.repeat				,		number=10				,		)

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





       def  lowercase		(   self:     Optional[int]				,		__A:     Any    ):

            '''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.'''    )
                           a__          =  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.'''    )
                                a__          =  '''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()
                                a__          =  nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx    )
                                a__          =  nvml.nvmlDeviceGetMemoryInfo(__a    )
                                a__          =  meminfo.used
                                a__          =  Memory(__a    )
                                # 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.'''    )
                                a__          =  None
                           else:
                                a__          =  measure_peak_memory_cpu(__a    )
                                a__          =  Memory(__a    ) if isinstance(__a				,		__a    ) else memory_bytes
                      if self.args.trace_memory_line_by_line:
                           a__          =  stop_memory_tracing(__a    )
                           if memory is None:
                                a__          =  summary.total
                      else:
                           a__          =  None

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










"""simple docstring"""







__a       :	List[Any]             = frozenset(
    [
        'prompt',
        'height',
        'width',
        'guidance_scale',
        'negative_prompt',
        'prompt_embeds',
        'negative_prompt_embeds',
        'cross_attention_kwargs',
    ]
)

__a       :	Union[str, Any]             = frozenset(['prompt', 'negative_prompt'])

__a       :	Any             = frozenset([])

__a       :	Union[str, Any]             = frozenset(['image'])

__a       :	Dict             = frozenset(
    [
        'image',
        'height',
        'width',
        'guidance_scale',
    ]
)

__a       :	Dict             = frozenset(['image'])

__a       :	Dict             = frozenset(
    [
        'prompt',
        'image',
        'height',
        'width',
        'guidance_scale',
        'negative_prompt',
        'prompt_embeds',
        'negative_prompt_embeds',
    ]
)

__a       :	Optional[Any]             = frozenset(['prompt', 'image', 'negative_prompt'])

__a       :	List[Any]             = frozenset(
    [
        # Text guided image variation with an image mask
        'prompt',
        'image',
        'mask_image',
        'height',
        'width',
        'guidance_scale',
        'negative_prompt',
        'prompt_embeds',
        'negative_prompt_embeds',
    ]
)

__a       :	Union[str, Any]             = frozenset(['prompt', 'image', 'mask_image', 'negative_prompt'])

__a       :	Optional[Any]             = frozenset(
    [
        # image variation with an image mask
        'image',
        'mask_image',
        'height',
        'width',
        'guidance_scale',
    ]
)

__a       :	int             = frozenset(['image', 'mask_image'])

__a       :	Tuple             = frozenset(
    [
        'example_image',
        'image',
        'mask_image',
        'height',
        'width',
        'guidance_scale',
    ]
)

__a       :	Optional[Any]             = frozenset(['example_image', 'image', 'mask_image'])

__a       :	Optional[Any]             = frozenset(['class_labels'])

__a       :	Tuple             = frozenset(['class_labels'])

__a       :	int             = frozenset(['batch_size'])

__a       :	int             = frozenset([])

__a       :	Union[str, Any]             = frozenset(['batch_size'])

__a       :	Tuple             = frozenset([])

__a       :	Dict             = frozenset(
    [
        'prompt',
        'audio_length_in_s',
        'guidance_scale',
        'negative_prompt',
        'prompt_embeds',
        'negative_prompt_embeds',
        'cross_attention_kwargs',
    ]
)

__a       :	Dict             = frozenset(['prompt', 'negative_prompt'])
__a       :	Optional[int]             = frozenset(['input_tokens'])

__a       :	str             = frozenset(['input_tokens'])


