





'''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__       :     Optional[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__       :     Union[str, Any]          ={
    '''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__       :     Dict          =sorted(arg_to_scheduler.keys())
A__       :     str          ='''{''' + ''', '''.join(arg_to_scheduler_choices) + '''}'''




class 				UpperCAmelCase  (    pl.LightningModule     ):




      def __init__(    self    :				Union[str, Any]    ,    __snake_case    :				Any    ,    __snake_case    :				List[Any]=None    ,    __snake_case    :				int="base"    ,    __snake_case    :				Optional[int]=None    ,    __snake_case    :				Tuple=None    ,    __snake_case    :				Union[str, Any]=None    ,    **__snake_case    :				Optional[Any]    ,    )       ->   Any:
            super().__init__()
            # TODO: move to self.save_hyperparameters()
            # self.save_hyperparameters()
            # can also expand arguments into trainer signature for easier reading

            self.save_hyperparameters(lowercase_						)
            _lowerCAmelCase							  =							0
            _lowerCAmelCase							  =							Path(self.hparams.output_dir						)
            _lowerCAmelCase							  =							self.hparams.cache_dir if self.hparams.cache_dir else None
            if config is None:
                  _lowerCAmelCase							  =							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=lowercase_    ,    **lowercase_    ,    )
            else:
                  _lowerCAmelCase							  =							config

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

            if tokenizer is None:
                  _lowerCAmelCase							  =							AutoTokenizer.from_pretrained(
                      self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path    ,    cache_dir=lowercase_    ,    )
            else:
                  _lowerCAmelCase							  =							tokenizer
            _lowerCAmelCase							  =							MODEL_MODES[mode]
            if model is None:
                  _lowerCAmelCase							  =							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=lowercase_    ,    )
            else:
                  _lowerCAmelCase							  =							model




      def        lowercase__					(    self    :				List[Any]    ,    *__snake_case    :				Any    ,    **__snake_case    :				Optional[Any]						)       ->   List[Any]:
            _lowerCAmelCase							  =							self.model_type.from_pretrained(*lowercase_    ,    **lowercase_						)




      def        lowercase__					(    self    :				int						)       ->   Dict:
            _lowerCAmelCase							  =							arg_to_scheduler[self.hparams.lr_scheduler]
            _lowerCAmelCase							  =							get_schedule_func(
                self.opt    ,    num_warmup_steps=self.hparams.warmup_steps    ,    num_training_steps=self.total_steps()						)
            _lowerCAmelCase							  =							{'''scheduler''': scheduler, '''interval''': '''step''', '''frequency''': 1}
            return scheduler




      def        lowercase__					(    self    :				str						)       ->   Optional[int]:
            _lowerCAmelCase							  =							self.model
            _lowerCAmelCase							  =							['''bias''', '''LayerNorm.weight''']
            _lowerCAmelCase							  =							[
                {
                    '''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:
                  _lowerCAmelCase							  =							Adafactor(
                      lowercase_    ,    lr=self.hparams.learning_rate    ,    scale_parameter=lowercase_    ,    relative_step=lowercase_						)

            else:
                  _lowerCAmelCase							  =							AdamW(
                      lowercase_    ,    lr=self.hparams.learning_rate    ,    eps=self.hparams.adam_epsilon						)
            _lowerCAmelCase							  =							optimizer

            _lowerCAmelCase							  =							self.get_lr_scheduler()

            return [optimizer], [scheduler]




      def        lowercase__					(    self    :				int    ,    __snake_case    :				Union[str, Any]    ,    __snake_case    :				List[Any]						)       ->   List[Any]:
            return self.validation_step(lowercase_    ,    lowercase_						)




      def        lowercase__					(    self    :				Union[str, Any]    ,    __snake_case    :				str						)       ->   Optional[int]:
            return self.validation_end(lowercase_						)




      def        lowercase__					(    self    :				Tuple						)       ->   int:
            _lowerCAmelCase							  =							max(1    ,    self.hparams.gpus						)  # TODO: consider num_tpu_cores
            _lowerCAmelCase							  =							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    :				Optional[int]    ,    __snake_case    :				Tuple						)       ->   Any:

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




      def        lowercase__					(    self    :				Optional[int]    ,    __snake_case    :				Optional[Any]    ,    __snake_case    :				str    ,    __snake_case    :				Optional[int] = False						)       ->   Optional[int]:
            raise NotImplementedError("""You must implement this for your task"""						)




      def        lowercase__					(    self    :				Tuple						)       ->   Tuple:
            return self.train_loader




      def        lowercase__					(    self    :				List[Any]						)       ->   Dict:
            return self.get_dataloader("""dev"""    ,    self.hparams.eval_batch_size    ,    shuffle=lowercase_						)




      def        lowercase__					(    self    :				Dict						)       ->   Dict:
            return self.get_dataloader("""test"""    ,    self.hparams.eval_batch_size    ,    shuffle=lowercase_						)




      def        lowercase__					(    self    :				Tuple    ,    __snake_case    :				Any						)       ->   Optional[int]:
            return os.path.join(
                self.hparams.data_dir    ,    """cached_{}_{}_{}""".format(
                    lowercase_    ,    list(filter(lowercase_    ,    self.hparams.model_name_or_path.split("""/"""						)						)						).pop()    ,    str(self.hparams.max_seq_length						)    ,    )    ,    )




      @pl.utilities.rank_zero_only
      def        lowercase__					(    self    :				int    ,    __snake_case    :				int						)       ->   None:
            _lowerCAmelCase							  =							self.output_dir.joinpath("""best_tfmr"""						)
            _lowerCAmelCase							  =							self.step_count
            self.model.save_pretrained(lowercase_						)
            self.tokenizer.save_pretrained(lowercase_						)




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




class 				UpperCAmelCase  (    pl.Callback     ):




      def        lowercase__					(    self    :				Union[str, Any]    ,    __snake_case    :				Any    ,    __snake_case    :				List[str]						)       ->   Union[str, Any]:
            _lowerCAmelCase							  =							trainer.lr_schedulers[0]['''scheduler''']
            _lowerCAmelCase							  =							{f"lr_group_{i}": lr for i, lr in enumerate(lr_scheduler.get_lr()						)}
            pl_module.logger.log_metrics(lowercase_						)




      def        lowercase__					(    self    :				Any    ,    __snake_case    :				List[str]    ,    __snake_case    :				Any						)       ->   Any:
            rank_zero_info("""***** Validation results *****"""						)
            _lowerCAmelCase							  =							trainer.callback_metrics
            # Log results
            for key in sorted(lowercase_						):
                  if key not in ["log", "progress_bar"]:
                        rank_zero_info("""{} = {}\n""".format(lowercase_    ,    str(metrics[key]						)						)						)




      def        lowercase__					(    self    :				Union[str, Any]    ,    __snake_case    :				Optional[int]    ,    __snake_case    :				Any						)       ->   str:
            rank_zero_info("""***** Test results *****"""						)
            _lowerCAmelCase							  =							trainer.callback_metrics
            # Log and save results to file
            _lowerCAmelCase							  =							os.path.join(pl_module.hparams.output_dir    ,    """test_results.txt"""						)
            with open(lowercase_    ,    """w"""						) as writer:
                  for key in sorted(lowercase_						):
                        if key not in ["log", "progress_bar"]:
                              rank_zero_info("""{} = {}\n""".format(lowercase_    ,    str(metrics[key]						)						)						)
                              writer.write("""{} = {}\n""".format(lowercase_    ,    str(metrics[key]						)						)						)




def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase						):





      """simple docstring"""


      parser.add_argument(
          """--output_dir"""							,			default=lowercase_							,			type=lowercase_							,			required=lowercase_							,			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=lowercase_							,			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=lowercase_						)
      parser.add_argument("""--max_grad_norm"""							,			dest="""gradient_clip_val"""							,			default=1.0							,			type=lowercase_							,			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=lowercase_							,			default=1							,			help="""Number of updates steps to accumulate before performing a backward/update pass."""							,			)
      parser.add_argument("""--seed"""							,			type=lowercase_							,			default=42							,			help="""random seed for initialization"""						)
      parser.add_argument(
          """--data_dir"""							,			default=lowercase_							,			type=lowercase_							,			required=lowercase_							,			help="""The input data dir. Should contain the training files for the CoNLL-2003 NER task."""							,			)





def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase=None							,			lowerCAmelCase=True							,			lowerCAmelCase=[]							,			lowerCAmelCase=None							,			lowerCAmelCase=None							,			**lowerCAmelCase							,			):





      """simple docstring"""


      pl.seed_everything(args.seed						)

      # init model
      _lowerCAmelCase							  =							Path(model.hparams.output_dir						)
      odir.mkdir(exist_ok=lowercase_						)

      # add custom checkpoints
      if checkpoint_callback is None:
            _lowerCAmelCase							  =							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(lowercase_						)
      if logging_callback is None:
            _lowerCAmelCase							  =							LoggingCallback()

      _lowerCAmelCase							  =							{}

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

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

      _lowerCAmelCase							  =							args.accumulate_grad_batches
      _lowerCAmelCase							  =							extra_train_kwargs.get("""accelerator"""							,			lowercase_						)
      _lowerCAmelCase							  =							extra_train_kwargs.get("""profiler"""							,			lowercase_						)

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

      if args.do_train:
            trainer.fit(lowercase_						)

      return trainer






'''simple docstring'''






import torch

from transformers import PreTrainedModel

from .custom_configuration import CustomConfig, NoSuperInitConfig




class 				UpperCAmelCase  (    snake_case_     ):
      _lowercase:    str		  =     CustomConfig




      def __init__(    self    :				Optional[int]    ,    __snake_case    :				Tuple						)       ->   Optional[Any]:
            super().__init__(__snake_case						)
            _lowerCAmelCase							  =							torch.nn.Linear(config.hidden_size    ,    config.hidden_size						)




      def        lowercase__					(    self    :				Dict    ,    __snake_case    :				List[Any]						)       ->   Dict:
            return self.linear(__snake_case						)




      def        lowercase__					(    self    :				List[Any]    ,    __snake_case    :				int						)       ->   Any:
            pass




class 				UpperCAmelCase  (    snake_case_     ):
      _lowercase:    List[str]		  =     NoSuperInitConfig




      def __init__(    self    :				List[str]    ,    __snake_case    :				List[str]						)       ->   int:
            super().__init__(__snake_case						)
            _lowerCAmelCase							  =							torch.nn.Linear(config.attribute    ,    config.attribute						)




      def        lowercase__					(    self    :				List[str]    ,    __snake_case    :				Dict						)       ->   Any:
            return self.linear(__snake_case						)




      def        lowercase__					(    self    :				Optional[int]    ,    __snake_case    :				List[Any]						)       ->   Dict:
            pass






