






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

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

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





class        _lowerCamelCase      (       pl.LightningModule ):


						def __init__(self   ,							__a   ,							__a=None   ,							__a="base"   ,							__a=None   ,							__a=None   ,							__a=None   ,							**__a   ,							) ->   str:
								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_		)
								UpperCamelCase   =	0
								UpperCamelCase   =	Path(self.hparams.output_dir		)
								UpperCamelCase   =	self.hparams.cache_dir if self.hparams.cache_dir else None
								if config is None:
										UpperCamelCase   =	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:
										UpperCamelCase   =	config

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


						def     snake_case_   (self   ,							*__a   ,							**__a		) ->   List[str]:
								UpperCamelCase   =	self.model_type.from_pretrained(*A_   ,							**A_		)


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


						def     snake_case_   (self		) ->   Union[str, Any]:
								UpperCamelCase   =	self.model
								UpperCamelCase   =	["bias", "LayerNorm.weight"]
								UpperCamelCase   =	[
								    {
								        "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:
										UpperCamelCase   =	Adafactor(
										    A_   ,							lr=self.hparams.learning_rate   ,							scale_parameter=A_   ,							relative_step=A_		)

								else:
										UpperCamelCase   =	AdamW(
										    A_   ,							lr=self.hparams.learning_rate   ,							eps=self.hparams.adam_epsilon		)
								UpperCamelCase   =	optimizer

								UpperCamelCase   =	self.get_lr_scheduler()

								return [optimizer], [scheduler]


						def     snake_case_   (self   ,							__a   ,							__a		) ->   str:
								return self.validation_step(A_   ,							A_		)


						def     snake_case_   (self   ,							__a		) ->   Optional[int]:
								return self.validation_end(A_		)


						def     snake_case_   (self		) ->   int:
								UpperCamelCase   =	max(1   ,							self.hparams.gpus		)  # TODO: consider num_tpu_cores
								UpperCamelCase   =	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   ,							__a		) ->   Union[str, Any]:
								if stage == "test":
										UpperCamelCase   =	len(self.test_dataloader().dataset		)
								else:
										UpperCamelCase   =	self.get_dataloader("train"   ,							self.hparams.train_batch_size   ,							shuffle=A_		)
										UpperCamelCase   =	len(self.train_dataloader().dataset		)


						def     snake_case_   (self   ,							__a   ,							__a   ,							__a = False		) ->   List[str]:
								raise NotImplementedError("You must implement this for your task"		)


						def     snake_case_   (self		) ->   Optional[int]:
								return self.train_loader


						def     snake_case_   (self		) ->   str:
								return self.get_dataloader("dev"   ,							self.hparams.eval_batch_size   ,							shuffle=A_		)


						def     snake_case_   (self		) ->   str:
								return self.get_dataloader("test"   ,							self.hparams.eval_batch_size   ,							shuffle=A_		)


						def     snake_case_   (self   ,							__a		) ->   str:
								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     snake_case_   (self   ,							__a		) ->   None:
								UpperCamelCase   =	self.output_dir.joinpath("best_tfmr"		)
								UpperCamelCase   =	self.step_count
								self.model.save_pretrained(A_		)
								self.tokenizer.save_pretrained(A_		)




						@staticmethod
						def     snake_case_   (__a   ,							__a		) ->   Any:
								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        _lowerCamelCase      (       pl.Callback ):


						def     snake_case_   (self   ,							__a   ,							__a		) ->   Optional[Any]:
								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     snake_case_   (self   ,							__a   ,							__a		) ->   Union[str, Any]:
								# print(pl_module.model.rag)
								for name, param in pl_module.model.rag.named_parameters():
										if param.grad is None:
												print(A_		)





class        _lowerCamelCase      (       pl.Callback ):


						def     snake_case_   (self   ,							__a   ,							__a		) ->   Optional[int]:
								UpperCamelCase   =	trainer.lr_schedulers[0]["scheduler"]
								UpperCamelCase   =	{F"lr_group_{i}": lr for i, lr in enumerate(lr_scheduler.get_lr()		)}
								pl_module.logger.log_metrics(A_		)


						def     snake_case_   (self   ,							__a   ,							__a		) ->   Dict:
								rank_zero_info("***** Validation results *****"		)
								UpperCamelCase   =	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     snake_case_   (self   ,							__a   ,							__a		) ->   str:
								rank_zero_info("***** Test results *****"		)
								UpperCamelCase   =	trainer.callback_metrics
								# Log and save results to file
								UpperCamelCase   =	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       a__					(					_SCREAMING_SNAKE_CASE  ,	_SCREAMING_SNAKE_CASE    ):




		"""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  ,	):




		"""simple docstring"""


		pl.seed_everything(args.seed    )

		# init model
		UpperCamelCase   =	Path(model.hparams.output_dir    )
		odir.mkdir(exist_ok=_SCREAMING_SNAKE_CASE    )

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

		UpperCamelCase   =	{}

		if args.fpaa:
				UpperCamelCase   =	16

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

		UpperCamelCase   =	args.accumulate_grad_batches
		UpperCamelCase   =	None
		UpperCamelCase   =	"auto"

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





def       a__					(					_SCREAMING_SNAKE_CASE    ):




		"""simple docstring"""


		return number & 1 == 0


if __name__ == "__main__":
						import doctest

						doctest.testmod()






