"""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       :			Union[str, Any]   	=logging.getLogger(__name__)

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

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





class 					_A		(			pl.LightningModule      ):

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







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

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

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

      def 							A__			(		self						,			*__lowerCAmelCase						,			**__lowerCAmelCase		):







            """simple docstring"""



            lowercase								=      self.model_type.from_pretrained(*__lowerCAmelCase						,			**__lowerCAmelCase		)

      def 							A__			(		self		):







            """simple docstring"""



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

      def 							A__			(		self		):







            """simple docstring"""



            lowercase								=      self.model
            lowercase								=      ["""bias""", """LayerNorm.weight"""]
            lowercase								=      [
                {
                    """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:
                  lowercase								=      Adafactor(
                      __lowerCAmelCase						,			lr=self.hparams.learning_rate						,			scale_parameter=__lowerCAmelCase						,			relative_step=__lowerCAmelCase		)

            else:
                  lowercase								=      AdamW(
                      __lowerCAmelCase						,			lr=self.hparams.learning_rate						,			eps=self.hparams.adam_epsilon		)
            lowercase								=      optimizer

            lowercase								=      self.get_lr_scheduler()

            return [optimizer], [scheduler]

      def 							A__			(		self						,			__lowerCAmelCase						,			__lowerCAmelCase		):







            """simple docstring"""



            return self.validation_step(__lowerCAmelCase						,			__lowerCAmelCase		)

      def 							A__			(		self						,			__lowerCAmelCase		):







            """simple docstring"""



            return self.validation_end(__lowerCAmelCase		)

      def 							A__			(		self		):







            """simple docstring"""



            lowercase								=      max(1						,			self.hparams.gpus		)  # TODO: consider num_tpu_cores
            lowercase								=      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						,			__lowerCAmelCase		):







            """simple docstring"""



            if stage == "test":
                  lowercase								=      len(self.test_dataloader().dataset		)
            else:
                  lowercase								=      self.get_dataloader("""train"""						,			self.hparams.train_batch_size						,			shuffle=__lowerCAmelCase		)
                  lowercase								=      len(self.train_dataloader().dataset		)

      def 							A__			(		self						,			__lowerCAmelCase						,			__lowerCAmelCase						,			__lowerCAmelCase = False		):







            """simple docstring"""



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

      def 							A__			(		self		):







            """simple docstring"""



            return self.train_loader

      def 							A__			(		self		):







            """simple docstring"""



            return self.get_dataloader("""dev"""						,			self.hparams.eval_batch_size						,			shuffle=__lowerCAmelCase		)

      def 							A__			(		self		):







            """simple docstring"""



            return self.get_dataloader("""test"""						,			self.hparams.eval_batch_size						,			shuffle=__lowerCAmelCase		)

      def 							A__			(		self						,			__lowerCAmelCase		):







            """simple docstring"""



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

      @pl.utilities.rank_zero_only
      def 							A__			(		self						,			__lowerCAmelCase		):







            """simple docstring"""



            lowercase								=      self.output_dir.joinpath("""best_tfmr"""		)
            lowercase								=      self.step_count
            self.model.save_pretrained(__lowerCAmelCase		)
            self.tokenizer.save_pretrained(__lowerCAmelCase		)







      @staticmethod
      def 							A__			(		__lowerCAmelCase						,			__lowerCAmelCase		):







            """simple docstring"""



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






class 					_A		(			pl.Callback      ):

      def 							A__			(		self						,			__lowerCAmelCase						,			__lowerCAmelCase		):







            """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 					_A		(			pl.Callback      ):

      def 							A__			(		self						,			__lowerCAmelCase						,			__lowerCAmelCase		):







            """simple docstring"""



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






class 					_A		(			pl.Callback      ):

      def 							A__			(		self						,			__lowerCAmelCase						,			__lowerCAmelCase		):







            """simple docstring"""



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

      def 							A__			(		self						,			__lowerCAmelCase						,			__lowerCAmelCase		):







            """simple docstring"""



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







      def 							A__			(		self						,			__lowerCAmelCase						,			__lowerCAmelCase		):







            """simple docstring"""



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






def      UpperCAmelCase__			(							lowerCAmelCase__	:Union[str, Any]   ,						lowerCAmelCase__	:Optional[Any]						)					->       None:
      '''simple docstring'''


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






def      UpperCAmelCase__			(							lowerCAmelCase__	:BaseTransformer   ,						lowerCAmelCase__	:argparse.Namespace   ,						lowerCAmelCase__	:List[str]=None   ,						lowerCAmelCase__	:int=True   ,						lowerCAmelCase__	:Dict=[]   ,						lowerCAmelCase__	:Optional[Any]=None   ,						lowerCAmelCase__	:Optional[Any]=None   ,						**lowerCAmelCase__	:int   ,						)					->       str:
      '''simple docstring'''


      pl.seed_everything(args.seed						)

      # init model
      lowercase								=      Path(model.hparams.output_dir						)
      odir.mkdir(exist_ok=lowerCAmelCase__						)

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

      lowercase								=      {}

      if args.fpaa:
            lowercase								=      1_6

      if args.gpus > 1:
            lowercase								=      """auto"""
            lowercase								=      """ddp"""

      lowercase								=      args.accumulate_grad_batches
      lowercase								=      None
      lowercase								=      """auto"""

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

      if args.do_train:
            trainer.fit(lowerCAmelCase__						)

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




"""simple docstring"""






import os
import unittest

from transformers.models.transfo_xl.tokenization_transfo_xl import VOCAB_FILES_NAMES, TransfoXLTokenizer

from ...test_tokenization_common import TokenizerTesterMixin





class 					_A		(			lowerCAmelCase    ,       unittest.TestCase      ):
      snake_case__     :  str				    =      TransfoXLTokenizer
      snake_case__     :  Union[str, Any]				    =      False
      snake_case__     :  Union[str, Any]				    =      False

      def 							A__			(		self		):







            """simple docstring"""



            super().setUp()

            lowercase								=      [
                """<unk>""",
                """[CLS]""",
                """[SEP]""",
                """want""",
                """unwanted""",
                """wa""",
                """un""",
                """running""",
                """,""",
                """low""",
                """l""",
            ]
            lowercase								=      os.path.join(self.tmpdirname						,			VOCAB_FILES_NAMES["""vocab_file"""]		)
            with open(self.vocab_file						,			"""w"""						,			encoding="""utf-8"""		) as vocab_writer:
                  vocab_writer.write("""""".join([x + """\n""" for x in vocab_tokens]		)		)

      def 							A__			(		self						,			**__lowerCAmelCase		):







            """simple docstring"""



            lowercase								=      True
            return TransfoXLTokenizer.from_pretrained(self.tmpdirname						,			**__lowerCAmelCase		)

      def 							A__			(		self						,			__lowerCAmelCase		):







            """simple docstring"""



            lowercase								=      """<unk> UNwanted , running"""
            lowercase								=      """<unk> unwanted, running"""
            return input_text, output_text

      def 							A__			(		self		):







            """simple docstring"""



            lowercase								=      TransfoXLTokenizer(vocab_file=self.vocab_file						,			lower_case=__lowerCAmelCase		)

            lowercase								=      tokenizer.tokenize("""<unk> UNwanted , running"""		)
            self.assertListEqual(__lowerCAmelCase						,			["""<unk>""", """unwanted""", """,""", """running"""]		)

            self.assertListEqual(tokenizer.convert_tokens_to_ids(__lowerCAmelCase		)						,			[0, 4, 8, 7]		)

      def 							A__			(		self		):







            """simple docstring"""



            lowercase								=      TransfoXLTokenizer(lower_case=__lowerCAmelCase		)

            self.assertListEqual(
                tokenizer.tokenize(""" \tHeLLo ! how  \n Are yoU ?  """		)						,			["""hello""", """!""", """how""", """are""", """you""", """?"""]		)

      def 							A__			(		self		):







            """simple docstring"""



            lowercase								=      TransfoXLTokenizer(lower_case=__lowerCAmelCase		)

            self.assertListEqual(
                tokenizer.tokenize(""" \tHeLLo ! how  \n Are yoU ?  """		)						,			["""HeLLo""", """!""", """how""", """Are""", """yoU""", """?"""]		)

      def 							A__			(		self		):







            """simple docstring"""



            lowercase								=      TransfoXLTokenizer(lower_case=__lowerCAmelCase		)
            lowercase								=      """Hello (bracket) and side-scrolled [and] Henry's $5,000 with 3.34 m. What's up!?"""
            lowercase								=      [
                """Hello""",
                """(""",
                """bracket""",
                """)""",
                """and""",
                """side""",
                """@-@""",
                """scrolled""",
                """[""",
                """and""",
                """]""",
                """Henry""",
                """'s""",
                """$""",
                """5""",
                """@,@""",
                """000""",
                """with""",
                """3""",
                """@.@""",
                """34""",
                """m""",
                """.""",
                """What""",
                """'s""",
                """up""",
                """!""",
                """?""",
            ]

            self.assertListEqual(tokenizer.tokenize(__lowerCAmelCase		)						,			__lowerCAmelCase		)

            self.assertEqual(tokenizer.convert_tokens_to_string(__lowerCAmelCase		)						,			__lowerCAmelCase		)







      def 							A__			(		self		):







            """simple docstring"""



            lowercase								=      self.get_tokenizer()
            lowercase								=      len(__lowerCAmelCase		)

            tokenizer.add_tokens(["""new1""", """new2"""]		)
            tokenizer.move_added_token("""new1"""						,			1		)

            # Check that moved token is not copied (duplicate)
            self.assertEqual(len(__lowerCAmelCase		)						,			original_len + 2		)
            # Check that token is moved to specified id
            self.assertEqual(tokenizer.encode("""new1"""		)						,			[1]		)
            self.assertEqual(tokenizer.decode([1]		)						,			"""new1"""		)



