





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

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

__A											=		{
    """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											=		{
    """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											=		sorted(arg_to_scheduler.keys())
__A											=		"""{""" + """, """.join(arg_to_scheduler_choices) + """}"""






class 					_lowerCAmelCase   (     pl.LightningModule       ):
       """simple docstring"""




       def __init__(    self    , __UpperCAmelCase    , __UpperCAmelCase=None    , __UpperCAmelCase="base"    , __UpperCAmelCase=None    , __UpperCAmelCase=None    , __UpperCAmelCase=None    , **__UpperCAmelCase    , ):



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

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

         if tokenizer is None:
           lowerCAmelCase__			:Optional[Any]											=							AutoTokenizer.from_pretrained(
               self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path    , cache_dir=__UpperCAmelCase    , )
         else:
           lowerCAmelCase__			:PreTrainedTokenizer											=							tokenizer
         lowerCAmelCase__			:Any											=							MODEL_MODES[mode]
         if model is None:
           lowerCAmelCase__			:Tuple											=							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=__UpperCAmelCase    , )
         else:
           lowerCAmelCase__			:Optional[Any]											=							model






       def      snake_case						(    self    , *__UpperCAmelCase    , **__UpperCAmelCase      ):



         '''simple docstring'''

         lowerCAmelCase__			:Optional[Any]											=							self.model_type.from_pretrained(*__UpperCAmelCase    , **__UpperCAmelCase      )






       def      snake_case						(    self      ):



         '''simple docstring'''

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






       def      snake_case						(    self      ):



         '''simple docstring'''

         lowerCAmelCase__			:List[str]											=							self.model
         lowerCAmelCase__			:List[str]											=							['bias', 'LayerNorm.weight']
         lowerCAmelCase__			:Any											=							[
             {
                 'params': [
                     p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay      )
                 ],  # check this named paramters
                 '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__			:Optional[int]											=							Adafactor(
               __UpperCAmelCase    , lr=self.hparams.learning_rate    , scale_parameter=__UpperCAmelCase    , relative_step=__UpperCAmelCase      )

         else:
           lowerCAmelCase__			:Optional[int]											=							AdamW(
               __UpperCAmelCase    , lr=self.hparams.learning_rate    , eps=self.hparams.adam_epsilon      )
         lowerCAmelCase__			:str											=							optimizer

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

         return [optimizer], [scheduler]






       def      snake_case						(    self    , __UpperCAmelCase    , __UpperCAmelCase      ):



         '''simple docstring'''

         return self.validation_step(__UpperCAmelCase    , __UpperCAmelCase      )






       def      snake_case						(    self    , __UpperCAmelCase      ):



         '''simple docstring'''

         return self.validation_end(__UpperCAmelCase      )






       def      snake_case						(    self      ):



         '''simple docstring'''

         lowerCAmelCase__			:Tuple											=							max(1    , self.hparams.gpus      )  # TODO: consider num_tpu_cores
         lowerCAmelCase__			:List[Any]											=							self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
         return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs






       def      snake_case						(    self    , __UpperCAmelCase      ):



         '''simple docstring'''

         if stage == "test":
           lowerCAmelCase__			:int											=							len(self.test_dataloader().dataset      )
         else:
           lowerCAmelCase__			:Any											=							self.get_dataloader('train'    , self.hparams.train_batch_size    , shuffle=__UpperCAmelCase      )
           lowerCAmelCase__			:Optional[int]											=							len(self.train_dataloader().dataset      )






       def      snake_case						(    self    , __UpperCAmelCase    , __UpperCAmelCase    , __UpperCAmelCase = False      ):



         '''simple docstring'''

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






       def      snake_case						(    self      ):



         '''simple docstring'''

         return self.train_loader






       def      snake_case						(    self      ):



         '''simple docstring'''

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






       def      snake_case						(    self      ):



         '''simple docstring'''

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






       def      snake_case						(    self    , __UpperCAmelCase      ):



         '''simple docstring'''

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






       @pl.utilities.rank_zero_only
       def      snake_case						(    self    , __UpperCAmelCase      ):



         '''simple docstring'''

         lowerCAmelCase__			:Optional[int]											=							self.output_dir.joinpath('best_tfmr'      )
         lowerCAmelCase__			:str											=							self.step_count
         self.model.save_pretrained(__UpperCAmelCase      )
         self.tokenizer.save_pretrained(__UpperCAmelCase      )






       @staticmethod
       def      snake_case						(    __UpperCAmelCase    , __UpperCAmelCase      ):



         '''simple docstring'''

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






class 					_lowerCAmelCase   (     pl.Callback       ):
       """simple docstring"""




       def      snake_case						(    self    , __UpperCAmelCase    , __UpperCAmelCase      ):



         '''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       ):
       """simple docstring"""




       def      snake_case						(    self    , __UpperCAmelCase    , __UpperCAmelCase      ):



         '''simple docstring'''

         for name, param in pl_module.model.rag.named_parameters():
           if param.grad is None:
             print(__UpperCAmelCase      )






class 					_lowerCAmelCase   (     pl.Callback       ):
       """simple docstring"""




       def      snake_case						(    self    , __UpperCAmelCase    , __UpperCAmelCase      ):



         '''simple docstring'''

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






       def      snake_case						(    self    , __UpperCAmelCase    , __UpperCAmelCase      ):



         '''simple docstring'''

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






       def      snake_case						(    self    , __UpperCAmelCase    , __UpperCAmelCase      ):



         '''simple docstring'''

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






def  __A   (_SCREAMING_SNAKE_CASE     ,   _SCREAMING_SNAKE_CASE			)					->None:





  """simple docstring"""


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




def  __A   (_SCREAMING_SNAKE_CASE     ,   _SCREAMING_SNAKE_CASE     ,   _SCREAMING_SNAKE_CASE=None     ,   _SCREAMING_SNAKE_CASE=True     ,   _SCREAMING_SNAKE_CASE=[]     ,   _SCREAMING_SNAKE_CASE=None     ,   _SCREAMING_SNAKE_CASE=None     ,   **_SCREAMING_SNAKE_CASE     ,   )					->Dict:





  """simple docstring"""


  pl.seed_everything(args.seed			)

  # init model
  lowerCAmelCase__			:Dict											=							Path(model.hparams.output_dir			)
  odir.mkdir(exist_ok=_SCREAMING_SNAKE_CASE			)

  # add custom checkpoints
  if checkpoint_callback is None:
    lowerCAmelCase__			:int											=							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(_SCREAMING_SNAKE_CASE			)
  if logging_callback is None:
    lowerCAmelCase__			:Tuple											=							LoggingCallback()

  lowerCAmelCase__			:Union[str, Any]											=							{}

  if args.fpaa:
    lowerCAmelCase__			:List[str]											=							16

  if args.gpus > 1:
    lowerCAmelCase__			:Union[str, Any]											=							'auto'
    lowerCAmelCase__			:Optional[int]											=							'ddp'

  lowerCAmelCase__			:Optional[int]											=							args.accumulate_grad_batches
  lowerCAmelCase__			:Optional[int]											=							None
  lowerCAmelCase__			:List[str]											=							'auto'

  lowerCAmelCase__			:Union[str, Any]											=							pl.Trainer.from_argparse_args(
      _SCREAMING_SNAKE_CASE     ,   weights_summary=_SCREAMING_SNAKE_CASE     ,   callbacks=[logging_callback] + extra_callbacks + [InitCallback()] + [checkpoint_callback]     ,   logger=_SCREAMING_SNAKE_CASE     ,   val_check_interval=1     ,   num_sanity_val_steps=2     ,   **_SCREAMING_SNAKE_CASE     ,   )

  if args.do_train:
    trainer.fit(_SCREAMING_SNAKE_CASE			)

  else:
    print('RAG modeling tests with new set functions successfuly executed!'			)
  return trainer







"""simple docstring"""





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tokenizers_available,
    is_torch_available,
)


__A											=		{"""configuration_reformer""": ["""REFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP""", """ReformerConfig"""]}

try:
  if not is_sentencepiece_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __A											=		["""ReformerTokenizer"""]

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __A											=		["""ReformerTokenizerFast"""]

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __A											=		[
      """REFORMER_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """ReformerAttention""",
      """ReformerForMaskedLM""",
      """ReformerForQuestionAnswering""",
      """ReformerForSequenceClassification""",
      """ReformerLayer""",
      """ReformerModel""",
      """ReformerModelWithLMHead""",
      """ReformerPreTrainedModel""",
  ]


if TYPE_CHECKING:
  from .configuration_reformer import REFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, ReformerConfig

  try:
    if not is_sentencepiece_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_reformer import ReformerTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_reformer_fast import ReformerTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_reformer import (
        REFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
        ReformerAttention,
        ReformerForMaskedLM,
        ReformerForQuestionAnswering,
        ReformerForSequenceClassification,
        ReformerLayer,
        ReformerModel,
        ReformerModelWithLMHead,
        ReformerPreTrainedModel,
    )

else:
  import sys

  __A											=		_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
