


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


def _lowerCamelCase	(		_UpperCamelCase    , _UpperCamelCase      ):

			'''simple docstring'''




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

						@wraps(A_      )
						@tf.function(experimental_compile=A_      )
						def run_in_graph_mode(*_UpperCamelCase    , **_UpperCamelCase      ):
									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 _lowerCamelCase	(		_UpperCamelCase    , _UpperCamelCase    , _UpperCamelCase      ):

			'''simple docstring'''




			__lowerCAmelCase				    =			random.Random()
			__lowerCAmelCase				    =			[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   _UpperCamelCase  (   _a		):
							'''simple docstring'''







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

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

							def 			snake_case				(     self    ,   __a    ,   __a    ,   __a	):
										# initialize GPU on separate process
										__lowerCAmelCase				    =			self.args.strategy
										if strategy is None:
													raise ValueError("A device strategy has to be initialized before using TensorFlow."	)
										__lowerCAmelCase				    =			self._prepare_inference_func(__UpperCAmelCase    ,   __UpperCAmelCase    ,   __UpperCAmelCase	)
										return self._measure_speed(_inference	)

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

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

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

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

							def 			snake_case				(     self    ,   __a    ,   __a    ,   __a	):
										__lowerCAmelCase				    =			self.config_dict[model_name]

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

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

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

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

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

										__lowerCAmelCase				    =			encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

										return _inference

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

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

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

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

										__lowerCAmelCase				    =			encoder_decoder_train if config.is_encoder_decoder else encoder_train

										return _train

							def 			snake_case				(     self    ,   __a	):
										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(__UpperCAmelCase    ,   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				    =			timeit.repeat(
																    __UpperCAmelCase    ,   repeat=self.args.repeat    ,   number=10    ,   )

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



							def 			snake_case				(     self    ,   __a	):
										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				    =			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				    =			"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				    =			nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx	)
																						__lowerCAmelCase				    =			nvml.nvmlDeviceGetMemoryInfo(__UpperCAmelCase	)
																						__lowerCAmelCase				    =			meminfo.used
																						__lowerCAmelCase				    =			Memory(__UpperCAmelCase	)
																						# 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				    =			None
																			else:
																						__lowerCAmelCase				    =			measure_peak_memory_cpu(__UpperCAmelCase	)
																						__lowerCAmelCase				    =			Memory(__UpperCAmelCase	) if isinstance(__UpperCAmelCase    ,   __UpperCAmelCase	) else memory_bytes
																if self.args.trace_memory_line_by_line:
																			__lowerCAmelCase				    =			stop_memory_tracing(__UpperCAmelCase	)
																			if memory is None:
																						__lowerCAmelCase				    =			summary.total
																else:
																			__lowerCAmelCase				    =			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 typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


__lowercase    					=      {
    """configuration_rag""": ["""RagConfig"""],
    """retrieval_rag""": ["""RagRetriever"""],
    """tokenization_rag""": ["""RagTokenizer"""],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __lowercase    					=      [
     """RagModel""",
     """RagPreTrainedModel""",
     """RagSequenceForGeneration""",
     """RagTokenForGeneration""",
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __lowercase    					=      [
     """TFRagModel""",
     """TFRagPreTrainedModel""",
     """TFRagSequenceForGeneration""",
     """TFRagTokenForGeneration""",
 ]


if TYPE_CHECKING:
 from .configuration_rag import RagConfig
 from .retrieval_rag import RagRetriever
 from .tokenization_rag import RagTokenizer

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_rag import RagModel, RagPreTrainedModel, RagSequenceForGeneration, RagTokenForGeneration

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_rag import (
      TFRagModel,
      TFRagPreTrainedModel,
      TFRagSequenceForGeneration,
      TFRagTokenForGeneration,
  )

else:
 import sys

 __lowercase    					=      _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)






