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





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

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

												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						(							lowercase__						,						lowercase__						,						lowercase__		):
						lowercase__							       =				random.Random()
						lowercase__							       =				[rng.randint(0						,						vocab_size - 1		) for i in range(batch_size * sequence_length		)]
						return tf.constant(lowercase__						,						shape=(batch_size, sequence_length)						,						dtype=tf.intaa		)







class 	A							(				lowerCAmelCase_ ):
			lowerCamelCase							:							TensorFlowBenchmarkArguments
			lowerCamelCase							:							PretrainedConfig
			lowerCamelCase							:							str										=     "TensorFlow"



			@property
			def A__       (    self				)	->						str:




									'''simple docstring'''



									return tf.__version__



			def A__       (    self				,   lowerCamelCase__				,   lowerCamelCase__				,   lowerCamelCase__				)	->						Any:




									'''simple docstring'''



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



			def A__       (    self				,   lowerCamelCase__				,   lowerCamelCase__				,   lowerCamelCase__				)	->						Any:




									'''simple docstring'''



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



			def A__       (    self				,   lowerCamelCase__				,   lowerCamelCase__				,   lowerCamelCase__				)	->						Tuple:




									'''simple docstring'''



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



			def A__       (    self				,   lowerCamelCase__				,   lowerCamelCase__				,   lowerCamelCase__				)	->						List[str]:




									'''simple docstring'''



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

									lowercase__							       =				self._prepare_train_func(__lowerCAmelCase				,   __lowerCAmelCase				,   __lowerCAmelCase				)
									return self._measure_memory(_train				)



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




									'''simple docstring'''



									lowercase__							       =				self.config_dict[model_name]

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

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

									# encoder-decoder has vocab size saved differently
									lowercase__							       =				config.vocab_size if hasattr(__lowerCAmelCase				,   """vocab_size"""				) else config.encoder.vocab_size
									lowercase__							       =				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				)

									lowercase__							       =				encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

									return _inference



			def A__       (    self				,   lowerCamelCase__				,   lowerCamelCase__				,   lowerCamelCase__				)	->						str:




									'''simple docstring'''



									lowercase__							       =				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."""				)

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

									# encoder-decoder has vocab size saved differently
									lowercase__							       =				config.vocab_size if hasattr(__lowerCAmelCase				,   """vocab_size"""				) else config.encoder.vocab_size
									lowercase__							       =				random_input_ids(__lowerCAmelCase				,   __lowerCAmelCase				,   __lowerCAmelCase				)

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

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

									lowercase__							       =				encoder_decoder_train if config.is_encoder_decoder else encoder_train

									return _train



			def A__       (    self				,   lowerCamelCase__				)	->						Any:




									'''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
																					lowercase__							       =				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 A__       (    self				,   lowerCamelCase__				)	->						Union[str, 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."""				)
																											lowercase__							       =				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."""				)
																																	lowercase__							       =				"""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()
																																	lowercase__							       =				nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx				)
																																	lowercase__							       =				nvml.nvmlDeviceGetMemoryInfo(__lowerCAmelCase				)
																																	lowercase__							       =				meminfo.used
																																	lowercase__							       =				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."""				)
																																	lowercase__							       =				None
																											else:
																																	lowercase__							       =				measure_peak_memory_cpu(__lowerCAmelCase				)
																																	lowercase__							       =				Memory(__lowerCAmelCase				) if isinstance(__lowerCAmelCase				,   __lowerCAmelCase				) else memory_bytes
																					if self.args.trace_memory_line_by_line:
																											lowercase__							       =				stop_memory_tracing(__lowerCAmelCase				)
																											if memory is None:
																																	lowercase__							       =				summary.total
																					else:
																											lowercase__							       =				None

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







"""simple docstring"""


from collections import OrderedDict
from typing import List, Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


UpperCAmelCase__     						=  logging.get_logger(__name__)

UpperCAmelCase__     						=  {
    """google/efficientnet-b7""": """https://huggingface.co/google/efficientnet-b7/resolve/main/config.json""",
}


class 				a (						lowerCAmelCase_      ):
    _snake_case   :     Any   							=							'efficientnet'


    def __init__(    self		:     Any      , __lowerCAmelCase		:     int = 3      , __lowerCAmelCase		:     int = 600      , __lowerCAmelCase		:     float = 2.0      , __lowerCAmelCase		:     float = 3.1      , __lowerCAmelCase		:     int = 8      , __lowerCAmelCase		:     List[int] = [3, 3, 5, 3, 5, 5, 3]      , __lowerCAmelCase		:     List[int] = [32, 16, 24, 40, 80, 112, 192]      , __lowerCAmelCase		:     List[int] = [16, 24, 40, 80, 112, 192, 320]      , __lowerCAmelCase		:     List[int] = []      , __lowerCAmelCase		:     List[int] = [1, 2, 2, 2, 1, 2, 1]      , __lowerCAmelCase		:     List[int] = [1, 2, 2, 3, 3, 4, 1]      , __lowerCAmelCase		:     List[int] = [1, 6, 6, 6, 6, 6, 6]      , __lowerCAmelCase		:     float = 0.25      , __lowerCAmelCase		:     str = "swish"      , __lowerCAmelCase		:     int = 2560      , __lowerCAmelCase		:     str = "mean"      , __lowerCAmelCase		:     float = 0.02      , __lowerCAmelCase		:     float = 0.001      , __lowerCAmelCase		:     float = 0.99      , __lowerCAmelCase		:     float = 0.5      , __lowerCAmelCase		:     float = 0.2      , **__lowerCAmelCase		:     List[Any]      , ):
      super().__init__(**__lowerCAmelCase    )

      _UpperCAmelCase         =					num_channels
      _UpperCAmelCase         =					image_size
      _UpperCAmelCase         =					width_coefficient
      _UpperCAmelCase         =					depth_coefficient
      _UpperCAmelCase         =					depth_divisor
      _UpperCAmelCase         =					kernel_sizes
      _UpperCAmelCase         =					in_channels
      _UpperCAmelCase         =					out_channels
      _UpperCAmelCase         =					depthwise_padding
      _UpperCAmelCase         =					strides
      _UpperCAmelCase         =					num_block_repeats
      _UpperCAmelCase         =					expand_ratios
      _UpperCAmelCase         =					squeeze_expansion_ratio
      _UpperCAmelCase         =					hidden_act
      _UpperCAmelCase         =					hidden_dim
      _UpperCAmelCase         =					pooling_type
      _UpperCAmelCase         =					initializer_range
      _UpperCAmelCase         =					batch_norm_eps
      _UpperCAmelCase         =					batch_norm_momentum
      _UpperCAmelCase         =					dropout_rate
      _UpperCAmelCase         =					drop_connect_rate
      _UpperCAmelCase         =					sum(__lowerCAmelCase    ) * 4



class 				a (						lowerCAmelCase_      ):
    _snake_case   :     Dict   							=							version.parse('1.11'      )


    @property
    def    lowerCAmelCase_     (    self		:     Any    ):
      return OrderedDict(
          [
              ("""pixel_values""", {0: """batch""", 1: """num_channels""", 2: """height""", 3: """width"""}),
          ]    )


    @property
    def    lowerCAmelCase_     (    self		:     int    ):
      return 1e-5
