

import argparse
import logging
import os
from pathlib import Path
from typing import Any, Dict

import pytorch_lightning as pl
from pytorch_lightning.utilities import rank_zero_info

from transformers import (
    AdamW,
    AutoConfig,
    AutoModel,
    AutoModelForPreTraining,
    AutoModelForQuestionAnswering,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
    AutoModelWithLMHead,
    AutoTokenizer,
    PretrainedConfig,
    PreTrainedTokenizer,
)
from transformers.optimization import (
    Adafactor,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.utils.versions import require_version


lowerCamelCase									= logging.getLogger(__name__)

require_version('''pytorch_lightning>=1.0.4''')

lowerCamelCase									= {
    '''base''': AutoModel,
    '''sequence-classification''': AutoModelForSequenceClassification,
    '''question-answering''': AutoModelForQuestionAnswering,
    '''pretraining''': AutoModelForPreTraining,
    '''token-classification''': AutoModelForTokenClassification,
    '''language-modeling''': AutoModelWithLMHead,
    '''summarization''': AutoModelForSeqaSeqLM,
    '''translation''': AutoModelForSeqaSeqLM,
}


# update this and the import above to support new schedulers from transformers.optimization
lowerCamelCase									= {
    '''linear''': get_linear_schedule_with_warmup,
    '''cosine''': get_cosine_schedule_with_warmup,
    '''cosine_w_restarts''': get_cosine_with_hard_restarts_schedule_with_warmup,
    '''polynomial''': get_polynomial_decay_schedule_with_warmup,
    # '': get_constant_schedule,             # not supported for now
    # '': get_constant_schedule_with_warmup, # not supported for now
}
lowerCamelCase									= sorted(arg_to_scheduler.keys())
lowerCamelCase									= '''{''' + ''', '''.join(arg_to_scheduler_choices) + '''}'''







class 						__magic_name__      (							pl.LightningModule							):



   def __init__(    self,   lowercase_,   lowercase_=None,   lowercase_="base",   lowercase_=None,   lowercase_=None,   lowercase_=None,   **lowercase_,   )							->   Dict:






     """simple docstring"""






     super().__init__()
     # TODO: move to self.save_hyperparameters()
     # self.save_hyperparameters()
     # can also expand arguments into trainer signature for easier reading

     self.save_hyperparameters(_A	)
     a__  =0
     a__  =Path(self.hparams.output_dir	)
     a__  =self.hparams.cache_dir if self.hparams.cache_dir else None
     if config is None:
       a__  =AutoConfig.from_pretrained(
           self.hparams.config_name if self.hparams.config_name else self.hparams.model_name_or_path,   **({'''num_labels''': num_labels} if num_labels is not None else {}),   cache_dir=_A,   **_A,   )
     else:
       a__  =config

     a__  =('encoder_layerdrop', 'decoder_layerdrop', 'dropout', 'attention_dropout')
     for p in extra_model_params:
       if getattr(self.hparams,   _A,   _A	):
         assert hasattr(self.config,   _A	), F"""model config doesn't have a `{p}` attribute"""
         setattr(self.config,   _A,   getattr(self.hparams,   _A	)	)

     if tokenizer is None:
       a__  =AutoTokenizer.from_pretrained(
           self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path,   cache_dir=_A,   )
     else:
       a__  =tokenizer
     a__  =MODEL_MODES[mode]
     if model is None:
       a__  =self.model_type.from_pretrained(
           self.hparams.model_name_or_path,   from_tf=bool('''.ckpt''' in self.hparams.model_name_or_path	),   config=self.config,   cache_dir=_A,   )
     else:
       a__  =model



   def 			_UpperCAmelCase			(    self,   *lowercase_,   **lowercase_	)							->   str:






     """simple docstring"""






     a__  =self.model_type.from_pretrained(*_A,   **_A	)



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






     """simple docstring"""






     a__  =arg_to_scheduler[self.hparams.lr_scheduler]
     a__  =get_schedule_func(
         self.opt,   num_warmup_steps=self.hparams.warmup_steps,   num_training_steps=self.total_steps()	)
     a__  ={'scheduler': scheduler, 'interval': 'step', 'frequency': 1}
     return scheduler



   def 			_UpperCAmelCase			(    self	)							->   Any:






     """simple docstring"""






     a__  =self.model
     a__  =['bias', 'LayerNorm.weight']
     a__  =[
         {
             'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay	)],
             'weight_decay': self.hparams.weight_decay,
         },
         {
             'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay	)],
             'weight_decay': 0.0,
         },
     ]
     if self.hparams.adafactor:
       a__  =Adafactor(
           _A,   lr=self.hparams.learning_rate,   scale_parameter=_A,   relative_step=_A	)

     else:
       a__  =AdamW(
           _A,   lr=self.hparams.learning_rate,   eps=self.hparams.adam_epsilon	)
     a__  =optimizer

     a__  =self.get_lr_scheduler()

     return [optimizer], [scheduler]



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






     """simple docstring"""






     return self.validation_step(_A,   _A	)



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






     """simple docstring"""






     return self.validation_end(_A	)



   def 			_UpperCAmelCase			(    self	)							->   Optional[Any]:






     """simple docstring"""






     a__  =max(1,   self.hparams.gpus	)  # TODO: consider num_tpu_cores
     a__  =self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
     return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs



   def 			_UpperCAmelCase			(    self,   lowercase_	)							->   Any:






     """simple docstring"""






     if mode == "test":
       a__  =len(self.test_dataloader().dataset	)
     else:
       a__  =self.get_dataloader('''train''',   self.hparams.train_batch_size,   shuffle=_A	)
       a__  =len(self.train_dataloader().dataset	)



   def 			_UpperCAmelCase			(    self,   lowercase_,   lowercase_,   lowercase_ = False	)							->   Union[str, Any]:






     """simple docstring"""






     raise NotImplementedError('''You must implement this for your task'''	)



   def 			_UpperCAmelCase			(    self	)							->   str:






     """simple docstring"""






     return self.train_loader



   def 			_UpperCAmelCase			(    self	)							->   List[str]:






     """simple docstring"""






     return self.get_dataloader('''dev''',   self.hparams.eval_batch_size,   shuffle=_A	)



   def 			_UpperCAmelCase			(    self	)							->   Tuple:






     """simple docstring"""






     return self.get_dataloader('''test''',   self.hparams.eval_batch_size,   shuffle=_A	)



   def 			_UpperCAmelCase			(    self,   lowercase_	)							->   int:






     """simple docstring"""






     return os.path.join(
         self.hparams.data_dir,   '''cached_{}_{}_{}'''.format(
             _A,   list(filter(_A,   self.hparams.model_name_or_path.split('''/'''	)	)	).pop(),   str(self.hparams.max_seq_length	),   ),   )



   @pl.utilities.rank_zero_only
   def 			_UpperCAmelCase			(    self,   lowercase_	)							->   Optional[int]:






     """simple docstring"""






     a__  =self.output_dir.joinpath('''best_tfmr'''	)
     a__  =self.step_count
     self.model.save_pretrained(_A	)
     self.tokenizer.save_pretrained(_A	)



   @staticmethod
   def 			_UpperCAmelCase			(    lowercase_,   lowercase_	)							->   int:






     """simple docstring"""






     parser.add_argument(
         '''--model_name_or_path''',   default=_A,   type=_A,   required=_A,   help='''Path to pretrained model or model identifier from huggingface.co/models''',   )
     parser.add_argument(
         '''--config_name''',   default='''''',   type=_A,   help='''Pretrained config name or path if not the same as model_name'''	)
     parser.add_argument(
         '''--tokenizer_name''',   default=_A,   type=_A,   help='''Pretrained tokenizer name or path if not the same as model_name''',   )
     parser.add_argument(
         '''--cache_dir''',   default='''''',   type=_A,   help='''Where do you want to store the pre-trained models downloaded from huggingface.co''',   )
     parser.add_argument(
         '''--encoder_layerdrop''',   type=_A,   help='''Encoder layer dropout probability (Optional). Goes into model.config''',   )
     parser.add_argument(
         '''--decoder_layerdrop''',   type=_A,   help='''Decoder layer dropout probability (Optional). Goes into model.config''',   )
     parser.add_argument(
         '''--dropout''',   type=_A,   help='''Dropout probability (Optional). Goes into model.config''',   )
     parser.add_argument(
         '''--attention_dropout''',   type=_A,   help='''Attention dropout probability (Optional). Goes into model.config''',   )
     parser.add_argument('''--learning_rate''',   default=5E-5,   type=_A,   help='''The initial learning rate for Adam.'''	)
     parser.add_argument(
         '''--lr_scheduler''',   default='''linear''',   choices=_A,   metavar=_A,   type=_A,   help='''Learning rate scheduler''',   )
     parser.add_argument('''--weight_decay''',   default=0.0,   type=_A,   help='''Weight decay if we apply some.'''	)
     parser.add_argument('''--adam_epsilon''',   default=1E-8,   type=_A,   help='''Epsilon for Adam optimizer.'''	)
     parser.add_argument('''--warmup_steps''',   default=0,   type=_A,   help='''Linear warmup over warmup_steps.'''	)
     parser.add_argument('''--num_workers''',   default=4,   type=_A,   help='''kwarg passed to DataLoader'''	)
     parser.add_argument('''--num_train_epochs''',   dest='''max_epochs''',   default=3,   type=_A	)
     parser.add_argument('''--train_batch_size''',   default=32,   type=_A	)
     parser.add_argument('''--eval_batch_size''',   default=32,   type=_A	)
     parser.add_argument('''--adafactor''',   action='''store_true'''	)







class 						__magic_name__      (							pl.Callback							):



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






     """simple docstring"""






     a__  =trainer.lr_schedulers[0]['scheduler']
     a__  ={F"""lr_group_{i}""": lr for i, lr in enumerate(lr_scheduler.get_lr()	)}
     pl_module.logger.log_metrics(_A	)



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






     """simple docstring"""






     rank_zero_info('''***** Validation results *****'''	)
     a__  =trainer.callback_metrics
     # Log results
     for key in sorted(_A	):
       if key not in ["log", "progress_bar"]:
         rank_zero_info('''{} = {}\n'''.format(_A,   str(metrics[key]	)	)	)



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






     """simple docstring"""






     rank_zero_info('''***** Test results *****'''	)
     a__  =trainer.callback_metrics
     # Log and save results to file
     a__  =os.path.join(pl_module.hparams.output_dir,   '''test_results.txt'''	)
     with open(_A,   '''w'''	) as writer:
       for key in sorted(_A	):
         if key not in ["log", "progress_bar"]:
           rank_zero_info('''{} = {}\n'''.format(_A,   str(metrics[key]	)	)	)
           writer.write('''{} = {}\n'''.format(_A,   str(metrics[key]	)	)	)





def       UpperCAmelCase__    (      _A				:			Tuple     ,				_A				:			Dict							):





  '''simple docstring'''







  parser.add_argument(
      '''--output_dir'''     ,				default=_A     ,				type=_A     ,				required=_A     ,				help='''The output directory where the model predictions and checkpoints will be written.'''     ,				)
  parser.add_argument(
      '''--fp16'''     ,				action='''store_true'''     ,				help='''Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit'''     ,				)

  parser.add_argument(
      '''--fp16_opt_level'''     ,				type=_A     ,				default='''O2'''     ,				help=(
          '''For fp16: Apex AMP optimization level selected in [\'O0\', \'O1\', \'O2\', and \'O3\'].'''
          '''See details at https://nvidia.github.io/apex/amp.html'''
      )     ,				)
  parser.add_argument('''--n_tpu_cores'''     ,				dest='''tpu_cores'''     ,				type=_A							)
  parser.add_argument('''--max_grad_norm'''     ,				dest='''gradient_clip_val'''     ,				default=1.0     ,				type=_A     ,				help='''Max gradient norm'''							)
  parser.add_argument('''--do_train'''     ,				action='''store_true'''     ,				help='''Whether to run training.'''							)
  parser.add_argument('''--do_predict'''     ,				action='''store_true'''     ,				help='''Whether to run predictions on the test set.'''							)
  parser.add_argument(
      '''--gradient_accumulation_steps'''     ,				dest='''accumulate_grad_batches'''     ,				type=_A     ,				default=1     ,				help='''Number of updates steps to accumulate before performing a backward/update pass.'''     ,				)
  parser.add_argument('''--seed'''     ,				type=_A     ,				default=42     ,				help='''random seed for initialization'''							)
  parser.add_argument(
      '''--data_dir'''     ,				default=_A     ,				type=_A     ,				required=_A     ,				help='''The input data dir. Should contain the training files for the CoNLL-2003 NER task.'''     ,				)




def       UpperCAmelCase__    (      _A				:			Union[str, Any]     ,				_A				:			int     ,				_A				:			List[Any]=None     ,				_A				:			str=True     ,				_A				:			Any=[]     ,				_A				:			Union[str, Any]=None     ,				_A				:			List[str]=None     ,				**_A				:			List[Any]     ,				):





  '''simple docstring'''







  pl.seed_everything(args.seed							)

  # init model
  a__  =Path(model.hparams.output_dir							)
  odir.mkdir(exist_ok=_A							)

  # add custom checkpoints
  if checkpoint_callback is None:
    a__  =pl.callbacks.ModelCheckpoint(
        filepath=args.output_dir     ,				prefix='''checkpoint'''     ,				monitor='''val_loss'''     ,				mode='''min'''     ,				save_top_k=1							)
  if early_stopping_callback:
    extra_callbacks.append(_A							)
  if logging_callback is None:
    a__  =LoggingCallback()

  a__  ={}

  # TODO: remove with PyTorch 1.6 since pl uses native amp
  if args.fpaa:
    a__  =16
    a__  =args.fpaa_opt_level

  if args.gpus > 1:
    a__  ='ddp'

  a__  =args.accumulate_grad_batches
  a__  =extra_train_kwargs.get('''accelerator'''     ,				_A							)
  a__  =extra_train_kwargs.get('''profiler'''     ,				_A							)

  a__  =pl.Trainer.from_argparse_args(
      _A     ,				weights_summary=_A     ,				callbacks=[logging_callback] + extra_callbacks     ,				logger=_A     ,				checkpoint_callback=_A     ,				**_A     ,				)

  if args.do_train:
    trainer.fit(_A							)

  return trainer



import os
import random
from pathlib import Path
from typing import List

import numpy as np
import torch
from torch.cuda.amp import GradScaler

from .utils import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    get_pretty_name,
    is_tpu_available,
    is_xpu_available,
    save,
)


if is_tpu_available(check_device=False):
       import torch_xla.core.xla_model as xm

from .logging import get_logger
from .state import PartialState


lowerCamelCase									= get_logger(__name__)




def       UpperCAmelCase__    (      _A				:			str     ,				_A				:			List[dict]     ,				_A				:			list     ,				_A				:			list     ,				_A				:			int     ,				_A				:			GradScaler = None     ,				):





  '''simple docstring'''







  for i, state in enumerate(_A							):
    a__  =F"""{MODEL_NAME}.bin""" if i == 0 else F"""{MODEL_NAME}_{i}.bin"""
    a__  =os.path.join(_A     ,				_A							)
    save(_A     ,				_A							)
    logger.info(F"""Model weights saved in {output_model_file}"""							)
  # Optimizer states
  for i, opt in enumerate(_A							):
    a__  =opt.state_dict()
    a__  =F"""{OPTIMIZER_NAME}.bin""" if i == 0 else F"""{OPTIMIZER_NAME}_{i}.bin"""
    a__  =os.path.join(_A     ,				_A							)
    save(_A     ,				_A							)
    logger.info(F"""Optimizer state saved in {output_optimizer_file}"""							)
  # Scheduler states
  for i, scheduler in enumerate(_A							):
    a__  =scheduler.state_dict()
    a__  =F"""{SCHEDULER_NAME}.bin""" if i == 0 else F"""{SCHEDULER_NAME}_{i}.bin"""
    a__  =os.path.join(_A     ,				_A							)
    save(_A     ,				_A							)
    logger.info(F"""Scheduler state saved in {output_scheduler_file}"""							)
  # GradScaler state
  if scaler is not None:
    a__  =scaler.state_dict()
    a__  =os.path.join(_A     ,				_A							)
    torch.save(_A     ,				_A							)
    logger.info(F"""Gradient scaler state saved in {output_scaler_file}"""							)
  # Random number generator states
  a__  ={}
  a__  =F"""{RNG_STATE_NAME}_{process_index}.pkl"""
  a__  =random.getstate()
  a__  =np.random.get_state()
  a__  =torch.get_rng_state()
  if is_xpu_available():
    a__  =torch.xpu.get_rng_state_all()
  else:
    a__  =torch.cuda.get_rng_state_all()
  if is_tpu_available():
    a__  =xm.get_rng_state()
  a__  =os.path.join(_A     ,				_A							)
  torch.save(_A     ,				_A							)
  logger.info(F"""Random states saved in {output_states_file}"""							)
  return output_dir




def       UpperCAmelCase__    (      _A				:			Tuple     ,				_A				:			List[Any]     ,				_A				:			int     ,				_A				:			Tuple     ,				_A				:			List[Any]     ,				_A				:			List[str]=None     ,				_A				:			Dict=None     ,				**_A				:			Tuple     ,				):





  '''simple docstring'''







  if map_location not in [None, "cpu", "on_device"]:
    raise TypeError(
        '''Unsupported optimizer map location passed, please choose one of `None`, `\'cpu\'`, or `\'on_device\'`'''							)
  if map_location is None:
    a__  ='''cpu'''
  elif map_location == "on_device":
    a__  =PartialState().device
  # Model states
  for i, model in enumerate(_A							):
    a__  =F"""{MODEL_NAME}.bin""" if i == 0 else F"""{MODEL_NAME}_{i}.bin"""
    a__  =os.path.join(_A     ,				_A							)
    models[i].load_state_dict(torch.load(_A     ,				map_location=_A							)     ,				**_A							)
  logger.info('''All model weights loaded successfully'''							)

  # Optimizer states
  for i, opt in enumerate(_A							):
    a__  =F"""{OPTIMIZER_NAME}.bin""" if i == 0 else F"""{OPTIMIZER_NAME}_{i}.bin"""
    a__  =os.path.join(_A     ,				_A							)
    a__  =torch.load(_A     ,				map_location=_A							)
    optimizers[i].load_state_dict(_A							)
  logger.info('''All optimizer states loaded successfully'''							)

  # Scheduler states
  for i, scheduler in enumerate(_A							):
    a__  =F"""{SCHEDULER_NAME}.bin""" if i == 0 else F"""{SCHEDULER_NAME}_{i}.bin"""
    a__  =os.path.join(_A     ,				_A							)
    scheduler.load_state_dict(torch.load(_A							)							)
  logger.info('''All scheduler states loaded successfully'''							)

  # GradScaler state
  if scaler is not None:
    a__  =os.path.join(_A     ,				_A							)
    scaler.load_state_dict(torch.load(_A							)							)
    logger.info('''GradScaler state loaded successfully'''							)

  # Random states
  try:
    a__  =torch.load(os.path.join(_A     ,				F"""{RNG_STATE_NAME}_{process_index}.pkl"""							)							)
    random.setstate(states['''random_state''']							)
    np.random.set_state(states['''numpy_random_seed''']							)
    torch.set_rng_state(states['''torch_manual_seed''']							)
    if is_xpu_available():
      torch.xpu.set_rng_state_all(states['''torch_xpu_manual_seed''']							)
    else:
      torch.cuda.set_rng_state_all(states['''torch_cuda_manual_seed''']							)
    if is_tpu_available():
      xm.set_rng_state(states['''xm_seed''']							)
    logger.info('''All random states loaded successfully'''							)
  except Exception:
    logger.info('''Could not load random states'''							)




def       UpperCAmelCase__    (      _A				:			Tuple     ,				_A				:			Tuple     ,				_A				:			int = 0							):





  '''simple docstring'''







  a__  =Path(_A							) / F"""custom_checkpoint_{index}.pkl"""
  logger.info(F"""Saving the state of {get_pretty_name(_A							)} to {save_location}"""							)
  torch.save(obj.state_dict()     ,				_A							)




def       UpperCAmelCase__    (      _A				:			int     ,				_A				:			Tuple     ,				_A				:			int = 0							):





  '''simple docstring'''







  a__  =F"""{path}/custom_checkpoint_{index}.pkl"""
  logger.info(F"""Loading the state of {get_pretty_name(_A							)} from {load_location}"""							)
  obj.load_state_dict(torch.load(_A     ,				map_location='''cpu'''							)							)
