



'''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							:str				    =					logging.getLogger(__name__)

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

lowerCamelCase							:Union[str, Any]				    =					{
    '''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							:List[str]				    =					{
    '''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							:Tuple				    =					'''{''' + ''', '''.join(arg_to_scheduler_choices) + '''}'''







class 	_lowerCAmelCase							(			pl.LightningModule ):


	def __init__(self       ,      lowercase       ,      lowercase=None       ,      lowercase="base"       ,      lowercase=None       ,      lowercase=None       ,      lowercase=None       ,      **lowercase       ,      ):
				super().__init__()
				# TODO: move to self.save_hyperparameters()
				# self.save_hyperparameters()
				# can also expand arguments into trainer signature for easier reading

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

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

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


	def     _a	(self       ,      *lowercase       ,      **lowercase							):
				A_	:						Any	=     self.model_type.from_pretrained(*lowercase       ,      **lowercase							)


	def     _a	(self							):
				A_	:						str	=     arg_to_scheduler[self.hparams.lr_scheduler]
				A_	:						Union[str, Any]	=     get_schedule_func(
				    self.opt       ,      num_warmup_steps=self.hparams.warmup_steps       ,      num_training_steps=self.total_steps()							)
				A_	:						Any	=     {"""scheduler""": scheduler, """interval""": """step""", """frequency""": 1}
				return scheduler


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

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

				A_	:						int	=     self.get_lr_scheduler()

				return [optimizer], [scheduler]


	def     _a	(self       ,      lowercase       ,      lowercase							):
				return self.validation_step(lowercase       ,      lowercase							)


	def     _a	(self       ,      lowercase							):
				return self.validation_end(lowercase							)


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


	def     _a	(self       ,      lowercase       ,      lowercase       ,      lowercase = False							):
				raise NotImplementedError("""You must implement this for your task"""							)


	def     _a	(self							):
				return self.train_loader


	def     _a	(self							):
				return self.get_dataloader("""dev"""       ,      self.hparams.eval_batch_size       ,      shuffle=lowercase							)


	def     _a	(self							):
				return self.get_dataloader("""test"""       ,      self.hparams.eval_batch_size       ,      shuffle=lowercase							)


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


	@pl.utilities.rank_zero_only
	def     _a	(self       ,      lowercase							):
				A_	:						Union[str, Any]	=     self.output_dir.joinpath("""best_tfmr"""							)
				A_	:						Union[str, Any]	=     self.step_count
				self.model.save_pretrained(lowercase							)
				self.tokenizer.save_pretrained(lowercase							)


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







class 	_lowerCAmelCase							(			pl.Callback ):


	def     _a	(self       ,      lowercase       ,      lowercase							):
				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     _a	(self       ,      lowercase       ,      lowercase							):
				# print(pl_module.model.rag)
				for name, param in pl_module.model.rag.named_parameters():
							if param.grad is None:
										print(lowercase							)







class 	_lowerCAmelCase							(			pl.Callback ):


	def     _a	(self       ,      lowercase       ,      lowercase							):
				A_	:						int	=     trainer.lr_schedulers[0]["""scheduler"""]
				A_	:						str	=     {F'lr_group_{i}': lr for i, lr in enumerate(lr_scheduler.get_lr()							)}
				pl_module.logger.log_metrics(lowercase							)


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


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






def 	a			( lowerCamelCase__						,      lowerCamelCase__    ):
			'''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 	a			( lowerCamelCase__						,      lowerCamelCase__						,      lowerCamelCase__=None						,      lowerCamelCase__=True						,      lowerCamelCase__=[]						,      lowerCamelCase__=None						,      lowerCamelCase__=None						,      **lowerCamelCase__						,      ):
			'''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_	:						Tuple	=     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_	:						str	=     LoggingCallback()

			A_	:						Optional[Any]	=     {}

			if args.fpaa:
						A_	:						Union[str, Any]	=     16

			if args.gpus > 1:
						A_	:						Optional[Any]	=     """auto"""
						A_	:						List[Any]	=     """ddp"""

			A_	:						Union[str, Any]	=     args.accumulate_grad_batches
			A_	:						Any	=     None
			A_	:						List[str]	=     """auto"""

			A_	:						Tuple	=     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'''




# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse

from .config import config_command_parser
from .config_args import default_config_file, load_config_from_file  # noqa: F401
from .default import default_command_parser
from .update import update_command_parser



def 	a			( lowerCamelCase__=None    ):
			'''simple docstring'''





			A_	:						int	=     argparse.ArgumentParser(add_help=lowerCamelCase__						,      allow_abbrev=lowerCamelCase__    )
			# The main config parser
			A_	:						int	=     config_command_parser(lowerCamelCase__    )
			# The subparser to add commands to
			A_	:						int	=     config_parser.add_subparsers(title="""subcommands"""						,      dest="""subcommand"""    )

			# Then add other parsers with the parent parser
			default_command_parser(lowerCamelCase__						,      parents=[parent_parser]    )
			update_command_parser(lowerCamelCase__						,      parents=[parent_parser]    )

			return config_parser



def 	a			( ):
			'''simple docstring'''





			A_	:						Optional[int]	=     get_config_parser()
			A_	:						List[str]	=     config_parser.parse_args()

			if not hasattr(lowerCamelCase__						,      """func"""    ):
						config_parser.print_help()
						exit(1    )

			# Run
			args.func(lowerCamelCase__    )


if __name__ == "__main__":
						main()