



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

snake_case_    					=	logging.get_logger(__name__)

def 			_lowerCAmelCase				(       lowercase_		,     lowercase_						):
						def run_func(lowercase_						):
												@wraps(lowerCamelCase__						)
												def run_in_eager_mode(*lowercase_		,     **lowercase_						):
																		return func(*lowerCamelCase__		,     **lowerCamelCase__						)

												@wraps(lowerCamelCase__						)
												@tf.function(experimental_compile=lowerCamelCase__						)
												def run_in_graph_mode(*lowercase_		,     **lowercase_						):
																		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 			_lowerCAmelCase				(       lowercase_		,     lowercase_		,     lowercase_						):
						UpperCAmelCase					=			random.Random()
						UpperCAmelCase					=			[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 				A_      (					a__    ):



						"""simple docstring"""
						__UpperCamelCase		    =							42
						__UpperCamelCase		    =							42
						__UpperCamelCase		    =							"""TensorFlow"""






						@property
						def      UpperCAmelCase__					( self			:Optional[int] )    ->      Optional[int]:
												return tf.__version__






						def      UpperCAmelCase__					( self			:List[Any]       ,  lowercase_			:str       ,  lowercase_			:int       ,  lowercase_			:int )    ->      List[Any]:
												UpperCAmelCase					=			self.args.strategy
												if strategy is None:
																		raise ValueError('A device strategy has to be initialized before using TensorFlow.' )
												UpperCAmelCase					=			self._prepare_inference_func(_lowerCamelCase       ,  _lowerCamelCase       ,  _lowerCamelCase )
												return self._measure_speed(_inference )






						def      UpperCAmelCase__					( self			:List[Any]       ,  lowercase_			:str       ,  lowercase_			:int       ,  lowercase_			:int )    ->      List[str]:
												UpperCAmelCase					=			self.args.strategy
												if strategy is None:
																		raise ValueError('A device strategy has to be initialized before using TensorFlow.' )
												UpperCAmelCase					=			self._prepare_train_func(_lowerCamelCase       ,  _lowerCamelCase       ,  _lowerCamelCase )
												return self._measure_speed(_train )






						def      UpperCAmelCase__					( self			:Any       ,  lowercase_			:str       ,  lowercase_			:int       ,  lowercase_			:int )    ->      Optional[Any]:







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






						def      UpperCAmelCase__					( self			:List[Any]       ,  lowercase_			:str       ,  lowercase_			:int       ,  lowercase_			:int )    ->      Dict:







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

												UpperCAmelCase					=			self._prepare_train_func(_lowerCamelCase       ,  _lowerCamelCase       ,  _lowerCamelCase )
												return self._measure_memory(_train )






						def      UpperCAmelCase__					( self			:List[str]       ,  lowercase_			:str       ,  lowercase_			:int       ,  lowercase_			:int )    ->      Tuple:
												UpperCAmelCase					=			self.config_dict[model_name]

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

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

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

												UpperCAmelCase					=			encoder_decoder_forward if config.is_encoder_decoder else encoder_forward

												return _inference






						def      UpperCAmelCase__					( self			:List[Any]       ,  lowercase_			:str       ,  lowercase_			:int       ,  lowercase_			:int )    ->      Optional[int]:
												UpperCAmelCase					=			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.' )

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

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

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

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

												UpperCAmelCase					=			encoder_decoder_train if config.is_encoder_decoder else encoder_train

												return _train






						def      UpperCAmelCase__					( self			:Union[str, Any]       ,  lowercase_			:Union[str, Any] )    ->      Dict:







												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
																								UpperCAmelCase					=			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      UpperCAmelCase__					( self			:Union[str, Any]       ,  lowercase_			:Callable[[], None] )    ->      Dict:
												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.' )
																														UpperCAmelCase					=			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.' )
																																				UpperCAmelCase					=			'''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()
																																				UpperCAmelCase					=			nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx )
																																				UpperCAmelCase					=			nvml.nvmlDeviceGetMemoryInfo(_lowerCamelCase )
																																				UpperCAmelCase					=			meminfo.used
																																				UpperCAmelCase					=			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.' )
																																				UpperCAmelCase					=			None
																														else:
																																				UpperCAmelCase					=			measure_peak_memory_cpu(_lowerCamelCase )
																																				UpperCAmelCase					=			Memory(_lowerCamelCase ) if isinstance(_lowerCamelCase       ,  _lowerCamelCase ) else memory_bytes
																								if self.args.trace_memory_line_by_line:
																														UpperCAmelCase					=			stop_memory_tracing(_lowerCamelCase )
																														if memory is None:
																																				UpperCAmelCase					=			summary.total
																								else:
																														UpperCAmelCase					=			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 importlib
import sys
from argparse import REMAINDER, ArgumentParser
from pathlib import Path

import torch_xla.distributed.xla_multiprocessing as xmp

def 			_lowerCAmelCase				(       ):
						UpperCAmelCase					=			ArgumentParser(
						    description=(
						        'PyTorch TPU distributed training launch helper utility that will spawn up multiple distributed processes'
						    )						)

						# Optional arguments for the launch helper
						parser.add_argument('--num_cores'		,     type=lowercase_		,     default=1		,     help='Number of TPU cores to use (1 or 8).'						)

						# positional
						parser.add_argument(
						    'training_script'		,     type=lowercase_		,     help=(
						        'The full path to the single TPU training '
						        'program/script to be launched in parallel, '
						        'followed by all the arguments for the '
						        'training script'
						    )		,     )

						# rest from the training program
						parser.add_argument('training_script_args'		,     nargs=lowercase_						)

						return parser.parse_args()

def 			_lowerCAmelCase				(       ):
						UpperCAmelCase					=			parse_args()

						# Import training_script as a module.
						UpperCAmelCase					=			Path(args.training_script						)
						sys.path.append(str(script_fpath.parent.resolve()						)						)
						UpperCAmelCase					=			script_fpath.stem
						UpperCAmelCase					=			importlib.import_module(lowercase_						)

						# Patch sys.argv
						UpperCAmelCase					=			[args.training_script] + args.training_script_args + ['--tpu_num_cores', str(args.num_cores						)]

						xmp.spawn(mod._mp_fn		,     args=()		,     nprocs=args.num_cores						)


if __name__ == "__main__":
				main()





