
'''simple docstring'''





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


a__		: Union[str, Any]  							=							logging.getLogger(__name__)

require_version("pytorch_lightning>=1.0.4")

a__		: List[str]  							=							{
    "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
a__		: List[str]  							=							{
    "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
}
a__		: Optional[Any]  							=							sorted(arg_to_scheduler.keys())
a__		: Tuple  							=							"{" + ", ".join(arg_to_scheduler_choices) + "}"


class 			UpperCamelCase__	(     pl.LightningModule):






 def __init__(      self	:Optional[Any]						,					_A	:argparse.Namespace						,					_A	:Optional[int]=None						,					_A	:List[Any]="base"						,					_A	:Tuple=None						,					_A	:List[Any]=None						,					_A	:Optional[Any]=None						,					**_A	:Optional[int]						,					)					->      Optional[Any]:







       '''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 						lowercase_    (      self	:str						,					*_A	:str						,					**_A	:Dict      )					->      Union[str, Any]:







       '''simple docstring'''



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






 def 						lowercase_    (      self	:Tuple      )					->      Optional[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 						lowercase_    (      self	:List[str]      )					->      Tuple:







       '''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 						lowercase_    (      self	:Tuple						,					_A	:Tuple						,					_A	:str      )					->      Union[str, Any]:







       '''simple docstring'''



       return self.validation_step(_A						,					_A      )






 def 						lowercase_    (      self	:Union[str, Any]						,					_A	:Optional[int]      )					->      List[Any]:







       '''simple docstring'''



       return self.validation_end(_A      )






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







       '''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 						lowercase_    (      self	:int						,					_A	:Optional[int]      )					->      Optional[int]:







       '''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 						lowercase_    (      self	:int						,					_A	:str						,					_A	:int						,					_A	:bool = False      )					->      Union[str, Any]:







       '''simple docstring'''



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






 def 						lowercase_    (      self	:Optional[Any]      )					->      str:







       '''simple docstring'''



       return self.train_loader






 def 						lowercase_    (      self	:Optional[int]      )					->      Union[str, Any]:







       '''simple docstring'''



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






 def 						lowercase_    (      self	:List[str]      )					->      Union[str, Any]:







       '''simple docstring'''



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






 def 						lowercase_    (      self	:str						,					_A	:Union[str, Any]      )					->      List[Any]:







       '''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 						lowercase_    (      self	:List[str]						,					_A	:Dict[str, Any]      )					->      None:







       '''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 						lowercase_    (      _A	:Dict						,					_A	:Union[str, Any]      )					->      Optional[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 			UpperCamelCase__	(     pl.Callback):






 def 						lowercase_    (      self	:Tuple						,					_A	:Tuple						,					_A	:Dict      )					->      int:







       '''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 						lowercase_    (      self	:Dict						,					_A	:pl.Trainer						,					_A	:pl.LightningModule      )					->      str:







       '''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 						lowercase_    (      self	:Union[str, Any]						,					_A	:pl.Trainer						,					_A	:pl.LightningModule      )					->      str:







       '''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 							snake_case				(		UpperCAmelCase   ,	UpperCAmelCase						)->      None:







      """simple docstring"""
      #  To allow all pl args uncomment the following line
      #  parser = pl.Trainer.add_argparse_args(parser)
      parser.add_argument(
          '--output_dir'   ,	default=UpperCAmelCase   ,	type=UpperCAmelCase   ,	required=UpperCAmelCase   ,	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=UpperCAmelCase   ,	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=UpperCAmelCase						)
      parser.add_argument('--max_grad_norm'   ,	dest='gradient_clip_val'   ,	default=1.0   ,	type=UpperCAmelCase   ,	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=UpperCAmelCase   ,	default=1   ,	help='Number of updates steps to accumulate before performing a backward/update pass.'   ,	)
      parser.add_argument('--seed'   ,	type=UpperCAmelCase   ,	default=4_2   ,	help='random seed for initialization'						)
      parser.add_argument(
          '--data_dir'   ,	default=UpperCAmelCase   ,	type=UpperCAmelCase   ,	required=UpperCAmelCase   ,	help='The input data dir. Should contain the training files for the CoNLL-2003 NER task.'   ,	)





def 							snake_case				(		UpperCAmelCase   ,	UpperCAmelCase   ,	UpperCAmelCase=None   ,	UpperCAmelCase=True   ,	UpperCAmelCase=[]   ,	UpperCAmelCase=None   ,	UpperCAmelCase=None   ,	**UpperCAmelCase   ,	)->      Optional[Any]:







      """simple docstring"""
      pl.seed_everything(args.seed						)

      # init model
      __A													=      Path(model.hparams.output_dir						)
      odir.mkdir(exist_ok=UpperCAmelCase						)

      # 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(UpperCAmelCase						)
      if logging_callback is None:
            __A													=      LoggingCallback()

      __A													=      {}

      # TODO: remove with PyTorch 1.6 since pl uses native amp
      if args.fpaa:
            __A													=      1_6
            __A													=      args.fpaa_opt_level

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

      __A													=      args.accumulate_grad_batches
      __A													=      extra_train_kwargs.get('accelerator'   ,	UpperCAmelCase						)
      __A													=      extra_train_kwargs.get('profiler'   ,	UpperCAmelCase						)

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

      if args.do_train:
            trainer.fit(UpperCAmelCase						)

      return trainer













'''simple docstring'''


from __future__ import annotations




def       A							(      _UpperCAmelCase   :				list					)  ->				int | float:


  '''simple docstring'''




  __lowerCAmelCase		:    List[Any]			     =  sorted(_UpperCAmelCase					)
  __lowerCAmelCase		:    List[str]			     =  len(_UpperCAmelCase					)
  __lowerCAmelCase		:    Any			     =  length >> 1
  return (
      (sorted_list[mid_index] + sorted_list[mid_index - 1]) / 2
      if length % 2 == 0
      else sorted_list[mid_index]
  )




def       A							(      )  ->				Union[str, Any]:


  '''simple docstring'''




  import doctest

  doctest.testmod()


if __name__ == "__main__":
    main()





