

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


lowerCamelCase 	=   logging.get_logger(__name__)







class 		__magic_name__     (				snake_case__				):

   '''simple docstring'''




   lowerCamelCase__					:      PyTorchBenchmarkArguments
   lowerCamelCase__					:      PretrainedConfig
   lowerCamelCase__					:      str        =  "PyTorch"



   @property
   def 					_UpperCAmelCase					(							self			)				->       List[Any]:






     """simple docstring"""






     return torch.__version__



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_			)				->       float:






     """simple docstring"""






     a__    =self._prepare_inference_func(lowercase_,       lowercase_,       lowercase_			)
     return self._measure_speed(_inference			)



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_			)				->       [Memory, Optional[MemorySummary]]:






     """simple docstring"""






     a__    =self._prepare_inference_func(lowercase_,       lowercase_,       lowercase_			)
     return self._measure_memory(_inference			)



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_			)				->       float:






     """simple docstring"""






     a__    =self._prepare_train_func(lowercase_,       lowercase_,       lowercase_			)
     return self._measure_speed(_train			)



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_			)				->       [Memory, Optional[MemorySummary]]:






     """simple docstring"""






     a__    =self._prepare_train_func(lowercase_,       lowercase_,       lowercase_			)
     return self._measure_memory(_train			)



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_			)				->       Callable[[], None]:






     """simple docstring"""






     a__    =self.config_dict[model_name]

     if self.args.torchscript:
       a__    =True

     a__    =(
         hasattr(lowercase_,       '''architectures'''			)
         and isinstance(config.architectures,       lowercase_			)
         and len(config.architectures			) > 0
     )
     if not self.args.only_pretrain_model and has_model_class_in_config:
       try:
         a__    =config.architectures[0]
         a__    =__import__('''transformers''',       fromlist=[model_class]			)
         a__    =getattr(lowercase_,       lowercase_			)
         a__    =model_cls(lowercase_			)
       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:
       a__    =MODEL_MAPPING[config.__class__](lowercase_			)

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

     # encoder-decoder has vocab size saved differently
     a__    =config.vocab_size if hasattr(lowercase_,       '''vocab_size'''			) else config.encoder.vocab_size
     a__    =torch.randint(lowercase_,       (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():
         a__    =torch.jit.trace(lowercase_,       lowercase_			)
     else:
       a__    =model

     def encoder_decoder_forward():
       with torch.no_grad():
         a__    =inference_model(lowercase_,       decoder_input_ids=lowercase_			)
       return outputs

     def encoder_forward():
       with torch.no_grad():
         a__    =inference_model(lowercase_			)
       return outputs

     a__    =encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
     return _forward



   def 					_UpperCAmelCase					(							self,       lowercase_,       lowercase_,       lowercase_			)				->       Callable[[], None]:






     """simple docstring"""






     a__    =self.config_dict[model_name]

     a__    =(
         hasattr(lowercase_,       '''architectures'''			)
         and isinstance(config.architectures,       lowercase_			)
         and len(config.architectures			) > 0
     )
     if not self.args.only_pretrain_model and has_model_class_in_config:
       try:
         a__    =config.architectures[0]
         a__    =__import__('''transformers''',       fromlist=[model_class]			)
         a__    =getattr(lowercase_,       lowercase_			)
         a__    =model_cls(lowercase_			)
       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:
       a__    =MODEL_WITH_LM_HEAD_MAPPING[config.__class__](lowercase_			)

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

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

     # encoder-decoder has vocab size saved differently
     a__    =config.vocab_size if hasattr(lowercase_,       '''vocab_size'''			) else config.encoder.vocab_size
     a__    =torch.randint(lowercase_,       (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():
       a__    =train_model(lowercase_,       labels=lowercase_			)[0]
       loss.backward()
       return loss

     def compute_loss_and_backprob_encoder_decoder():
       a__    =train_model(lowercase_,       decoder_input_ids=lowercase_,       labels=lowercase_			)[0]
       loss.backward()
       return loss

     a__    =(
         compute_loss_and_backprob_encoder_decoder
         if config.is_encoder_decoder
         else compute_loss_and_backprob_encoder
     )
     return _train



   def 					_UpperCAmelCase					(							self,       lowercase_			)				->       float:






     """simple docstring"""






     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(
             lowercase_,       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
       a__    =timeit.repeat(
           lowercase_,       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(lowercase_			) / 10.0
     except RuntimeError as e:
       self.print_fn(F"""Doesn\'t fit on GPU. {e}"""			)
       return "N/A"



   def 					_UpperCAmelCase					(							self,       lowercase_			)				->       [Memory, MemorySummary]:






     """simple docstring"""






     try:
       if self.args.trace_memory_line_by_line:
         a__    =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.'''			)
           a__    ='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()
           a__    =nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx			)
           a__    =nvml.nvmlDeviceGetMemoryInfo(lowercase_			)
           a__    =meminfo.used
           a__    =Memory(lowercase_			)
           # shutdown nvml
           nvml.nvmlShutdown()
       else:
         # cpu
         a__    =measure_peak_memory_cpu(lowercase_			)
         a__    =Memory(lowercase_			) if isinstance(lowercase_,       lowercase_			) else memory_bytes

       if self.args.trace_memory_line_by_line:
         a__    =stop_memory_tracing(lowercase_			)
       else:
         a__    =None

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



def _a					(			SCREAMING_SNAKE_CASE						:     int						)       ->       int:

 """simple docstring"""






 return n * (2 * n - 1)







def _a					(			SCREAMING_SNAKE_CASE						:     int						)       ->       bool:

 """simple docstring"""






 __lowerCAmelCase:  List[Any]        = (1 + 24 * n) ** 0.5
 return ((1 + root) / 6) % 1 == 0







def _a					(			SCREAMING_SNAKE_CASE						:     int = 1_44						)       ->       int:

 """simple docstring"""






 __lowerCAmelCase:  Dict        = start
 __lowerCAmelCase:  List[str]        = hexagonal_num(SCREAMING_SNAKE_CASE						)
 while not is_pentagonal(SCREAMING_SNAKE_CASE						):
  n += 1
  __lowerCAmelCase:  Any        = hexagonal_num(SCREAMING_SNAKE_CASE						)
 return num


if __name__ == "__main__":
  print(f"{solution()} = ")





