





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


__UpperCAmelCase						=			logging.getLogger(__name__)

require_version('''pytorch_lightning>=1.0.4''')

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



class 					__UpperCAmelCase							(						pl.LightningModule  ):
				def __init__(		self    :			Any	, a_    :			argparse.Namespace	, a_    :			Optional[int]=None	, a_    :			List[Any]="base"	, a_    :			str=None	, a_    :			Union[str, Any]=None	, a_    :			Optional[Any]=None	, **a_    :			Optional[int]	, )	->	Union[str, 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__  :	Optional[Any]		 =			0
							a__  :	List[str]		 =			Path(self.hparams.output_dir    )
							a__  :	str		 =			self.hparams.cache_dir if self.hparams.cache_dir else None
							if config is None:
										a__  :	str		 =			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__  :	PretrainedConfig		 =			config

							a__  :	Tuple		 =			("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__  :	Optional[int]		 =			AutoTokenizer.from_pretrained(
										    self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path	, cache_dir=a_	, )
							else:
										a__  :	PreTrainedTokenizer		 =			tokenizer
							a__  :	Dict		 =			MODEL_MODES[mode]
							if model is None:
										a__  :	Optional[int]		 =			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__  :	Tuple		 =			model
				def       UpperCAmelCase      (		self    :			Optional[int]	, *a_    :			int	, **a_    :			Optional[int]    )	->	Optional[Any]:



							'''simple docstring'''



							a__  :	Tuple		 =			self.model_type.from_pretrained(*a_	, **a_    )
				def       UpperCAmelCase      (		self    :			Tuple    )	->	Any:



							'''simple docstring'''



							a__  :	Dict		 =			arg_to_scheduler[self.hparams.lr_scheduler]
							a__  :	Optional[Any]		 =			get_schedule_func(
							    self.opt	, num_warmup_steps=self.hparams.warmup_steps	, num_training_steps=self.total_steps()    )
							a__  :	int		 =			{"scheduler": scheduler, "interval": "step", "frequency": 1}
							return scheduler
				def       UpperCAmelCase      (		self    :			Any    )	->	Optional[int]:



							'''simple docstring'''



							a__  :	Tuple		 =			self.model
							a__  :	str		 =			["bias", "LayerNorm.weight"]
							a__  :	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:
										a__  :	Optional[Any]		 =			Adafactor(
										    a_	, lr=self.hparams.learning_rate	, scale_parameter=a_	, relative_step=a_    )

							else:
										a__  :	Tuple		 =			AdamW(
										    a_	, lr=self.hparams.learning_rate	, eps=self.hparams.adam_epsilon    )
							a__  :	Dict		 =			optimizer

							a__  :	List[Any]		 =			self.get_lr_scheduler()

							return [optimizer], [scheduler]
				def       UpperCAmelCase      (		self    :			int	, a_    :			List[str]	, a_    :			Dict    )	->	Tuple:



							'''simple docstring'''



							return self.validation_step(a_	, a_    )
				def       UpperCAmelCase      (		self    :			List[str]	, a_    :			int    )	->	Tuple:



							'''simple docstring'''



							return self.validation_end(a_    )
				def       UpperCAmelCase      (		self    :			str    )	->	int:



							'''simple docstring'''



							a__  :	Optional[int]		 =			max(1	, self.hparams.gpus    )  # TODO: consider num_tpu_cores
							a__  :	Optional[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       UpperCAmelCase      (		self    :			List[str]	, a_    :			Any    )	->	Union[str, Any]:



							'''simple docstring'''



							if stage == "test":
										a__  :	str		 =			len(self.test_dataloader().dataset    )
							else:
										a__  :	Optional[int]		 =			self.get_dataloader("train"	, self.hparams.train_batch_size	, shuffle=a_    )
										a__  :	Any		 =			len(self.train_dataloader().dataset    )
				def       UpperCAmelCase      (		self    :			str	, a_    :			str	, a_    :			int	, a_    :			bool = False    )	->	Dict:



							'''simple docstring'''



							raise NotImplementedError("You must implement this for your task"    )
				def       UpperCAmelCase      (		self    :			str    )	->	str:



							'''simple docstring'''



							return self.train_loader
				def       UpperCAmelCase      (		self    :			List[Any]    )	->	Optional[int]:



							'''simple docstring'''



							return self.get_dataloader("dev"	, self.hparams.eval_batch_size	, shuffle=a_    )
				def       UpperCAmelCase      (		self    :			List[Any]    )	->	Optional[int]:



							'''simple docstring'''



							return self.get_dataloader("test"	, self.hparams.eval_batch_size	, shuffle=a_    )
				def       UpperCAmelCase      (		self    :			str	, a_    :			List[Any]    )	->	Dict:



							'''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       UpperCAmelCase      (		self    :			Union[str, Any]	, a_    :			Dict[str, Any]    )	->	None:



							'''simple docstring'''



							a__  :	int		 =			self.output_dir.joinpath("best_tfmr"    )
							a__  :	Any		 =			self.step_count
							self.model.save_pretrained(a_    )
							self.tokenizer.save_pretrained(a_    )






				@staticmethod
				def       UpperCAmelCase      (		a_    :			Union[str, Any]	, a_    :			str    )	->	Union[str, Any]:



							'''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=str(Path(a_    ).parent / "test_run" / "cache"    )	, 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       UpperCAmelCase      (		self    :			Tuple	, a_    :			Optional[int]	, a_    :			Dict    )	->	Optional[int]:



							'''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 					__UpperCAmelCase							(						pl.Callback  ):
				def       UpperCAmelCase      (		self    :			Dict	, a_    :			Union[str, Any]	, a_    :			Dict    )	->	Optional[Any]:



							'''simple docstring'''



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



class 					__UpperCAmelCase							(						pl.Callback  ):
				def       UpperCAmelCase      (		self    :			Union[str, Any]	, a_    :			Tuple	, a_    :			str    )	->	Union[str, Any]:



							'''simple docstring'''



							a__  :	Optional[Any]		 =			trainer.lr_schedulers[0]["scheduler"]
							a__  :	Optional[Any]		 =			{F"lr_group_{i}": lr for i, lr in enumerate(lr_scheduler.get_lr()    )}
							pl_module.logger.log_metrics(a_    )
				def       UpperCAmelCase      (		self    :			Tuple	, a_    :			pl.Trainer	, a_    :			pl.LightningModule    )	->	Optional[Any]:



							'''simple docstring'''



							rank_zero_info("***** Validation results *****"    )
							a__  :	Optional[int]		 =			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       UpperCAmelCase      (		self    :			List[Any]	, a_    :			pl.Trainer	, a_    :			pl.LightningModule    )	->	Union[str, Any]:



							'''simple docstring'''



							rank_zero_info("***** Test results *****"    )
							a__  :	List[str]		 =			trainer.callback_metrics
							# Log and save results to file
							a__  :	List[str]		 =			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      lowercase__ (					lowerCAmelCase__   :			Dict							,			lowerCAmelCase__   :			Optional[int]			)    ->			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=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      lowercase__ (					lowerCAmelCase__   :			BaseTransformer							,			lowerCAmelCase__   :			argparse.Namespace							,			lowerCAmelCase__   :			int=None							,			lowerCAmelCase__   :			Optional[Any]=True							,			lowerCAmelCase__   :			List[str]=[]							,			lowerCAmelCase__   :			Any=None							,			lowerCAmelCase__   :			Optional[int]=None							,			**lowerCAmelCase__   :			Any							,			)    ->			Tuple:



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

			# init model
			a__  :	Any		 =			Path(model.hparams.output_dir			)
			odir.mkdir(exist_ok=lowerCAmelCase__			)

			# add custom checkpoints
			if checkpoint_callback is None:
						a__  :	List[str]		 =			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:
						a__  :	Union[str, Any]		 =			LoggingCallback()

			a__  :	int		 =			{}

			if args.fpaa:
						a__  :	Optional[int]		 =			1_6

			if args.gpus > 1:
						a__  :	List[str]		 =			"auto"
						a__  :	Tuple		 =			"ddp"

			a__  :	Dict		 =			args.accumulate_grad_batches
			a__  :	str		 =			None
			a__  :	List[str]		 =			"auto"

			a__  :	Optional[int]		 =			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"""






from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__UpperCAmelCase						=			{
    '''configuration_informer''': [
        '''INFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''InformerConfig''',
    ],
}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__UpperCAmelCase						=			[
		    '''INFORMER_PRETRAINED_MODEL_ARCHIVE_LIST''',
		    '''InformerForPrediction''',
		    '''InformerModel''',
		    '''InformerPreTrainedModel''',
		]


if TYPE_CHECKING:
		from .configuration_informer import INFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, InformerConfig

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_informer import (
				    INFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
				    InformerForPrediction,
				    InformerModel,
				    InformerPreTrainedModel,
				)

else:
		import sys

		__UpperCAmelCase						=			_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)