

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    	=			logging.get_logger(__name__)


class        _UpperCAmelCase (   __lowercase			):
  '''simple docstring'''







  SCREAMING_SNAKE_CASE							:				PyTorchBenchmarkArguments
  SCREAMING_SNAKE_CASE							:				PretrainedConfig
  SCREAMING_SNAKE_CASE							:				str					      =				"PyTorch"







  @property
  def 	UpperCamelCase       ( self	: Optional[Any] ):
        return torch.__version__







  def 	UpperCamelCase       ( self	: Tuple		,						UpperCamelCase__	: str		,						UpperCamelCase__	: int		,						UpperCamelCase__	: int ):
        A    =				self._prepare_inference_func(UpperCamelCase__		,						UpperCamelCase__		,						UpperCamelCase__ )
        return self._measure_speed(_inference )







  def 	UpperCamelCase       ( self	: Dict		,						UpperCamelCase__	: str		,						UpperCamelCase__	: int		,						UpperCamelCase__	: int ):
        A    =				self._prepare_inference_func(UpperCamelCase__		,						UpperCamelCase__		,						UpperCamelCase__ )
        return self._measure_memory(_inference )







  def 	UpperCamelCase       ( self	: Optional[int]		,						UpperCamelCase__	: str		,						UpperCamelCase__	: int		,						UpperCamelCase__	: int ):
        A    =				self._prepare_train_func(UpperCamelCase__		,						UpperCamelCase__		,						UpperCamelCase__ )
        return self._measure_speed(_train )







  def 	UpperCamelCase       ( self	: Optional[Any]		,						UpperCamelCase__	: str		,						UpperCamelCase__	: int		,						UpperCamelCase__	: int ):
        A    =				self._prepare_train_func(UpperCamelCase__		,						UpperCamelCase__		,						UpperCamelCase__ )
        return self._measure_memory(_train )







  def 	UpperCamelCase       ( self	: Tuple		,						UpperCamelCase__	: str		,						UpperCamelCase__	: int		,						UpperCamelCase__	: int ):
        A    =				self.config_dict[model_name]

        if self.args.torchscript:
              A    =				True

        A    =				(
            hasattr(UpperCamelCase__		,						'architectures' )
            and isinstance(config.architectures		,						UpperCamelCase__ )
            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(UpperCamelCase__		,						UpperCamelCase__ )
                    A    =				model_cls(UpperCamelCase__ )
              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__](UpperCamelCase__ )

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

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

        def encoder_decoder_forward():
              with torch.no_grad():
                    A    =				inference_model(UpperCamelCase__		,						decoder_input_ids=UpperCamelCase__ )
              return outputs

        def encoder_forward():
              with torch.no_grad():
                    A    =				inference_model(UpperCamelCase__ )
              return outputs

        A    =				encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
        return _forward







  def 	UpperCamelCase       ( self	: Optional[Any]		,						UpperCamelCase__	: str		,						UpperCamelCase__	: int		,						UpperCamelCase__	: int ):
        A    =				self.config_dict[model_name]

        A    =				(
            hasattr(UpperCamelCase__		,						'architectures' )
            and isinstance(config.architectures		,						UpperCamelCase__ )
            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(UpperCamelCase__		,						UpperCamelCase__ )
                    A    =				model_cls(UpperCamelCase__ )
              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__](UpperCamelCase__ )

        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(UpperCamelCase__		,						'vocab_size' ) else config.encoder.vocab_size
        A    =				torch.randint(UpperCamelCase__		,						(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(UpperCamelCase__		,						labels=UpperCamelCase__ )[0]
              loss.backward()
              return loss

        def compute_loss_and_backprob_encoder_decoder():
              A    =				train_model(UpperCamelCase__		,						decoder_input_ids=UpperCamelCase__		,						labels=UpperCamelCase__ )[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	: List[Any]		,						UpperCamelCase__	: str ):
        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(
                        UpperCamelCase__		,						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(
                  UpperCamelCase__		,						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(UpperCamelCase__ ) / 10.0
        except RuntimeError as e:
              self.print_fn(f'''Doesn\'t fit on GPU. {e}''' )
              return "N/A"







  def 	UpperCamelCase       ( self	: Tuple		,						UpperCamelCase__	: Callable[[], None] ):
        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(UpperCamelCase__ )
                          A    =				meminfo.used
                          A    =				Memory(UpperCamelCase__ )
                          # shutdown nvml
                          nvml.nvmlShutdown()
              else:
                    # cpu
                    A    =				measure_peak_memory_cpu(UpperCamelCase__ )
                    A    =				Memory(UpperCamelCase__ ) if isinstance(UpperCamelCase__		,						UpperCamelCase__ ) else memory_bytes

              if self.args.trace_memory_line_by_line:
                    A    =				stop_memory_tracing(UpperCamelCase__ )
              else:
                    A    =				None

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






import logging
import os
import sys
from dataclasses import dataclass, field
from typing import Optional

import transformers
from transformers import (
    AutoConfig,
    AutoModelForQuestionAnswering,
    AutoTokenizer,
    DataCollatorWithPadding,
    HfArgumentParser,
    SquadDataset,
    Trainer,
    TrainingArguments,
)
from transformers import SquadDataTrainingArguments as DataTrainingArguments
from transformers.trainer_utils import is_main_process


_UpperCAmelCase    	=			logging.getLogger(__name__)



@dataclass
class        _UpperCAmelCase :
  '''simple docstring'''







  SCREAMING_SNAKE_CASE							:				str					      =				field(
      metadata={'''help''': '''Path to pretrained model or model identifier from huggingface.co/models'''}			)
  SCREAMING_SNAKE_CASE							:				Optional[str]					      =				field(
      default=__lowercase     ,				metadata={'''help''': '''Pretrained config name or path if not the same as model_name'''}			)
  SCREAMING_SNAKE_CASE							:				Optional[str]					      =				field(
      default=__lowercase     ,				metadata={'''help''': '''Pretrained tokenizer name or path if not the same as model_name'''}			)
  SCREAMING_SNAKE_CASE							:				bool					      =				field(default=__lowercase     ,				metadata={'''help''': '''Set this flag to use fast tokenization.'''}			)
  # If you want to tweak more attributes on your tokenizer, you should do it in a distinct script,
  # or just modify its tokenizer_config.json.
  SCREAMING_SNAKE_CASE							:				Optional[str]					      =				field(
      default=__lowercase     ,				metadata={'''help''': '''Where do you want to store the pretrained models downloaded from huggingface.co'''}     ,				)


def 			__UpperCamelCase    ()     ->  Tuple:
      # See all possible arguments in src/transformers/training_args.py
      # or by passing the --help flag to this script.
      # We now keep distinct sets of args, for a cleaner separation of concerns.

      A    =				HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)     )

      if len(sys.argv     ) == 2 and sys.argv[1].endswith('.json'     ):
            # If we pass only one argument to the script and it's the path to a json file,
            # let's parse it to get our arguments.
            A ,	A ,	A    =				parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]     )     )
      else:
            A ,	A ,	A    =				parser.parse_args_into_dataclasses()

      if (
          os.path.exists(training_args.output_dir     )
          and os.listdir(training_args.output_dir     )
          and training_args.do_train
          and not training_args.overwrite_output_dir
      ):
            raise ValueError(
                f'''Output directory ({training_args.output_dir}) already exists and is not empty. Use'''
                ' --overwrite_output_dir to overcome.'     )

      # Setup logging
      logging.basicConfig(
          format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',    datefmt='%m/%d/%Y %H:%M:%S',    level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,    )
      logger.warning(
          'Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s',    training_args.local_rank,    training_args.device,    training_args.n_gpu,    bool(training_args.local_rank != -1     ),    training_args.fpaa,    )
      # Set the verbosity to info of the Transformers logger (on main process only):
      if is_main_process(training_args.local_rank     ):
            transformers.utils.logging.set_verbosity_info()
            transformers.utils.logging.enable_default_handler()
            transformers.utils.logging.enable_explicit_format()
      logger.info('Training/evaluation parameters %s',    lowerCAmelCase     )

      # Prepare Question-Answering task
      # Load pretrained model and tokenizer
      #
      # Distributed training:
      # The .from_pretrained methods guarantee that only one local process can concurrently
      # download model & vocab.

      A    =				AutoConfig.from_pretrained(
          model_args.config_name if model_args.config_name else model_args.model_name_or_path,    cache_dir=model_args.cache_dir,    )
      A    =				AutoTokenizer.from_pretrained(
          model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,    cache_dir=model_args.cache_dir,    use_fast=lowerCAmelCase,    )
      A    =				AutoModelForQuestionAnswering.from_pretrained(
          model_args.model_name_or_path,    from_tf=bool('.ckpt' in model_args.model_name_or_path     ),    config=lowerCAmelCase,    cache_dir=model_args.cache_dir,    )

      # Get datasets
      A    =				hasattr(model.config,    'lang2id'     )
      A    =				(
          SquadDataset(
              lowerCAmelCase,    tokenizer=lowerCAmelCase,    is_language_sensitive=lowerCAmelCase,    cache_dir=model_args.cache_dir     )
          if training_args.do_train
          else None
      )
      A    =				(
          SquadDataset(
              lowerCAmelCase,    tokenizer=lowerCAmelCase,    mode='dev',    is_language_sensitive=lowerCAmelCase,    cache_dir=model_args.cache_dir,    )
          if training_args.do_eval
          else None
      )

      # Data collator
      A    =				DataCollatorWithPadding(lowerCAmelCase,    pad_to_multiple_of=8     ) if training_args.fpaa else None

      # Initialize our Trainer
      A    =				Trainer(
          model=lowerCAmelCase,    args=lowerCAmelCase,    train_dataset=lowerCAmelCase,    eval_dataset=lowerCAmelCase,    data_collator=lowerCAmelCase,    )

      # Training
      if training_args.do_train:
            trainer.train(
                model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path     ) else None     )
            trainer.save_model()
            # For convenience, we also re-save the tokenizer to the same directory,
            # so that you can share your model easily on huggingface.co/models =)
            if trainer.is_world_master():
                  tokenizer.save_pretrained(training_args.output_dir     )


def 			__UpperCamelCase    (lowerCAmelCase					:  Tuple     )     ->  List[Any]:
      # For xla_spawn (TPUs)
      main()


if __name__ == "__main__":
  main()



