





'''simple docstring'''


import json
import os
from functools import lru_cache
from typing import TYPE_CHECKING, List, Optional, Tuple

import regex as re

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


if TYPE_CHECKING:
    from transformers.pipelines.conversational import Conversation

A_				     =							logging.get_logger(__name__)

A_				     =							{
    "vocab_file": "vocab.json",
    "merges_file": "merges.txt",
}

A_				     =							{
    "vocab_file": {
        "gpt2": "https://huggingface.co/gpt2/resolve/main/vocab.json",
        "gpt2-medium": "https://huggingface.co/gpt2-medium/resolve/main/vocab.json",
        "gpt2-large": "https://huggingface.co/gpt2-large/resolve/main/vocab.json",
        "gpt2-xl": "https://huggingface.co/gpt2-xl/resolve/main/vocab.json",
        "distilgpt2": "https://huggingface.co/distilgpt2/resolve/main/vocab.json",
    },
    "merges_file": {
        "gpt2": "https://huggingface.co/gpt2/resolve/main/merges.txt",
        "gpt2-medium": "https://huggingface.co/gpt2-medium/resolve/main/merges.txt",
        "gpt2-large": "https://huggingface.co/gpt2-large/resolve/main/merges.txt",
        "gpt2-xl": "https://huggingface.co/gpt2-xl/resolve/main/merges.txt",
        "distilgpt2": "https://huggingface.co/distilgpt2/resolve/main/merges.txt",
    },
}

A_				     =							{
    "gpt2": 10_24,
    "gpt2-medium": 10_24,
    "gpt2-large": 10_24,
    "gpt2-xl": 10_24,
    "distilgpt2": 10_24,
}




@lru_cache()
def       A							(      )  ->				Optional[Any]:


  '''simple docstring'''




  __lowerCAmelCase		:    List[str]			     =  (
      list(range(ord('!'					)						,ord('~'					) + 1					)					) + list(range(ord('¡'					)						,ord('¬'					) + 1					)					) + list(range(ord('®'					)						,ord('ÿ'					) + 1					)					)
  )
  __lowerCAmelCase		:    Tuple			     =  bs[:]
  __lowerCAmelCase		:    Union[str, Any]			     =  0
  for b in range(2**8					):
    if b not in bs:
      bs.append(_UpperCAmelCase					)
      cs.append(2**8 + n					)
      n += 1
  __lowerCAmelCase		:    Tuple			     =  [chr(_UpperCAmelCase					) for n in cs]
  return dict(zip(_UpperCAmelCase						,_UpperCAmelCase					)					)




def       A							(      _UpperCAmelCase   :				Dict					)  ->				Dict:


  '''simple docstring'''




  __lowerCAmelCase		:    List[str]			     =  set()
  __lowerCAmelCase		:    List[Any]			     =  word[0]
  for char in word[1:]:
    pairs.add((prev_char, char)					)
    __lowerCAmelCase		:    Any			     =  char
  return pairs


class 		UpperCamelCase__     (       a							):
      '''simple docstring'''






      _snake_case           =				VOCAB_FILES_NAMES
      _snake_case           =				PRETRAINED_VOCAB_FILES_MAP
      _snake_case           =				PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      _snake_case           =				['''input_ids''', '''attention_mask''']




      def __init__(     self     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE="replace"     , SCREAMING_SNAKE_CASE="<|endoftext|>"     , SCREAMING_SNAKE_CASE="<|endoftext|>"     , SCREAMING_SNAKE_CASE="<|endoftext|>"     , SCREAMING_SNAKE_CASE=None     , SCREAMING_SNAKE_CASE=False     , SCREAMING_SNAKE_CASE=False     , **SCREAMING_SNAKE_CASE     , )				->  Optional[Any]:
        __lowerCAmelCase		:    int			     =  AddedToken(SCREAMING_SNAKE_CASE     , lstrip=SCREAMING_SNAKE_CASE     , rstrip=SCREAMING_SNAKE_CASE ) if isinstance(SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE ) else bos_token
        __lowerCAmelCase		:    List[str]			     =  AddedToken(SCREAMING_SNAKE_CASE     , lstrip=SCREAMING_SNAKE_CASE     , rstrip=SCREAMING_SNAKE_CASE ) if isinstance(SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE ) else eos_token
        __lowerCAmelCase		:    Optional[Any]			     =  AddedToken(SCREAMING_SNAKE_CASE     , lstrip=SCREAMING_SNAKE_CASE     , rstrip=SCREAMING_SNAKE_CASE ) if isinstance(SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE ) else unk_token
        __lowerCAmelCase		:    str			     =  AddedToken(SCREAMING_SNAKE_CASE     , lstrip=SCREAMING_SNAKE_CASE     , rstrip=SCREAMING_SNAKE_CASE ) if isinstance(SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE ) else pad_token
        super().__init__(
            errors=SCREAMING_SNAKE_CASE     , unk_token=SCREAMING_SNAKE_CASE     , bos_token=SCREAMING_SNAKE_CASE     , eos_token=SCREAMING_SNAKE_CASE     , pad_token=SCREAMING_SNAKE_CASE     , add_prefix_space=SCREAMING_SNAKE_CASE     , add_bos_token=SCREAMING_SNAKE_CASE     , **SCREAMING_SNAKE_CASE     , )
        __lowerCAmelCase		:    str			     =  add_bos_token

        with open(SCREAMING_SNAKE_CASE     , encoding='utf-8' ) as vocab_handle:
          __lowerCAmelCase		:    int			     =  json.load(SCREAMING_SNAKE_CASE )
        __lowerCAmelCase		:    List[Any]			     =  {v: k for k, v in self.encoder.items()}
        __lowerCAmelCase		:    List[Any]			     =  errors  # how to handle errors in decoding
        __lowerCAmelCase		:    List[Any]			     =  bytes_to_unicode()
        __lowerCAmelCase		:    str			     =  {v: k for k, v in self.byte_encoder.items()}
        with open(SCREAMING_SNAKE_CASE     , encoding='utf-8' ) as merges_handle:
          __lowerCAmelCase		:    List[Any]			     =  merges_handle.read().split('\n' )[1:-1]
        __lowerCAmelCase		:    int			     =  [tuple(merge.split() ) for merge in bpe_merges]
        __lowerCAmelCase		:    Any			     =  dict(zip(SCREAMING_SNAKE_CASE     , range(len(SCREAMING_SNAKE_CASE ) ) ) )
        __lowerCAmelCase		:    Optional[Any]			     =  {}
        __lowerCAmelCase		:    Optional[int]			     =  add_prefix_space

        # Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
        __lowerCAmelCase		:    Tuple			     =  re.compile(r'\'s|\'t|\'re|\'ve|\'m|\'ll|\'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+' )




      @property
      def 	snake_case  (     self )				->  Union[str, Any]:
        return len(self.encoder )




      def 	snake_case  (     self )				->  int:
        return dict(self.encoder     , **self.added_tokens_encoder )




      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE )				->  Tuple:
        if token in self.cache:
          return self.cache[token]
        __lowerCAmelCase		:    Any			     =  tuple(SCREAMING_SNAKE_CASE )
        __lowerCAmelCase		:    Union[str, Any]			     =  get_pairs(SCREAMING_SNAKE_CASE )

        if not pairs:
          return token

        while True:
          __lowerCAmelCase		:    List[Any]			     =  min(SCREAMING_SNAKE_CASE     , key=lambda SCREAMING_SNAKE_CASE : self.bpe_ranks.get(SCREAMING_SNAKE_CASE     , float('inf' ) ) )
          if bigram not in self.bpe_ranks:
            break
          __lowerCAmelCase     ,					__lowerCAmelCase		:    Union[str, Any]			     =  bigram
          __lowerCAmelCase		:    List[str]			     =  []
          __lowerCAmelCase		:    List[Any]			     =  0
          while i < len(SCREAMING_SNAKE_CASE ):
            try:
              __lowerCAmelCase		:    List[str]			     =  word.index(SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )
            except ValueError:
              new_word.extend(word[i:] )
              break
            else:
              new_word.extend(word[i:j] )
              __lowerCAmelCase		:    Optional[Any]			     =  j

            if word[i] == first and i < len(SCREAMING_SNAKE_CASE ) - 1 and word[i + 1] == second:
              new_word.append(first + second )
              i += 2
            else:
              new_word.append(word[i] )
              i += 1
          __lowerCAmelCase		:    Optional[int]			     =  tuple(SCREAMING_SNAKE_CASE )
          __lowerCAmelCase		:    Optional[int]			     =  new_word
          if len(SCREAMING_SNAKE_CASE ) == 1:
            break
          else:
            __lowerCAmelCase		:    Union[str, Any]			     =  get_pairs(SCREAMING_SNAKE_CASE )
        __lowerCAmelCase		:    Union[str, Any]			     =  ' '.join(SCREAMING_SNAKE_CASE )
        __lowerCAmelCase		:    Dict			     =  word
        return word




      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE=None )				->  Union[str, Any]:
        if self.add_bos_token:
          __lowerCAmelCase		:    Any			     =  [self.bos_token_id]
        else:
          __lowerCAmelCase		:    int			     =  []

        __lowerCAmelCase		:    str			     =  bos_token_ids + token_ids_a

        if token_ids_a is None:
          return output

        return output + bos_token_ids + token_ids_a




      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE = None     , SCREAMING_SNAKE_CASE = False )				->  List[int]:
        if already_has_special_tokens:
          return super().get_special_tokens_mask(
              token_ids_a=SCREAMING_SNAKE_CASE     , token_ids_a=SCREAMING_SNAKE_CASE     , already_has_special_tokens=SCREAMING_SNAKE_CASE )

        if not self.add_bos_token:
          return super().get_special_tokens_mask(
              token_ids_a=SCREAMING_SNAKE_CASE     , token_ids_a=SCREAMING_SNAKE_CASE     , already_has_special_tokens=SCREAMING_SNAKE_CASE )

        if token_ids_a is None:
          return [1] + ([0] * len(SCREAMING_SNAKE_CASE ))
        return [1] + ([0] * len(SCREAMING_SNAKE_CASE )) + [1] + ([0] * len(SCREAMING_SNAKE_CASE ))




      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE )				->  Any:
        __lowerCAmelCase		:    Dict			     =  []
        for token in re.findall(self.pat     , SCREAMING_SNAKE_CASE ):
          __lowerCAmelCase		:    List[str]			     =  ''.join(
              self.byte_encoder[b] for b in token.encode('utf-8' ) )  # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
          bpe_tokens.extend(bpe_token for bpe_token in self.bpe(SCREAMING_SNAKE_CASE ).split(' ' ) )
        return bpe_tokens




      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE )				->  str:
        return self.encoder.get(SCREAMING_SNAKE_CASE     , self.encoder.get(self.unk_token ) )




      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE )				->  str:
        return self.decoder.get(SCREAMING_SNAKE_CASE )




      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE )				->  int:
        __lowerCAmelCase		:    str			     =  ''.join(SCREAMING_SNAKE_CASE )
        __lowerCAmelCase		:    str			     =  bytearray([self.byte_decoder[c] for c in text] ).decode('utf-8'     , errors=self.errors )
        return text




      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE = None )				->  Tuple[str]:
        if not os.path.isdir(SCREAMING_SNAKE_CASE ):
          logger.error(F"""Vocabulary path ({save_directory}) should be a directory""" )
          return
        __lowerCAmelCase		:    Tuple			     =  os.path.join(
            SCREAMING_SNAKE_CASE     , (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file'] )
        __lowerCAmelCase		:    List[Any]			     =  os.path.join(
            SCREAMING_SNAKE_CASE     , (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['merges_file'] )

        with open(SCREAMING_SNAKE_CASE     , 'w'     , encoding='utf-8' ) as f:
          f.write(json.dumps(self.encoder     , indent=2     , sort_keys=SCREAMING_SNAKE_CASE     , ensure_ascii=SCREAMING_SNAKE_CASE ) + '\n' )

        __lowerCAmelCase		:    Union[str, Any]			     =  0
        with open(SCREAMING_SNAKE_CASE     , 'w'     , encoding='utf-8' ) as writer:
          writer.write('#version: 0.2\n' )
          for bpe_tokens, token_index in sorted(self.bpe_ranks.items()     , key=lambda SCREAMING_SNAKE_CASE : kv[1] ):
            if index != token_index:
              logger.warning(
                  F"""Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."""
                  ' Please check that the tokenizer is not corrupted!' )
              __lowerCAmelCase		:    List[Any]			     =  token_index
            writer.write(' '.join(SCREAMING_SNAKE_CASE ) + '\n' )
            index += 1

        return vocab_file, merge_file




      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE=False     , **SCREAMING_SNAKE_CASE )				->  Dict:
        __lowerCAmelCase		:    Any			     =  kwargs.pop('add_prefix_space'     , self.add_prefix_space )
        if is_split_into_words or add_prefix_space:
          __lowerCAmelCase		:    Any			     =  ' ' + text
        return (text, kwargs)







      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE )				->  List[int]:
        __lowerCAmelCase		:    List[Any]			     =  []
        for is_user, text in conversation.iter_texts():
          input_ids.extend(self.encode(SCREAMING_SNAKE_CASE     , add_special_tokens=SCREAMING_SNAKE_CASE ) + [self.eos_token_id] )
        if len(SCREAMING_SNAKE_CASE ) > self.model_max_length:
          __lowerCAmelCase		:    List[str]			     =  input_ids[-self.model_max_length :]
        return input_ids












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


A_				     =							logging.get_logger(__name__)


class 		UpperCamelCase__     (       a							):
      '''simple docstring'''






      _snake_case           =				42
      _snake_case           =				42
      _snake_case           =				"PyTorch"




      @property
      def 	snake_case  (     self )				->  Dict:
        return torch.__version__




      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )				->  float:
        __lowerCAmelCase		:    Any			     =  self._prepare_inference_func(SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )
        return self._measure_speed(_inference )




      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )				->  [Memory, Optional[MemorySummary]]:
        __lowerCAmelCase		:    Dict			     =  self._prepare_inference_func(SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )
        return self._measure_memory(_inference )




      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )				->  float:
        __lowerCAmelCase		:    Optional[Any]			     =  self._prepare_train_func(SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )
        return self._measure_speed(_train )




      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )				->  [Memory, Optional[MemorySummary]]:
        __lowerCAmelCase		:    Optional[int]			     =  self._prepare_train_func(SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )
        return self._measure_memory(_train )




      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )				->  Callable[[], None]:
        __lowerCAmelCase		:    str			     =  self.config_dict[model_name]

        if self.args.torchscript:
          __lowerCAmelCase		:    Tuple			     =  True

        __lowerCAmelCase		:    Union[str, Any]			     =  (
            hasattr(SCREAMING_SNAKE_CASE     , 'architectures' )
            and isinstance(config.architectures     , SCREAMING_SNAKE_CASE )
            and len(config.architectures ) > 0
        )
        if not self.args.only_pretrain_model and has_model_class_in_config:
          try:
            __lowerCAmelCase		:    Dict			     =  config.architectures[0]
            __lowerCAmelCase		:    List[str]			     =  __import__('transformers'     , fromlist=[model_class] )
            __lowerCAmelCase		:    Tuple			     =  getattr(SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )
            __lowerCAmelCase		:    Dict			     =  model_cls(SCREAMING_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:
          __lowerCAmelCase		:    List[Any]			     =  MODEL_MAPPING[config.__class__](SCREAMING_SNAKE_CASE )

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

        # encoder-decoder has vocab size saved differently
        __lowerCAmelCase		:    List[str]			     =  config.vocab_size if hasattr(SCREAMING_SNAKE_CASE     , 'vocab_size' ) else config.encoder.vocab_size
        __lowerCAmelCase		:    Optional[Any]			     =  torch.randint(SCREAMING_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():
            __lowerCAmelCase		:    List[Any]			     =  torch.jit.trace(SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )
        else:
          __lowerCAmelCase		:    str			     =  model

        def encoder_decoder_forward():
          with torch.no_grad():
            __lowerCAmelCase		:    str			     =  inference_model(SCREAMING_SNAKE_CASE     , decoder_input_ids=SCREAMING_SNAKE_CASE )
          return outputs

        def encoder_forward():
          with torch.no_grad():
            __lowerCAmelCase		:    str			     =  inference_model(SCREAMING_SNAKE_CASE )
          return outputs

        __lowerCAmelCase		:    Optional[Any]			     =  encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
        return _forward




      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )				->  Callable[[], None]:
        __lowerCAmelCase		:    int			     =  self.config_dict[model_name]

        __lowerCAmelCase		:    Optional[int]			     =  (
            hasattr(SCREAMING_SNAKE_CASE     , 'architectures' )
            and isinstance(config.architectures     , SCREAMING_SNAKE_CASE )
            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		:    List[str]			     =  __import__('transformers'     , fromlist=[model_class] )
            __lowerCAmelCase		:    Tuple			     =  getattr(SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE )
            __lowerCAmelCase		:    Any			     =  model_cls(SCREAMING_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:
          __lowerCAmelCase		:    List[str]			     =  MODEL_WITH_LM_HEAD_MAPPING[config.__class__](SCREAMING_SNAKE_CASE )

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

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

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

        def compute_loss_and_backprob_encoder_decoder():
          __lowerCAmelCase		:    List[Any]			     =  train_model(SCREAMING_SNAKE_CASE     , decoder_input_ids=SCREAMING_SNAKE_CASE     , labels=SCREAMING_SNAKE_CASE )[0]
          loss.backward()
          return loss

        __lowerCAmelCase		:    Dict			     =  (
            compute_loss_and_backprob_encoder_decoder
            if config.is_encoder_decoder
            else compute_loss_and_backprob_encoder
        )
        return _train




      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE )				->  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(
                SCREAMING_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
          __lowerCAmelCase		:    Tuple			     =  timeit.repeat(
              SCREAMING_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(SCREAMING_SNAKE_CASE ) / 1_0.0
        except RuntimeError as e:
          self.print_fn(F"""Doesn't fit on GPU. {e}""" )
          return "N/A"







      def 	snake_case  (     self     , SCREAMING_SNAKE_CASE )				->  [Memory, MemorySummary]:
        try:
          if self.args.trace_memory_line_by_line:
            __lowerCAmelCase		:    Tuple			     =  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		:    Union[str, Any]			     =  '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		:    List[str]			     =  nvml.nvmlDeviceGetMemoryInfo(SCREAMING_SNAKE_CASE )
              __lowerCAmelCase		:    Tuple			     =  meminfo.used
              __lowerCAmelCase		:    Dict			     =  Memory(SCREAMING_SNAKE_CASE )
              # shutdown nvml
              nvml.nvmlShutdown()
          else:
            # cpu
            __lowerCAmelCase		:    Tuple			     =  measure_peak_memory_cpu(SCREAMING_SNAKE_CASE )
            __lowerCAmelCase		:    str			     =  Memory(SCREAMING_SNAKE_CASE ) if isinstance(SCREAMING_SNAKE_CASE     , SCREAMING_SNAKE_CASE ) else memory_bytes

          if self.args.trace_memory_line_by_line:
            __lowerCAmelCase		:    int			     =  stop_memory_tracing(SCREAMING_SNAKE_CASE )
          else:
            __lowerCAmelCase		:    Any			     =  None

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





