






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


a_      = logging.getLogger(__name__)

require_version('pytorch_lightning>=1.0.4')

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





class 		SCREAMING_SNAKE_CASE__   (     pl.LightningModule    ):



		def __init__(    self:    Union[str, Any]				,   a:    Any				,   a:    Optional[int]=None				,   a:    int="base"				,   a:    str=None				,   a:    List[Any]=None				,   a:    str=None				,   **a:    Optional[int]				,   )  ->Optional[int]:






								'''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(lowercase_)
								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=lowercase_				,   **lowercase_				,   )
								else:
														a_				     =   config

								a_				     =   ('''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_				     =   AutoTokenizer.from_pretrained(
														    self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path				,   cache_dir=lowercase_				,   )
								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=lowercase_				,   )
								else:
														a_				     =   model



		def _lowerCAmelCase     (    self:    List[str]				,   *a:    Optional[int]				,   **a:    Optional[Any])  ->Any:






								'''simple docstring'''


								a_				     =   self.model_type.from_pretrained(*lowercase_				,   **lowercase_)



		def _lowerCAmelCase     (    self:    Optional[int])  ->Any:






								'''simple docstring'''


								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 _lowerCAmelCase     (    self:    List[str])  ->Optional[int]:






								'''simple docstring'''


								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)],
								        '''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(
														    lowercase_				,   lr=self.hparams.learning_rate				,   scale_parameter=lowercase_				,   relative_step=lowercase_)

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

								a_				     =   self.get_lr_scheduler()

								return [optimizer], [scheduler]



		def _lowerCAmelCase     (    self:    List[str]				,   a:    int				,   a:    str)  ->Any:






								'''simple docstring'''


								return self.validation_step(lowercase_				,   lowercase_)



		def _lowerCAmelCase     (    self:    int				,   a:    Dict)  ->Optional[Any]:






								'''simple docstring'''


								return self.validation_end(lowercase_)



		def _lowerCAmelCase     (    self:    Tuple)  ->int:






								'''simple docstring'''


								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 _lowerCAmelCase     (    self:    str				,   a:    Tuple)  ->int:






								'''simple docstring'''


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



		def _lowerCAmelCase     (    self:    Union[str, Any]				,   a:    Tuple				,   a:    Any				,   a:    Union[str, Any] = False)  ->Union[str, Any]:






								'''simple docstring'''


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



		def _lowerCAmelCase     (    self:    Tuple)  ->str:






								'''simple docstring'''


								return self.train_loader



		def _lowerCAmelCase     (    self:    List[Any])  ->Dict:






								'''simple docstring'''


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



		def _lowerCAmelCase     (    self:    List[str])  ->Optional[Any]:






								'''simple docstring'''


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



		def _lowerCAmelCase     (    self:    str				,   a:    Dict)  ->Dict:






								'''simple docstring'''


								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 _lowerCAmelCase     (    self:    List[Any]				,   a:    List[Any])  ->None:






								'''simple docstring'''


								a_				     =   self.output_dir.joinpath("best_tfmr")
								a_				     =   self.step_count
								self.model.save_pretrained(lowercase_)
								self.tokenizer.save_pretrained(lowercase_)




		@staticmethod
		def _lowerCAmelCase     (    a:    str				,   a:    Any)  ->Dict:






								'''simple docstring'''


								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=""				,   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 		SCREAMING_SNAKE_CASE__   (     pl.Callback    ):



		def _lowerCAmelCase     (    self:    Any				,   a:    List[str]				,   a:    Dict)  ->List[Any]:






								'''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 		SCREAMING_SNAKE_CASE__   (     pl.Callback    ):



		def _lowerCAmelCase     (    self:    str				,   a:    Dict				,   a:    Optional[int])  ->Dict:






								'''simple docstring'''


								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(lowercase_)



		def _lowerCAmelCase     (    self:    int				,   a:    List[str]				,   a:    List[Any])  ->Optional[int]:






								'''simple docstring'''


								rank_zero_info("***** Validation results *****")
								a_				     =   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 _lowerCAmelCase     (    self:    Tuple				,   a:    Optional[int]				,   a:    Union[str, Any])  ->Optional[Any]:






								'''simple docstring'''


								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(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 			__UpperCAmelCase				(lowercase__	,lowercase__				)		->				None:

						'''simple docstring'''


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



def 			__UpperCAmelCase				(lowercase__	,lowercase__	,lowercase__=None	,lowercase__=True	,lowercase__=None	,lowercase__=[]	,lowercase__=None	,lowercase__=None	,**lowercase__	,)		->				Optional[Any]:

						'''simple docstring'''


						pl.seed_everything(args.seed				)

						# init model
						a_				     =   Path(model.hparams.output_dir				)
						odir.mkdir(exist_ok=lowercase_				)

						# 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(lowercase_				)
						if logging_callback is None:
												a_				     =   LoggingCallback()

						a_				     =   {}

						# TODO: remove with PyTorch 1.6 since pl uses native amp
						if args.fpaa:
												a_				     =   16
												# train_params["amp_level"] = args.fp16_opt_level

						if args.gpus > 1:
												a_				     =   '''auto'''  # "ddp"
												a_				     =   '''ddp'''

						a_				     =   args.accumulate_grad_batches
						a_				     =   None  # extra_train_kwargs.get("profiler", None) #get unwanted logs
						a_				     =   '''auto'''

						a_				     =   pl.Trainer.from_argparse_args(
						    lowercase_	,weights_summary=lowercase_	,callbacks=[logging_callback] + extra_callbacks + [checkpoint_callback] + [InitCallback()]	,logger=lowercase_	,**lowercase_	,)

						if args.do_train:
												trainer.fit(lowercase_				)

						return trainer














'''simple docstring'''





import argparse
import collections
from pathlib import Path

import torch
from torch.serialization import default_restore_location

from transformers import BertConfig, DPRConfig, DPRContextEncoder, DPRQuestionEncoder, DPRReader


a_      = collections.namedtuple(
    'CheckpointState', ['model_dict', 'optimizer_dict', 'scheduler_dict', 'offset', 'epoch', 'encoder_params']
)



def 			__UpperCAmelCase				(lowercase__				)		->				CheckpointState:

						'''simple docstring'''


						print(F"""Reading saved model from {model_file}"""				)
						a_				     =   torch.load(lowercase__	,map_location=lambda lowercase__	,lowercase__				: default_restore_location(lowercase__	,"cpu"				)				)
						return CheckpointState(**lowercase__				)





class 		SCREAMING_SNAKE_CASE__   :



		def __init__(    self:    int				,   a:    Path)  ->str:






								'''simple docstring'''


								a_				     =   src_file



		def _lowerCAmelCase     (    self:    Dict)  ->Optional[Any]:






								'''simple docstring'''


								raise NotImplementedError




		@staticmethod
		def _lowerCAmelCase     (    a:    str				,   *a:    int				,   **a:    int)  ->"DPRState":






								'''simple docstring'''


								if comp_type.startswith("c"):
														return DPRContextEncoderState(*a				,   **a)
								if comp_type.startswith("q"):
														return DPRQuestionEncoderState(*a				,   **a)
								if comp_type.startswith("r"):
														return DPRReaderState(*a				,   **a)
								else:
														raise ValueError("Component type must be either 'ctx_encoder', 'question_encoder' or 'reader'.")





class 		SCREAMING_SNAKE_CASE__   (     lowercase_    ):



		def _lowerCAmelCase     (    self:    Union[str, Any])  ->Any:






								'''simple docstring'''


								a_				     =   DPRContextEncoder(DPRConfig(**BertConfig.get_config_dict("bert-base-uncased")[0]))
								print(f"""Loading DPR biencoder from {self.src_file}""")
								a_				     =   load_states_from_checkpoint(self.src_file)
								a_     ,       a_				     =   model.ctx_encoder, "ctx_model."
								# Fix changes from https://github.com/huggingface/transformers/commit/614fef1691edb806de976756d4948ecbcd0c0ca3
								a_				     =   {"bert_model.embeddings.position_ids": model.ctx_encoder.bert_model.embeddings.position_ids}
								for key, value in saved_state.model_dict.items():
														if key.startswith(a):
																				a_				     =   key[len(a) :]
																				if not key.startswith("encode_proj."):
																										a_				     =   "bert_model." + key
																				a_				     =   value
								encoder.load_state_dict(a)
								return model





class 		SCREAMING_SNAKE_CASE__   (     lowercase_    ):



		def _lowerCAmelCase     (    self:    Optional[int])  ->Optional[int]:






								'''simple docstring'''


								a_				     =   DPRQuestionEncoder(DPRConfig(**BertConfig.get_config_dict("bert-base-uncased")[0]))
								print(f"""Loading DPR biencoder from {self.src_file}""")
								a_				     =   load_states_from_checkpoint(self.src_file)
								a_     ,       a_				     =   model.question_encoder, "question_model."
								# Fix changes from https://github.com/huggingface/transformers/commit/614fef1691edb806de976756d4948ecbcd0c0ca3
								a_				     =   {"bert_model.embeddings.position_ids": model.question_encoder.bert_model.embeddings.position_ids}
								for key, value in saved_state.model_dict.items():
														if key.startswith(a):
																				a_				     =   key[len(a) :]
																				if not key.startswith("encode_proj."):
																										a_				     =   "bert_model." + key
																				a_				     =   value
								encoder.load_state_dict(a)
								return model





class 		SCREAMING_SNAKE_CASE__   (     lowercase_    ):



		def _lowerCAmelCase     (    self:    Tuple)  ->Optional[Any]:






								'''simple docstring'''


								a_				     =   DPRReader(DPRConfig(**BertConfig.get_config_dict("bert-base-uncased")[0]))
								print(f"""Loading DPR reader from {self.src_file}""")
								a_				     =   load_states_from_checkpoint(self.src_file)
								# Fix changes from https://github.com/huggingface/transformers/commit/614fef1691edb806de976756d4948ecbcd0c0ca3
								a_				     =   {
								    "encoder.bert_model.embeddings.position_ids": model.span_predictor.encoder.bert_model.embeddings.position_ids
								}
								for key, value in saved_state.model_dict.items():
														if key.startswith("encoder.") and not key.startswith("encoder.encode_proj"):
																				a_				     =   "encoder.bert_model." + key[len("encoder.") :]
														a_				     =   value
								model.span_predictor.load_state_dict(a)
								return model





def 			__UpperCAmelCase				(lowercase__	,lowercase__	,lowercase__				)		->				Any:

						'''simple docstring'''


						a_				     =   Path(lowercase__				)
						dest_dir.mkdir(exist_ok=lowercase__				)

						a_				     =   DPRState.from_type(lowercase__	,src_file=lowercase__				)
						a_				     =   dpr_state.load_dpr_model()
						model.save_pretrained(lowercase__				)
						model.from_pretrained(lowercase__				)  # sanity check


if __name__ == "__main__":
							a_      = argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    '--type', type=str, help='Type of the component to convert: \'ctx_encoder\', \'question_encoder\' or \'reader\'.'
							)
							parser.add_argument(
							    '--src',
							    type=str,
							    help=(
							        'Path to the dpr checkpoint file. They can be downloaded from the official DPR repo'
							        ' https://github.com/facebookresearch/DPR. Note that in the official repo, both encoders are stored in the'
							        ' \'retriever\' checkpoints.'
							    ),
							)
							parser.add_argument('--dest', type=str, default=None, help='Path to the output PyTorch model directory.')
							a_      = parser.parse_args()

							a_      = Path(args.src)
							a_      = F'converted-{src_file.name}' if args.dest is None else args.dest
							a_      = Path(dest_dir)
							assert src_file.exists()
							assert (
							    args.type is not None
							), "Please specify the component type of the DPR model to convert: 'ctx_encoder', 'question_encoder' or 'reader'."
							convert(args.type, src_file, dest_dir)



