


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


UpperCAmelCase_		:       Union[str, Any]							       =		logging.get_logger(__name__)





class    lowercase__  (						_snake_case ):





							'''simple docstring'''


							A_							:				PyTorchBenchmarkArguments
							A_							:				PretrainedConfig
							A_							:				str				     =  "PyTorch"






							@property
							def UpperCAmelCase_     (   self  ):
											return torch.__version__






							def UpperCAmelCase_     (   self   , __snake_case   , __snake_case   , __snake_case  ):
											_SCREAMING_SNAKE_CASE			:	List[str]							=				self._prepare_inference_func(__snake_case   , __snake_case   , __snake_case  )
											return self._measure_speed(_inference  )






							def UpperCAmelCase_     (   self   , __snake_case   , __snake_case   , __snake_case  ):
											_SCREAMING_SNAKE_CASE			:	Optional[Any]							=				self._prepare_inference_func(__snake_case   , __snake_case   , __snake_case  )
											return self._measure_memory(_inference  )






							def UpperCAmelCase_     (   self   , __snake_case   , __snake_case   , __snake_case  ):
											_SCREAMING_SNAKE_CASE			:	Union[str, Any]							=				self._prepare_train_func(__snake_case   , __snake_case   , __snake_case  )
											return self._measure_speed(_train  )






							def UpperCAmelCase_     (   self   , __snake_case   , __snake_case   , __snake_case  ):
											_SCREAMING_SNAKE_CASE			:	List[Any]							=				self._prepare_train_func(__snake_case   , __snake_case   , __snake_case  )
											return self._measure_memory(_train  )






							def UpperCAmelCase_     (   self   , __snake_case   , __snake_case   , __snake_case  ):
											_SCREAMING_SNAKE_CASE			:	Union[str, Any]							=				self.config_dict[model_name]

											if self.args.torchscript:
															_SCREAMING_SNAKE_CASE			:	Dict							=				True

											_SCREAMING_SNAKE_CASE			:	List[Any]							=				(
											    hasattr(__snake_case   , """architectures"""  )
											    and isinstance(config.architectures   , __snake_case  )
											    and len(config.architectures  ) > 0
											)
											if not self.args.only_pretrain_model and has_model_class_in_config:
															try:
																			_SCREAMING_SNAKE_CASE			:	Dict							=				config.architectures[0]
																			_SCREAMING_SNAKE_CASE			:	Any							=				__import__("""transformers"""   , fromlist=[model_class]  )
																			_SCREAMING_SNAKE_CASE			:	Optional[int]							=				getattr(__snake_case   , __snake_case  )
																			_SCREAMING_SNAKE_CASE			:	List[Any]							=				model_cls(__snake_case  )
															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:
															_SCREAMING_SNAKE_CASE			:	int							=				MODEL_MAPPING[config.__class__](__snake_case  )

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

											# encoder-decoder has vocab size saved differently
											_SCREAMING_SNAKE_CASE			:	Optional[int]							=				config.vocab_size if hasattr(__snake_case   , """vocab_size"""  ) else config.encoder.vocab_size
											_SCREAMING_SNAKE_CASE			:	Any							=				torch.randint(__snake_case   , (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():
																			_SCREAMING_SNAKE_CASE			:	int							=				torch.jit.trace(__snake_case   , __snake_case  )
											else:
															_SCREAMING_SNAKE_CASE			:	Union[str, Any]							=				model

											def encoder_decoder_forward():
															with torch.no_grad():
																			_SCREAMING_SNAKE_CASE			:	Any							=				inference_model(__snake_case   , decoder_input_ids=__snake_case  )
															return outputs

											def encoder_forward():
															with torch.no_grad():
																			_SCREAMING_SNAKE_CASE			:	Any							=				inference_model(__snake_case  )
															return outputs

											_SCREAMING_SNAKE_CASE			:	int							=				encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
											return _forward






							def UpperCAmelCase_     (   self   , __snake_case   , __snake_case   , __snake_case  ):
											_SCREAMING_SNAKE_CASE			:	Optional[Any]							=				self.config_dict[model_name]

											_SCREAMING_SNAKE_CASE			:	List[Any]							=				(
											    hasattr(__snake_case   , """architectures"""  )
											    and isinstance(config.architectures   , __snake_case  )
											    and len(config.architectures  ) > 0
											)
											if not self.args.only_pretrain_model and has_model_class_in_config:
															try:
																			_SCREAMING_SNAKE_CASE			:	Optional[Any]							=				config.architectures[0]
																			_SCREAMING_SNAKE_CASE			:	Optional[Any]							=				__import__("""transformers"""   , fromlist=[model_class]  )
																			_SCREAMING_SNAKE_CASE			:	Tuple							=				getattr(__snake_case   , __snake_case  )
																			_SCREAMING_SNAKE_CASE			:	Optional[int]							=				model_cls(__snake_case  )
															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:
															_SCREAMING_SNAKE_CASE			:	List[str]							=				MODEL_WITH_LM_HEAD_MAPPING[config.__class__](__snake_case  )

											if self.args.torchscript:
															raise NotImplementedError("""Training for torchscript is currently not implemented"""  )
											else:
															_SCREAMING_SNAKE_CASE			:	List[str]							=				model

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

											# encoder-decoder has vocab size saved differently
											_SCREAMING_SNAKE_CASE			:	List[str]							=				config.vocab_size if hasattr(__snake_case   , """vocab_size"""  ) else config.encoder.vocab_size
											_SCREAMING_SNAKE_CASE			:	Tuple							=				torch.randint(__snake_case   , (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():
															_SCREAMING_SNAKE_CASE			:	int							=				train_model(__snake_case   , labels=__snake_case  )[0]
															loss.backward()
															return loss

											def compute_loss_and_backprob_encoder_decoder():
															_SCREAMING_SNAKE_CASE			:	Any							=				train_model(__snake_case   , decoder_input_ids=__snake_case   , labels=__snake_case  )[0]
															loss.backward()
															return loss

											_SCREAMING_SNAKE_CASE			:	Tuple							=				(
											    compute_loss_and_backprob_encoder_decoder
											    if config.is_encoder_decoder
											    else compute_loss_and_backprob_encoder
											)
											return _train






							def UpperCAmelCase_     (   self   , __snake_case  ):
											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(
																			    __snake_case   , 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
															_SCREAMING_SNAKE_CASE			:	Optional[Any]							=				timeit.repeat(
															    __snake_case   , 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(__snake_case  ) / 10.0
											except RuntimeError as e:
															self.print_fn(f"""Doesn't fit on GPU. {e}"""  )
															return "N/A"






							def UpperCAmelCase_     (   self   , __snake_case  ):
											try:
															if self.args.trace_memory_line_by_line:
																			_SCREAMING_SNAKE_CASE			:	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."""  )
																							_SCREAMING_SNAKE_CASE			:	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()
																							_SCREAMING_SNAKE_CASE			:	Optional[Any]							=				nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx  )
																							_SCREAMING_SNAKE_CASE			:	Optional[int]							=				nvml.nvmlDeviceGetMemoryInfo(__snake_case  )
																							_SCREAMING_SNAKE_CASE			:	Any							=				meminfo.used
																							_SCREAMING_SNAKE_CASE			:	Any							=				Memory(__snake_case  )
																							# shutdown nvml
																							nvml.nvmlShutdown()
															else:
																			# cpu
																			_SCREAMING_SNAKE_CASE			:	int							=				measure_peak_memory_cpu(__snake_case  )
																			_SCREAMING_SNAKE_CASE			:	str							=				Memory(__snake_case  ) if isinstance(__snake_case   , __snake_case  ) else memory_bytes

															if self.args.trace_memory_line_by_line:
																			_SCREAMING_SNAKE_CASE			:	Union[str, Any]							=				stop_memory_tracing(__snake_case  )
															else:
																			_SCREAMING_SNAKE_CASE			:	Any							=				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 tempfile
import unittest
from typing import List

from transformers.models.esm.tokenization_esm import VOCAB_FILES_NAMES, EsmTokenizer
from transformers.testing_utils import require_tokenizers
from transformers.tokenization_utils import PreTrainedTokenizer
from transformers.tokenization_utils_base import PreTrainedTokenizerBase





@require_tokenizers
class    lowercase__  (						unittest.TestCase ):





							'''simple docstring'''


							A_							:				Optional[Any]				     =  EsmTokenizer






							def UpperCAmelCase_     (   self  ):
											super().setUp()
											_SCREAMING_SNAKE_CASE			:	List[Any]							=				tempfile.mkdtemp()
											# fmt: off
											_SCREAMING_SNAKE_CASE			:	List[str]							=				["<cls>", "<pad>", "<eos>", "<unk>", "L", "A", "G", "V", "S", "E", "R", "T", "I", "D", "P", "K", "Q", "N", "F", "Y", "M", "H", "W", "C", "X", "B", "U", "Z", "O", ".", "-", "<null_1>", "<mask>"]  # noqa: E501
											# fmt: on
											_SCREAMING_SNAKE_CASE			:	Any							=				os.path.join(self.tmpdirname   , VOCAB_FILES_NAMES["""vocab_file"""]  )
											with open(self.vocab_file   , """w"""   , encoding="""utf-8"""  ) as vocab_writer:
															vocab_writer.write("""""".join([x + """\n""" for x in vocab_tokens]  )  )






							def UpperCAmelCase_     (   self   , **__snake_case  ):
											return [self.get_tokenizer(**__snake_case  )]






							def UpperCAmelCase_     (   self   , **__snake_case  ):
											return self.tokenizer_class.from_pretrained(self.tmpdirname   , **__snake_case  )






							def UpperCAmelCase_     (   self  ):
											_SCREAMING_SNAKE_CASE			:	List[str]							=				self.tokenizer_class(self.vocab_file  )

											_SCREAMING_SNAKE_CASE			:	int							=				tokenizer.tokenize("""LAGVS"""  )
											self.assertListEqual(__snake_case   , ["""L""", """A""", """G""", """V""", """S"""]  )
											self.assertListEqual(tokenizer.convert_tokens_to_ids(__snake_case  )   , [4, 5, 6, 7, 8]  )






							def UpperCAmelCase_     (   self  ):
											_SCREAMING_SNAKE_CASE			:	int							=				self.tokenizer_class(self.vocab_file  )

											_SCREAMING_SNAKE_CASE			:	Optional[Any]							=				"""LAGVS"""
											self.assertListEqual(tokenizer.encode(__snake_case  )   , [0, 4, 5, 6, 7, 8, 2]  )






							def UpperCAmelCase_     (   self  ):
											_SCREAMING_SNAKE_CASE			:	Any							=				self.tokenizer_class(self.vocab_file  )

											_SCREAMING_SNAKE_CASE			:	List[Any]							=				["""LAGVS""", """WCB"""]
											_SCREAMING_SNAKE_CASE			:	Union[str, Any]							=				tokenizer(__snake_case   , padding=__snake_case  )["""input_ids"""]

											self.assertListEqual(__snake_case   , [[0, 4, 5, 6, 7, 8, 2], [0, 22, 23, 25, 2]]  )






							def UpperCAmelCase_     (   self  ):
											_SCREAMING_SNAKE_CASE			:	Union[str, Any]							=				self.tokenizer_class(self.vocab_file  )

											_SCREAMING_SNAKE_CASE			:	str							=				["""LAGVS""", """WCB"""]
											_SCREAMING_SNAKE_CASE			:	Optional[int]							=				tokenizer(__snake_case   , padding=__snake_case  )["""input_ids"""]

											self.assertListEqual(__snake_case   , [[0, 4, 5, 6, 7, 8, 2], [0, 22, 23, 25, 2, 1, 1]]  )






							def UpperCAmelCase_     (   self  ):
											_SCREAMING_SNAKE_CASE			:	str							=				self.get_tokenizers(fast=__snake_case  )
											for tokenizer in tokenizers:
															with self.subTest(f"""{tokenizer.__class__.__name__}"""  ):
																			_SCREAMING_SNAKE_CASE			:	Optional[int]							=				"""<unk>"""
																			_SCREAMING_SNAKE_CASE			:	Optional[int]							=				"""<mask>"""

																			_SCREAMING_SNAKE_CASE			:	Union[str, Any]							=				tokenizer.tokenize(__snake_case  )
																			_SCREAMING_SNAKE_CASE			:	Optional[Any]							=				tokenizer.tokenize(__snake_case  )

																			self.assertEqual(len(__snake_case  )   , 1  )
																			self.assertEqual(len(__snake_case  )   , 1  )
																			self.assertEqual(token_a[0]   , __snake_case  )
																			self.assertEqual(token_a[0]   , __snake_case  )




