



"""simple docstring"""






import timeit
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_auto import MODEL_MAPPING, MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_torch_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_torch_available():
		import torch

		from .benchmark_args import PyTorchBenchmarkArguments


if is_pyanvml_available():
		import pyanvml.pyanvml as nvml


_lowercase					:					Optional[Any] 		=		logging.get_logger(__name__)

class 			__SCREAMING_SNAKE_CASE      ( lowerCAmelCase_  ):


					'''simple docstring'''





					_a       =				42
					_a       =				42
					_a       =				"PyTorch"



					@property
					def 						snake_case			(			self				:    Dict	)->   List[str]:
								return torch.__version__



					def 						snake_case			(			self				:    Optional[Any],    lowerCamelCase				:    str,    lowerCamelCase				:    int,    lowerCamelCase				:    int	)->   float:
								lowerCamelCase__						:						List[Any]						      =self._prepare_inference_func(lowerCamelCase,    lowerCamelCase,    lowerCamelCase	)
								return self._measure_speed(_inference	)



					def 						snake_case			(			self				:    List[str],    lowerCamelCase				:    str,    lowerCamelCase				:    int,    lowerCamelCase				:    int	)->   [Memory, Optional[MemorySummary]]:
								lowerCamelCase__						:						List[Any]						      =self._prepare_inference_func(lowerCamelCase,    lowerCamelCase,    lowerCamelCase	)
								return self._measure_memory(_inference	)



					def 						snake_case			(			self				:    Any,    lowerCamelCase				:    str,    lowerCamelCase				:    int,    lowerCamelCase				:    int	)->   float:
								lowerCamelCase__						:						int						      =self._prepare_train_func(lowerCamelCase,    lowerCamelCase,    lowerCamelCase	)
								return self._measure_speed(_train	)



					def 						snake_case			(			self				:    List[str],    lowerCamelCase				:    str,    lowerCamelCase				:    int,    lowerCamelCase				:    int	)->   [Memory, Optional[MemorySummary]]:
								lowerCamelCase__						:						str						      =self._prepare_train_func(lowerCamelCase,    lowerCamelCase,    lowerCamelCase	)
								return self._measure_memory(_train	)



					def 						snake_case			(			self				:    List[str],    lowerCamelCase				:    str,    lowerCamelCase				:    int,    lowerCamelCase				:    int	)->   Callable[[], None]:
								lowerCamelCase__						:						List[str]						      =self.config_dict[model_name]

								if self.args.torchscript:
											lowerCamelCase__						:						int						      =True

								lowerCamelCase__						:						Optional[Any]						      =(
								    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:
														lowerCamelCase__						:						int						      =config.architectures[0]
														lowerCamelCase__						:						Dict						      =__import__('''transformers''',    fromlist=[model_class]	)
														lowerCamelCase__						:						str						      =getattr(lowerCamelCase,    lowerCamelCase	)
														lowerCamelCase__						:						Tuple						      =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:
											lowerCamelCase__						:						List[str]						      =MODEL_MAPPING[config.__class__](lowerCamelCase	)

								model.eval()
								model.to(self.args.device	)

								# encoder-decoder has vocab size saved differently
								lowerCamelCase__						:						Optional[Any]						      =config.vocab_size if hasattr(lowerCamelCase,    '''vocab_size'''	) else config.encoder.vocab_size
								lowerCamelCase__						:						int						      =torch.randint(lowerCamelCase,    (batch_size, sequence_length),    dtype=torch.long,    device=self.args.device	)

								if self.args.fpaa:
											logger.info('''Running training in Mixed Precision...'''	)
											if not self.args.is_gpu:
														raise ValueError('''Mixed precision is possible only for GPU.'''	)
											# amp seems to have memory leaks so that memory usage
											# is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
											model.half()

								if self.args.torchscript:
											with torch.no_grad():
														lowerCamelCase__						:						Optional[Any]						      =torch.jit.trace(lowerCamelCase,    lowerCamelCase	)
								else:
											lowerCamelCase__						:						Tuple						      =model

								def encoder_decoder_forward():
											with torch.no_grad():
														lowerCamelCase__						:						Optional[Any]						      =inference_model(lowerCamelCase,    decoder_input_ids=lowerCamelCase	)
											return outputs

								def encoder_forward():
											with torch.no_grad():
														lowerCamelCase__						:						Any						      =inference_model(lowerCamelCase	)
											return outputs

								lowerCamelCase__						:						Optional[int]						      =encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
								return _forward



					def 						snake_case			(			self				:    Union[str, Any],    lowerCamelCase				:    str,    lowerCamelCase				:    int,    lowerCamelCase				:    int	)->   Callable[[], None]:
								lowerCamelCase__						:						Any						      =self.config_dict[model_name]

								lowerCamelCase__						:						Tuple						      =(
								    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:
														lowerCamelCase__						:						List[Any]						      =config.architectures[0]
														lowerCamelCase__						:						int						      =__import__('''transformers''',    fromlist=[model_class]	)
														lowerCamelCase__						:						int						      =getattr(lowerCamelCase,    lowerCamelCase	)
														lowerCamelCase__						:						Dict						      =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:
											lowerCamelCase__						:						List[str]						      =MODEL_WITH_LM_HEAD_MAPPING[config.__class__](lowerCamelCase	)

								if self.args.torchscript:
											raise NotImplementedError('''Training for torchscript is currently not implemented'''	)
								else:
											lowerCamelCase__						:						str						      =model

								model.train()
								model.to(self.args.device	)

								# encoder-decoder has vocab size saved differently
								lowerCamelCase__						:						Dict						      =config.vocab_size if hasattr(lowerCamelCase,    '''vocab_size'''	) else config.encoder.vocab_size
								lowerCamelCase__						:						int						      =torch.randint(lowerCamelCase,    (batch_size, sequence_length),    dtype=torch.long,    device=self.args.device	)

								if self.args.fpaa:
											logger.info('''Running training in Mixed Precision...'''	)
											if not self.args.is_gpu:
														raise ValueError('''Mixed precision is possible only for GPU.'''	)

											# amp seems to have memory leaks so that memory usage
											# is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
											model.half()

								def compute_loss_and_backprob_encoder():
											lowerCamelCase__						:						Optional[int]						      =train_model(lowerCamelCase,    labels=lowerCamelCase	)[0]
											loss.backward()
											return loss

								def compute_loss_and_backprob_encoder_decoder():
											lowerCamelCase__						:						Dict						      =train_model(lowerCamelCase,    decoder_input_ids=lowerCamelCase,    labels=lowerCamelCase	)[0]
											loss.backward()
											return loss

								lowerCamelCase__						:						Optional[Any]						      =(
								    compute_loss_and_backprob_encoder_decoder
								    if config.is_encoder_decoder
								    else compute_loss_and_backprob_encoder
								)
								return _train



					def 						snake_case			(			self				:    Optional[int],    lowerCamelCase				:    Optional[int]	)->   float:
								try:
											if self.args.is_tpu or self.args.torchscript:
														# run additional 10 times to stabilize compilation for tpu and torchscript
														logger.info('''Do inference on TPU or torchscript. 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
											lowerCamelCase__						:						List[str]						      =timeit.repeat(
											    lowerCamelCase,    repeat=self.args.repeat,    number=10,    )

											if self.args.is_tpu and self.args.torch_xla_tpu_print_metrics:
														import torch_xla.debug.metrics as met

														self.print_fn(met.metrics_report()	)

											return min(lowerCamelCase	) / 10.0
								except RuntimeError as e:
											self.print_fn(F'''Doesn\'t fit on GPU. {e}'''	)
											return "N/A"




					def 						snake_case			(			self				:    Optional[int],    lowerCamelCase				:    Callable[[], None]	)->   [Memory, MemorySummary]:
								try:
											if self.args.trace_memory_line_by_line:
														lowerCamelCase__						:						Optional[int]						      =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'''
														    ''' `--no-memory` or `args.memory=False`'''	)
											elif self.args.is_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__						:						str						      ='''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__						:						int						      =nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx	)
																	lowerCamelCase__						:						Optional[Any]						      =nvml.nvmlDeviceGetMemoryInfo(lowerCamelCase	)
																	lowerCamelCase__						:						str						      =meminfo.used
																	lowerCamelCase__						:						List[str]						      =Memory(lowerCamelCase	)
																	# shutdown nvml
																	nvml.nvmlShutdown()
											else:
														# cpu
														lowerCamelCase__						:						List[Any]						      =measure_peak_memory_cpu(lowerCamelCase	)
														lowerCamelCase__						:						List[Any]						      =Memory(lowerCamelCase	) if isinstance(lowerCamelCase,    lowerCamelCase	) else memory_bytes

											if self.args.trace_memory_line_by_line:
														lowerCamelCase__						:						List[str]						      =stop_memory_tracing(lowerCamelCase	)
											else:
														lowerCamelCase__						:						Optional[int]						      =None

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





"""simple docstring"""




import os
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
					import torch

					from transformers import GPTaLMHeadModel

@require_torch
@require_sentencepiece
@require_tokenizers
class        __SCREAMING_SNAKE_CASE     (   unittest.TestCase							):







				@slow
				@unittest.skip('''Model is not available.'''  )
				def   lowercase_		(    self    :Union[str, Any]  )      ->					Optional[int]:




									"""simple docstring"""



									lowerCamelCase__		:   Any  				=					'''nvidia/megatron-gpt2-345m/'''
									if "MYDIR" in os.environ:
														lowerCamelCase__		:   str  				=					os.path.join(os.environ['''MYDIR''']      ,__UpperCAmelCase  )
									lowerCamelCase__		:   Dict  				=					GPTaLMHeadModel.from_pretrained(__UpperCAmelCase  )
									model.to(__UpperCAmelCase  )
									model.half()

									lowerCamelCase__		:   Dict  				=					torch.tensor(
									    [[1_01, 71_10, 10_05, 10_56, 20_23, 1_13_33, 1_74_13, 10_29, 1_02]]      ,device=__UpperCAmelCase      ,dtype=torch.long      ,)

									with torch.no_grad():
														lowerCamelCase__		:   Tuple  				=					model(__UpperCAmelCase  ).logits

									lowerCamelCase__		:   List[str]  				=					torch.Size((1, 9, 5_02_57)  )
									self.assertEqual(output.shape      ,__UpperCAmelCase  )

									lowerCamelCase__		:   Tuple  				=					torch.tensor(
									    [
									        4.9_414,
									        -0.2_920,
									        -1.2_148,
									        -4.0_273,
									        -0.5_161,
									        -5.2_109,
									        -1.2_412,
									        -1.8_301,
									        -1.7_734,
									        -4.7_148,
									        -0.2_317,
									        -1.0_811,
									        -2.1_777,
									        0.4_141,
									        -3.7_969,
									        -4.0_586,
									        -2.5_332,
									        -3.3_809,
									        4.3_867,
									    ]      ,device=__UpperCAmelCase      ,dtype=torch.half      ,)

									for i in range(19  ):
														lowerCamelCase__				,						lowerCamelCase__		:   Dict  				=					8 * i // 17, 27_92 * i  # along the diagonal
														lowerCamelCase__				,						lowerCamelCase__		:   Optional[int]  				=					output[0, r, c], expected_diag[i]
														lowerCamelCase__		:   Dict  				=					F"""row={r} col={c} computed={computed} expected={expected}"""
														self.assertAlmostEqual(__UpperCAmelCase      ,__UpperCAmelCase      ,delta=1E-4      ,msg=__UpperCAmelCase  )






