






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


_lowerCAmelCase     :     List[Any]            =	logging.getLogger(__name__)

require_version("""pytorch_lightning>=1.0.4""")

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




class        lowerCAmelCase__       (  pl.LightningModule					):
   def __init__(						self	:     int				,       snake_case__	:     argparse.Namespace				,       snake_case__	:     Optional[Any]=None				,       snake_case__	:     List[Any]="base"				,       snake_case__	:     Tuple=None				,       snake_case__	:     int=None				,       snake_case__	:     Tuple=None				,       **snake_case__	:     int				,       ):


        '''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(snake_case__				)
        UpperCAmelCase__		:  int				=				0
        UpperCAmelCase__		:  Union[str, Any]				=				Path(self.hparams.output_dir				)
        UpperCAmelCase__		:  List[str]				=				self.hparams.cache_dir if self.hparams.cache_dir else None
        if config is None:
             UpperCAmelCase__		:  Dict				=				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=snake_case__				,       **snake_case__				,       )
        else:
             UpperCAmelCase__		:  PretrainedConfig				=				config

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

        if tokenizer is None:
             UpperCAmelCase__		:  Union[str, Any]				=				AutoTokenizer.from_pretrained(
                 self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path				,       cache_dir=snake_case__				,       )
        else:
             UpperCAmelCase__		:  PreTrainedTokenizer				=				tokenizer
        UpperCAmelCase__		:  str				=				MODEL_MODES[mode]
        if model is None:
             UpperCAmelCase__		:  List[str]				=				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=snake_case__				,       )
        else:
             UpperCAmelCase__		:  Tuple				=				model
   def 	__a				(						self	:     Dict				,       *snake_case__	:     List[Any]				,       **snake_case__	:     int				):


        '''simple docstring'''
        UpperCAmelCase__		:  int				=				self.model_type.from_pretrained(*snake_case__				,       **snake_case__				)
   def 	__a				(						self	:     Union[str, Any]				):


        '''simple docstring'''
        UpperCAmelCase__		:  Dict				=				arg_to_scheduler[self.hparams.lr_scheduler]
        UpperCAmelCase__		:  Any				=				get_schedule_func(
            self.opt				,       num_warmup_steps=self.hparams.warmup_steps				,       num_training_steps=self.total_steps()				)
        UpperCAmelCase__		:  Optional[Any]				=				{"scheduler": scheduler, "interval": "step", "frequency": 1}
        return scheduler
   def 	__a				(						self	:     List[Any]				):


        '''simple docstring'''
        UpperCAmelCase__		:  Union[str, Any]				=				self.model
        UpperCAmelCase__		:  Optional[int]				=				["bias", "LayerNorm.weight"]
        UpperCAmelCase__		:  Optional[Any]				=				[
            {
                "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:
             UpperCAmelCase__		:  Tuple				=				Adafactor(
                 snake_case__				,       lr=self.hparams.learning_rate				,       scale_parameter=snake_case__				,       relative_step=snake_case__				)

        else:
             UpperCAmelCase__		:  List[str]				=				AdamW(
                 snake_case__				,       lr=self.hparams.learning_rate				,       eps=self.hparams.adam_epsilon				)
        UpperCAmelCase__		:  List[str]				=				optimizer

        UpperCAmelCase__		:  Optional[Any]				=				self.get_lr_scheduler()

        return [optimizer], [scheduler]
   def 	__a				(						self	:     Any				,       snake_case__	:     Optional[int]				,       snake_case__	:     Any				):


        '''simple docstring'''
        return self.validation_step(snake_case__				,       snake_case__				)
   def 	__a				(						self	:     Union[str, Any]				,       snake_case__	:     Union[str, Any]				):


        '''simple docstring'''
        return self.validation_end(snake_case__				)
   def 	__a				(						self	:     Dict				):


        '''simple docstring'''
        UpperCAmelCase__		:  str				=				max(1				,       self.hparams.gpus				)  # TODO: consider num_tpu_cores
        UpperCAmelCase__		:  Dict				=				self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
        return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs
   def 	__a				(						self	:     List[str]				,       snake_case__	:     Tuple				):


        '''simple docstring'''
        if stage == "test":
             UpperCAmelCase__		:  Optional[Any]				=				len(self.test_dataloader().dataset				)
        else:
             UpperCAmelCase__		:  Optional[int]				=				self.get_dataloader("train"				,       self.hparams.train_batch_size				,       shuffle=snake_case__				)
             UpperCAmelCase__		:  List[str]				=				len(self.train_dataloader().dataset				)
   def 	__a				(						self	:     List[str]				,       snake_case__	:     str				,       snake_case__	:     int				,       snake_case__	:     bool = False				):


        '''simple docstring'''
        raise NotImplementedError("You must implement this for your task"				)
   def 	__a				(						self	:     Any				):


        '''simple docstring'''
        return self.train_loader
   def 	__a				(						self	:     Optional[int]				):


        '''simple docstring'''
        return self.get_dataloader("dev"				,       self.hparams.eval_batch_size				,       shuffle=snake_case__				)
   def 	__a				(						self	:     Optional[Any]				):


        '''simple docstring'''
        return self.get_dataloader("test"				,       self.hparams.eval_batch_size				,       shuffle=snake_case__				)
   def 	__a				(						self	:     str				,       snake_case__	:     Union[str, Any]				):


        '''simple docstring'''
        return os.path.join(
            self.hparams.data_dir				,       "cached_{}_{}_{}".format(
                snake_case__				,       list(filter(snake_case__				,       self.hparams.model_name_or_path.split("/"				)				)				).pop()				,       str(self.hparams.max_seq_length				)				,       )				,       )
   @pl.utilities.rank_zero_only
   def 	__a				(						self	:     Dict				,       snake_case__	:     Dict[str, Any]				):


        '''simple docstring'''
        UpperCAmelCase__		:  str				=				self.output_dir.joinpath("best_tfmr"				)
        UpperCAmelCase__		:  Any				=				self.step_count
        self.model.save_pretrained(snake_case__				)
        self.tokenizer.save_pretrained(snake_case__				)




   @staticmethod
   def 	__a				(						snake_case__	:     Union[str, Any]				,       snake_case__	:     Union[str, Any]				):


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




class        lowerCAmelCase__       (  pl.Callback					):
   def 	__a				(						self	:     Tuple				,       snake_case__	:     str				,       snake_case__	:     Optional[Any]				):


        '''simple docstring'''
        if (
            trainer.is_global_zero and trainer.global_rank == 0
        ):  # we initialize the retriever only on master worker with RAY. In new pytorch-lightning accelorators are removed.
             pl_module.model.rag.retriever.init_retrieval()  # better to use hook functions.




class        lowerCAmelCase__       (  pl.Callback					):
   def 	__a				(						self	:     List[Any]				,       snake_case__	:     Union[str, Any]				,       snake_case__	:     Tuple				):


        '''simple docstring'''
        UpperCAmelCase__		:  Optional[int]				=				trainer.lr_schedulers[0]["scheduler"]
        UpperCAmelCase__		:  Dict				=				{f'lr_group_{i}': lr for i, lr in enumerate(lr_scheduler.get_lr()				)}
        pl_module.logger.log_metrics(snake_case__				)
   def 	__a				(						self	:     Dict				,       snake_case__	:     pl.Trainer				,       snake_case__	:     pl.LightningModule				):


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




   def 	__a				(						self	:     int				,       snake_case__	:     pl.Trainer				,       snake_case__	:     pl.LightningModule				):


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




def      SCREAMING_SNAKE_CASE__    (   snake_case    :      Dict       , snake_case    :      Union[str, Any]							)->						None:






     '''simple docstring'''



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



def      SCREAMING_SNAKE_CASE__    (   snake_case    :      BaseTransformer       , snake_case    :      argparse.Namespace       , snake_case    :      Any=None       , snake_case    :      Tuple=True       , snake_case    :      Optional[Any]=None       , snake_case    :      Optional[int]=[]       , snake_case    :      Any=None       , snake_case    :      str=None       , **snake_case    :      Optional[Any]       , )->						str:






     '''simple docstring'''



     pl.seed_everything(args.seed							)

     # init model
     UpperCAmelCase__		:  List[Any]				=				Path(model.hparams.output_dir							)
     odir.mkdir(exist_ok=snake_case							)

     # add custom checkpoints
     if checkpoint_callback is None:
          UpperCAmelCase__		:  Tuple				=				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(snake_case							)
     if logging_callback is None:
          UpperCAmelCase__		:  List[str]				=				LoggingCallback()

     UpperCAmelCase__		:  int				=				{}

     # TODO: remove with PyTorch 1.6 since pl uses native amp
     if args.fpaa:
          UpperCAmelCase__		:  Tuple				=				16
          # train_params["amp_level"] = args.fp16_opt_level

     if args.gpus > 1:
          UpperCAmelCase__		:  List[Any]				=				"auto"  # "ddp"
          UpperCAmelCase__		:  List[Any]				=				"ddp"

     UpperCAmelCase__		:  Dict				=				args.accumulate_grad_batches
     UpperCAmelCase__		:  List[Any]				=				None  # extra_train_kwargs.get("profiler", None) #get unwanted logs
     UpperCAmelCase__		:  List[str]				=				"auto"

     UpperCAmelCase__		:  Tuple				=				pl.Trainer.from_argparse_args(
         snake_case       , weights_summary=snake_case       , callbacks=[logging_callback] + extra_callbacks + [checkpoint_callback] + [InitCallback()]       , logger=snake_case       , **snake_case       , )

     if args.do_train:
          trainer.fit(snake_case							)

     return trainer










"""simple docstring"""



from doctest import testmod
from math import sqrt


def      SCREAMING_SNAKE_CASE__    (   snake_case    :      int							)->						list:






     '''simple docstring'''



     UpperCAmelCase__		:  list[int]				=				[]
     if num < 1:
          return facs
     facs.append(1							)
     if num == 1:
          return facs
     facs.append(snake_case							)
     for i in range(2       , int(sqrt(snake_case							)							) + 1							):
          if num % i == 0:  # If i is a factor of num
               facs.append(snake_case							)
               UpperCAmelCase__		:  Any				=				num // i  # num//i is the other factor of num
               if d != i:  # If d and i are distinct
                    facs.append(snake_case							)  # we have found another factor
     facs.sort()
     return facs


if __name__ == "__main__":
 testmod(name="""factors_of_a_number""", verbose=True)


