




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

def 						snake_case__       (					_lowerCamelCase,  _lowerCamelCase				)			->int:






		"""simple docstring"""


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

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

				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__       (					_lowerCamelCase,  _lowerCamelCase,  _lowerCamelCase				)			->["tf.Tensor"]:






		"""simple docstring"""


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





class       lowerCAmelCase__       (			lowerCAmelCase_    ):


				"""simple docstring"""





				__UpperCAmelCase :     TensorFlowBenchmarkArguments
				__UpperCAmelCase :     PretrainedConfig
				__UpperCAmelCase :     str       =      "TensorFlow"



				@property
				def       snake_case    (     self					:	Any						):
						return tf.__version__



				def       snake_case    (     self					:	int    ,   lowercase__					:	str    ,   lowercase__					:	int    ,   lowercase__					:	int						):
						# initialize GPU on separate process
						__lowercase				:		List[str]												=				self.args.strategy
						if strategy is None:
								raise ValueError("A device strategy has to be initialized before using TensorFlow."						)
						__lowercase				:		List[Any]												=				self._prepare_inference_func(lowercase__    ,   lowercase__    ,   lowercase__						)
						return self._measure_speed(_inference						)



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



				def       snake_case    (     self					:	Dict    ,   lowercase__					:	str    ,   lowercase__					:	int    ,   lowercase__					:	int						):
						# initialize GPU on separate process
						if self.args.is_gpu:
								tf.config.experimental.set_memory_growth(self.args.gpu_list[self.args.device_idx]    ,   lowercase__						)
						__lowercase				:		Any												=				self.args.strategy
						if strategy is None:
								raise ValueError("A device strategy has to be initialized before using TensorFlow."						)
						__lowercase				:		Dict												=				self._prepare_inference_func(lowercase__    ,   lowercase__    ,   lowercase__						)
						return self._measure_memory(_inference						)



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

						__lowercase				:		int												=				self._prepare_train_func(lowercase__    ,   lowercase__    ,   lowercase__						)
						return self._measure_memory(_train						)



				def       snake_case    (     self					:	Optional[int]    ,   lowercase__					:	str    ,   lowercase__					:	int    ,   lowercase__					:	int						):
						__lowercase				:		Optional[Any]												=				self.config_dict[model_name]

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

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

						# encoder-decoder has vocab size saved differently
						__lowercase				:		Tuple												=				config.vocab_size if hasattr(lowercase__    ,   "vocab_size"						) else config.encoder.vocab_size
						__lowercase				:		Optional[int]												=				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__						)

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

						return _inference



				def       snake_case    (     self					:	str    ,   lowercase__					:	str    ,   lowercase__					:	int    ,   lowercase__					:	int						):
						__lowercase				:		List[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."						)

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

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

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

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

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

						return _train



				def       snake_case    (     self					:	Tuple    ,   lowercase__					:	int						):
						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
										__lowercase				:		Optional[Any]												=				timeit.repeat(
										    lowercase__    ,   repeat=self.args.repeat    ,   number=1_0    ,   )

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



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

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






"""simple docstring"""







import os
import sys
import unittest


__A   :  Optional[Any]					=      os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
sys.path.append(os.path.join(git_repo_path, 'utils'))

import check_dummies  # noqa: E402
from check_dummies import create_dummy_files, create_dummy_object, find_backend, read_init  # noqa: E402


# Align TRANSFORMERS_PATH in check_dummies with the current path
__A   :  Optional[Any]					=      os.path.join(git_repo_path, 'src', 'diffusers')




class       lowerCAmelCase__       (			unittest.TestCase    ):


				"""simple docstring"""





				def       snake_case    (     self					:	str						):
						__lowercase				:		int												=				find_backend("    if not is_torch_available():"						)
						self.assertEqual(lowercase__    ,   "torch"						)

						# backend_with_underscore = find_backend("    if not is_tensorflow_text_available():")
						# self.assertEqual(backend_with_underscore, "tensorflow_text")

						__lowercase				:		int												=				find_backend("    if not (is_torch_available() and is_transformers_available()):"						)
						self.assertEqual(lowercase__    ,   "torch_and_transformers"						)

						# double_backend_with_underscore = find_backend(
						#    "    if not (is_sentencepiece_available() and is_tensorflow_text_available()):"
						# )
						# self.assertEqual(double_backend_with_underscore, "sentencepiece_and_tensorflow_text")

						__lowercase				:		str												=				find_backend(
						    "    if not (is_torch_available() and is_transformers_available() and is_onnx_available()):"						)
						self.assertEqual(lowercase__    ,   "torch_and_transformers_and_onnx"						)



				def       snake_case    (     self					:	Any						):
						__lowercase				:		List[Any]												=				read_init()
						# We don't assert on the exact list of keys to allow for smooth grow of backend-specific objects
						self.assertIn("torch"    ,   lowercase__						)
						self.assertIn("torch_and_transformers"    ,   lowercase__						)
						self.assertIn("flax_and_transformers"    ,   lowercase__						)
						self.assertIn("torch_and_transformers_and_onnx"    ,   lowercase__						)

						# Likewise, we can't assert on the exact content of a key
						self.assertIn("UNet2DModel"    ,   objects["torch"]						)
						self.assertIn("FlaxUNet2DConditionModel"    ,   objects["flax"]						)
						self.assertIn("StableDiffusionPipeline"    ,   objects["torch_and_transformers"]						)
						self.assertIn("FlaxStableDiffusionPipeline"    ,   objects["flax_and_transformers"]						)
						self.assertIn("LMSDiscreteScheduler"    ,   objects["torch_and_scipy"]						)
						self.assertIn("OnnxStableDiffusionPipeline"    ,   objects["torch_and_transformers_and_onnx"]						)



				def       snake_case    (     self					:	Dict						):
						__lowercase				:		Tuple												=				create_dummy_object("CONSTANT"    ,   "'torch'"						)
						self.assertEqual(lowercase__    ,   "\nCONSTANT = None\n"						)

						__lowercase				:		Union[str, Any]												=				create_dummy_object("function"    ,   "'torch'"						)
						self.assertEqual(
						    lowercase__    ,   "\ndef function(*args, **kwargs):\n    requires_backends(function, 'torch')\n"						)

						__lowercase				:		Tuple												=				"\nclass FakeClass(metaclass=DummyObject):\n    _backends = 'torch'\n\n    def __init__(self, *args, **kwargs):\n        requires_backends(self, 'torch')\n\n    @classmethod\n    def from_config(cls, *args, **kwargs):\n        requires_backends(cls, 'torch')\n\n    @classmethod\n    def from_pretrained(cls, *args, **kwargs):\n        requires_backends(cls, 'torch')\n"
						__lowercase				:		Dict												=				create_dummy_object("FakeClass"    ,   "'torch'"						)
						self.assertEqual(lowercase__    ,   lowercase__						)



				def       snake_case    (     self					:	List[Any]						):
						__lowercase				:		Optional[int]												=				"# This file is autogenerated by the command `make fix-copies`, do not edit.\nfrom ..utils import DummyObject, requires_backends\n\n\nCONSTANT = None\n\n\ndef function(*args, **kwargs):\n    requires_backends(function, [\"torch\"])\n\n\nclass FakeClass(metaclass=DummyObject):\n    _backends = [\"torch\"]\n\n    def __init__(self, *args, **kwargs):\n        requires_backends(self, [\"torch\"])\n\n    @classmethod\n    def from_config(cls, *args, **kwargs):\n        requires_backends(cls, [\"torch\"])\n\n    @classmethod\n    def from_pretrained(cls, *args, **kwargs):\n        requires_backends(cls, [\"torch\"])\n"
						__lowercase				:		List[str]												=				create_dummy_files({"torch": ["CONSTANT", "function", "FakeClass"]}						)
						self.assertEqual(dummy_files["torch"]    ,   lowercase__						)
