






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





def 					snake_case	(  lowerCamelCase					,	lowerCamelCase       ):




	'''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       ):




	'''simple docstring'''







	__lowercase   							=	random.Random()
	__lowercase   							=	[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 							__UpperCamelCase     (  _lowerCAmelCase ):
	__snake_case						:TensorFlowBenchmarkArguments
	__snake_case						:PretrainedConfig
	__snake_case						:str =							"TensorFlow"







	@property
	def 			_a (	self   :				str      )		->			Any:
		"""simple docstring"""







		return tf.__version__







	def 			_a (	self   :				List[str]							,					_lowerCAmelCase   :				str							,					_lowerCAmelCase   :				int							,					_lowerCAmelCase   :				int      )		->			float:
		"""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   :				Union[str, Any]							,					_lowerCAmelCase   :				str							,					_lowerCAmelCase   :				int							,					_lowerCAmelCase   :				int      )		->			float:
		"""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   :				Dict							,					_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      )
		__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   :				Optional[int]							,					_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      )
		__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   :				Tuple							,					_lowerCAmelCase   :				str							,					_lowerCAmelCase   :				int							,					_lowerCAmelCase   :				int      )		->			Callable[[], None]:
		"""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   :				Union[str, Any]							,					_lowerCAmelCase   :				str							,					_lowerCAmelCase   :				int							,					_lowerCAmelCase   :				int      )		->			Callable[[], None]:
		"""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   :				Dict							,					_lowerCAmelCase   :				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
				__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   :				str							,					_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."""      )
					__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"""




import os

import torch

from ..logging import get_logger
from .constants import FSDP_PYTORCH_VERSION, MODEL_NAME, OPTIMIZER_NAME
from .versions import is_torch_version


if is_torch_version('>=', FSDP_PYTORCH_VERSION):
						import torch.distributed.checkpoint as dist_cp
						from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner, DefaultSavePlanner
						from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_dict
						from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
						from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType


lowerCAmelCase__				=					get_logger(__name__)



def 							_lowerCamelCase			(			__a,				__a,				__a,				__a,				__a=0     ):
		os.makedirs(__a,				exist_ok=__a     )
		with FSDP.state_dict_type(
		    __a,				fsdp_plugin.state_dict_type,				fsdp_plugin.state_dict_config,				fsdp_plugin.optim_state_dict_config     ):
				SCREAMING_SNAKE_CASE_    			=			model.state_dict()
				if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
						SCREAMING_SNAKE_CASE_    			=			F'{MODEL_NAME}.bin' if model_index == 0 else F'{MODEL_NAME}_{model_index}.bin'
						SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				__a     )
						if accelerator.process_index == 0:
								logger.info(F'Saving model to {output_model_file}'     )
								torch.save(__a,				__a     )
								logger.info(F'Model saved to {output_model_file}'     )
				elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
						SCREAMING_SNAKE_CASE_    			=			(
						    F'{MODEL_NAME}_rank{accelerator.process_index}.bin'
						    if model_index == 0
						    else F'{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin'
						)
						SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				__a     )
						logger.info(F'Saving model to {output_model_file}'     )
						torch.save(__a,				__a     )
						logger.info(F'Model saved to {output_model_file}'     )
				elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
						SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				F'{MODEL_NAME}_{model_index}'     )
						os.makedirs(__a,				exist_ok=__a     )
						logger.info(F'Saving model to {ckpt_dir}'     )
						SCREAMING_SNAKE_CASE_    			=			{'''model''': state_dict}

						dist_cp.save_state_dict(
						    state_dict=__a,				storage_writer=dist_cp.FileSystemWriter(__a     ),				planner=DefaultSavePlanner(),				)
						logger.info(F'Model saved to {ckpt_dir}'     )



def 							_lowerCamelCase			(			__a,				__a,				__a,				__a,				__a=0     ):
		accelerator.wait_for_everyone()
		with FSDP.state_dict_type(
		    __a,				fsdp_plugin.state_dict_type,				fsdp_plugin.state_dict_config,				fsdp_plugin.optim_state_dict_config     ):
				if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
						if type(__a     ) != FSDP and accelerator.process_index != 0:
								if not fsdp_plugin.sync_module_states:
										raise ValueError(
										    '''Set the `sync_module_states` flag to `True` so that model states are synced across processes when '''
										    '''initializing FSDP object'''     )
								return
						SCREAMING_SNAKE_CASE_    			=			F'{MODEL_NAME}.bin' if model_index == 0 else F'{MODEL_NAME}_{model_index}.bin'
						SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				__a     )
						logger.info(F'Loading model from {input_model_file}'     )
						SCREAMING_SNAKE_CASE_    			=			torch.load(__a     )
						logger.info(F'Model loaded from {input_model_file}'     )
				elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
						SCREAMING_SNAKE_CASE_    			=			(
						    F'{MODEL_NAME}_rank{accelerator.process_index}.bin'
						    if model_index == 0
						    else F'{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin'
						)
						SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				__a     )
						logger.info(F'Loading model from {input_model_file}'     )
						SCREAMING_SNAKE_CASE_    			=			torch.load(__a     )
						logger.info(F'Model loaded from {input_model_file}'     )
				elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
						SCREAMING_SNAKE_CASE_    			=			(
						    os.path.join(__a,				F'{MODEL_NAME}_{model_index}'     )
						    if F'{MODEL_NAME}' not in input_dir
						    else input_dir
						)
						logger.info(F'Loading model from {ckpt_dir}'     )
						SCREAMING_SNAKE_CASE_    			=			{'''model''': model.state_dict()}
						dist_cp.load_state_dict(
						    state_dict=__a,				storage_reader=dist_cp.FileSystemReader(__a     ),				planner=DefaultLoadPlanner(),				)
						SCREAMING_SNAKE_CASE_    			=			state_dict['''model''']
						logger.info(F'Model loaded from {ckpt_dir}'     )
				model.load_state_dict(__a     )



def 							_lowerCamelCase			(			__a,				__a,				__a,				__a,				__a,				__a=0     ):
		os.makedirs(__a,				exist_ok=__a     )
		with FSDP.state_dict_type(
		    __a,				fsdp_plugin.state_dict_type,				fsdp_plugin.state_dict_config,				fsdp_plugin.optim_state_dict_config     ):
				SCREAMING_SNAKE_CASE_    			=			FSDP.optim_state_dict(__a,				__a     )
				if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
						if accelerator.process_index == 0:
								SCREAMING_SNAKE_CASE_    			=			(
								    F'{OPTIMIZER_NAME}.bin' if optimizer_index == 0 else F'{OPTIMIZER_NAME}_{optimizer_index}.bin'
								)
								SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				__a     )
								logger.info(F'Saving Optimizer state to {output_optimizer_file}'     )
								torch.save(__a,				__a     )
								logger.info(F'Optimizer state saved in {output_optimizer_file}'     )
				else:
						SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				F'{OPTIMIZER_NAME}_{optimizer_index}'     )
						os.makedirs(__a,				exist_ok=__a     )
						logger.info(F'Saving Optimizer state to {ckpt_dir}'     )
						dist_cp.save_state_dict(
						    state_dict={'''optimizer''': optim_state},				storage_writer=dist_cp.FileSystemWriter(__a     ),				planner=DefaultSavePlanner(),				)
						logger.info(F'Optimizer state saved in {ckpt_dir}'     )



def 							_lowerCamelCase			(			__a,				__a,				__a,				__a,				__a,				__a=0     ):
		accelerator.wait_for_everyone()
		with FSDP.state_dict_type(
		    __a,				fsdp_plugin.state_dict_type,				fsdp_plugin.state_dict_config,				fsdp_plugin.optim_state_dict_config     ):
				if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
						SCREAMING_SNAKE_CASE_    			=			None
						# below check should work but currently it isn't working (mostly opytorch issue),
						# in the meantime disabling it at the cost of excess memory usage
						# if accelerator.process_index == 0 or not fsdp_plugin.optim_state_dict_config.rank0_only:
						SCREAMING_SNAKE_CASE_    			=			(
						    F'{OPTIMIZER_NAME}.bin' if optimizer_index == 0 else F'{OPTIMIZER_NAME}_{optimizer_index}.bin'
						)
						SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				__a     )
						logger.info(F'Loading Optimizer state from {input_optimizer_file}'     )
						SCREAMING_SNAKE_CASE_    			=			torch.load(__a     )
						logger.info(F'Optimizer state loaded from {input_optimizer_file}'     )
				else:
						SCREAMING_SNAKE_CASE_    			=			(
						    os.path.join(__a,				F'{OPTIMIZER_NAME}_{optimizer_index}'     )
						    if F'{OPTIMIZER_NAME}' not in input_dir
						    else input_dir
						)
						logger.info(F'Loading Optimizer from {ckpt_dir}'     )
						SCREAMING_SNAKE_CASE_    			=			load_sharded_optimizer_state_dict(
						    model_state_dict=model.state_dict(),				optimizer_key='''optimizer''',				storage_reader=dist_cp.FileSystemReader(__a     ),				)
						SCREAMING_SNAKE_CASE_    			=			optim_state['''optimizer''']
						logger.info(F'Optimizer loaded from {ckpt_dir}'     )
				SCREAMING_SNAKE_CASE_    			=			FSDP.optim_state_dict_to_load(__a,				__a,				__a     )
				optimizer.load_state_dict(__a     )