



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

			"""simple docstring"""


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

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


			def   snake_case__	(      self						,	*SCREAMING_SNAKE_CASE__						,	**SCREAMING_SNAKE_CASE__		)      ->	Optional[int]:
				A__											=					self.model_type.from_pretrained(*snake_case_						,	**snake_case_		)


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


			def   snake_case__	(      self		)      ->	int:
				A__											=					self.model
				A__											=					["bias", "LayerNorm.weight"]
				A__											=					[
				    {
				        "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__											=					Adafactor(
					    snake_case_						,	lr=self.hparams.learning_rate						,	scale_parameter=snake_case_						,	relative_step=snake_case_		)

				else:
					A__											=					AdamW(
					    snake_case_						,	lr=self.hparams.learning_rate						,	eps=self.hparams.adam_epsilon		)
				A__											=					optimizer

				A__											=					self.get_lr_scheduler()

				return [optimizer], [scheduler]


			def   snake_case__	(      self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__		)      ->	List[str]:
				return self.validation_step(snake_case_						,	snake_case_		)


			def   snake_case__	(      self						,	SCREAMING_SNAKE_CASE__		)      ->	Any:
				return self.validation_end(snake_case_		)


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


			def   snake_case__	(      self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__ = False		)      ->	Any:
				raise NotImplementedError("You must implement this for your task"		)


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


			def   snake_case__	(      self		)      ->	List[Any]:
				return self.get_dataloader("dev"						,	self.hparams.eval_batch_size						,	shuffle=snake_case_		)


			def   snake_case__	(      self		)      ->	Optional[Any]:
				return self.get_dataloader("test"						,	self.hparams.eval_batch_size						,	shuffle=snake_case_		)


			def   snake_case__	(      self						,	SCREAMING_SNAKE_CASE__		)      ->	Dict:
				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   snake_case__	(      self						,	SCREAMING_SNAKE_CASE__		)      ->	None:
				A__											=					self.output_dir.joinpath("best_tfmr"		)
				A__											=					self.step_count
				self.model.save_pretrained(snake_case_		)
				self.tokenizer.save_pretrained(snake_case_		)







			@staticmethod
			def   snake_case__	(      SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__		)      ->	Optional[Any]:
				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=str(Path(snake_case_		).parent / "test_run" / "cache"		)						,	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=32						,	type=snake_case_		)
				parser.add_argument("--eval_batch_size"						,	default=32						,	type=snake_case_		)
				parser.add_argument("--adafactor"						,	action="store_true"		)






class       UpperCamelCase__			(   pl.Callback      ):

			"""simple docstring"""


			def   snake_case__	(      self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__		)      ->	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       UpperCamelCase__			(   pl.Callback      ):

			"""simple docstring"""


			def   snake_case__	(      self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__		)      ->	str:
				# print(pl_module.model.rag)
				for name, param in pl_module.model.rag.named_parameters():
					if param.grad is None:
						print(snake_case_		)






class       UpperCamelCase__			(   pl.Callback      ):

			"""simple docstring"""


			def   snake_case__	(      self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__		)      ->	str:
				A__											=					trainer.lr_schedulers[0]["scheduler"]
				A__											=					{f"""lr_group_{i}""": lr for i, lr in enumerate(lr_scheduler.get_lr()		)}
				pl_module.logger.log_metrics(snake_case_		)


			def   snake_case__	(      self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__		)      ->	str:
				rank_zero_info("***** Validation results *****"		)
				A__											=					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   snake_case__	(      self						,	SCREAMING_SNAKE_CASE__						,	SCREAMING_SNAKE_CASE__		)      ->	Optional[Any]:
				rank_zero_info("***** Test results *****"		)
				A__											=					trainer.callback_metrics
				# Log and save results to file
				A__											=					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       _lowerCamelCase				(							UpperCAmelCase_ :       List[str],	UpperCAmelCase_ :       List[Any]     )	-> List[str]:






	"""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=42,	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       _lowerCamelCase				(							UpperCAmelCase_ :       Any,	UpperCAmelCase_ :       List[Any],	UpperCAmelCase_ :       Union[str, Any]=None,	UpperCAmelCase_ :       Optional[Any]=True,	UpperCAmelCase_ :       List[Any]=[],	UpperCAmelCase_ :       List[str]=None,	UpperCAmelCase_ :       List[str]=None,	**UpperCAmelCase_ :       Dict,	)	-> Dict:






	"""simple docstring"""





	pl.seed_everything(args.seed     )

	# init model
	A__											=					Path(model.hparams.output_dir     )
	odir.mkdir(exist_ok=_lowerCAmelCase     )

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

	A__											=					{}

	if args.fpaa:
		A__											=					16

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

	A__											=					args.accumulate_grad_batches
	A__											=					None
	A__											=					"auto"

	A__											=					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 operator as op


SCREAMING_SNAKE_CASE_					     =	'''scaler.pt'''
SCREAMING_SNAKE_CASE_					     =	'''pytorch_model'''
SCREAMING_SNAKE_CASE_					     =	'''random_states'''
SCREAMING_SNAKE_CASE_					     =	'''optimizer'''
SCREAMING_SNAKE_CASE_					     =	'''scheduler'''
SCREAMING_SNAKE_CASE_					     =	'''pytorch_model.bin'''
SCREAMING_SNAKE_CASE_					     =	'''pytorch_model.bin.index.json'''
SCREAMING_SNAKE_CASE_					     =	'''model.safetensors'''
SCREAMING_SNAKE_CASE_					     =	'''model.safetensors.index.json'''
SCREAMING_SNAKE_CASE_					     =	'''1.10.2'''
SCREAMING_SNAKE_CASE_					     =	'''py38'''
SCREAMING_SNAKE_CASE_					     =	'''4.17.0'''
SCREAMING_SNAKE_CASE_					     =	['''ml.p3.16xlarge''', '''ml.p3dn.24xlarge''', '''ml.p4dn.24xlarge''']
SCREAMING_SNAKE_CASE_					     =	['''FULL_SHARD''', '''SHARD_GRAD_OP''', '''NO_SHARD''', '''HYBRID_SHARD''', '''HYBRID_SHARD_ZERO2''']
SCREAMING_SNAKE_CASE_					     =	['''TRANSFORMER_BASED_WRAP''', '''SIZE_BASED_WRAP''', '''NO_WRAP''']
SCREAMING_SNAKE_CASE_					     =	['''BACKWARD_PRE''', '''BACKWARD_POST''', '''NO_PREFETCH''']
SCREAMING_SNAKE_CASE_					     =	['''FULL_STATE_DICT''', '''LOCAL_STATE_DICT''', '''SHARDED_STATE_DICT''']
SCREAMING_SNAKE_CASE_					     =	'''2.0.1'''
SCREAMING_SNAKE_CASE_					     =	['''pdsh''', '''standard''', '''openmpi''', '''mvapich''']
SCREAMING_SNAKE_CASE_					     =	['''default''', '''reduce-overhead''', '''max-autotune''']

SCREAMING_SNAKE_CASE_					     =	{'''>''': op.gt, '''>=''': op.ge, '''==''': op.eq, '''!=''': op.ne, '''<=''': op.le, '''<''': op.lt}

# These are the args for `torch.distributed.launch` for pytorch < 1.9
SCREAMING_SNAKE_CASE_					     =	[
    '''nnodes''',
    '''nproc_per_node''',
    '''rdzv_backend''',
    '''rdzv_endpoint''',
    '''rdzv_id''',
    '''rdzv_conf''',
    '''standalone''',
    '''max_restarts''',
    '''monitor_interval''',
    '''start_method''',
    '''role''',
    '''module''',
    '''m''',
    '''no_python''',
    '''run_path''',
    '''log_dir''',
    '''r''',
    '''redirects''',
    '''t''',
    '''tee''',
    '''node_rank''',
    '''master_addr''',
    '''master_port''',
]

SCREAMING_SNAKE_CASE_					     =	['''DEEPSPEED''', '''MULTI_GPU''', '''FSDP''', '''MEGATRON_LM''']
SCREAMING_SNAKE_CASE_					     =	['''DEEPSPEED''', '''MULTI_XPU''', '''FSDP''']
