





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 					UpperCamelCase_    (       pl.LightningModule      ):


							def __init__(							self      :Optional[Any]							,			__A      :argparse.Namespace							,			__A      :Any=None							,			__A      :Optional[int]="base"							,			__A      :int=None							,			__A      :Optional[int]=None							,			__A      :Union[str, Any]=None							,			**__A      :Optional[int]							,			)				->				str:



									"""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						)
									SCREAMING_SNAKE_CASE__								= 0
									SCREAMING_SNAKE_CASE__								= Path(self.hparams.output_dir						)
									SCREAMING_SNAKE_CASE__								= self.hparams.cache_dir if self.hparams.cache_dir else None
									if config is None:
											SCREAMING_SNAKE_CASE__								= 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:
											SCREAMING_SNAKE_CASE__								= config

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


							def        _snake_case					(							self      :str							,			*__A      :Any							,			**__A      :List[Any]						)				->				Any:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								= self.model_type.from_pretrained(*__A							,			**__A						)


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



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								= arg_to_scheduler[self.hparams.lr_scheduler]
									SCREAMING_SNAKE_CASE__								= get_schedule_func(
									    self.opt							,			num_warmup_steps=self.hparams.warmup_steps							,			num_training_steps=self.total_steps()						)
									SCREAMING_SNAKE_CASE__								= {"""scheduler""": scheduler, """interval""": """step""", """frequency""": 1}
									return scheduler


							def        _snake_case					(							self      :Tuple						)				->				List[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								= self.model
									SCREAMING_SNAKE_CASE__								= ["""bias""", """LayerNorm.weight"""]
									SCREAMING_SNAKE_CASE__								= [
									    {
									        """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:
											SCREAMING_SNAKE_CASE__								= Adafactor(
											    __A							,			lr=self.hparams.learning_rate							,			scale_parameter=__A							,			relative_step=__A						)

									else:
											SCREAMING_SNAKE_CASE__								= AdamW(
											    __A							,			lr=self.hparams.learning_rate							,			eps=self.hparams.adam_epsilon						)
									SCREAMING_SNAKE_CASE__								= optimizer

									SCREAMING_SNAKE_CASE__								= self.get_lr_scheduler()

									return [optimizer], [scheduler]


							def        _snake_case					(							self      :str							,			__A      :Tuple							,			__A      :Tuple						)				->				Union[str, Any]:



									"""simple docstring"""



									return self.validation_step(__A							,			__A						)


							def        _snake_case					(							self      :Tuple							,			__A      :Optional[int]						)				->				Any:



									"""simple docstring"""



									return self.validation_end(__A						)


							def        _snake_case					(							self      :Dict						)				->				int:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								= max(1							,			self.hparams.gpus						)  # TODO: consider num_tpu_cores
									SCREAMING_SNAKE_CASE__								= 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      :Dict							,			__A      :List[str]						)				->				List[Any]:



									"""simple docstring"""



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


							def        _snake_case					(							self      :Optional[Any]							,			__A      :str							,			__A      :int							,			__A      :bool = False						)				->				Union[str, Any]:



									"""simple docstring"""



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


							def        _snake_case					(							self      :Any						)				->				Dict:



									"""simple docstring"""



									return self.train_loader


							def        _snake_case					(							self      :Any						)				->				int:



									"""simple docstring"""



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


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



									"""simple docstring"""



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


							def        _snake_case					(							self      :Any							,			__A      :Optional[int]						)				->				Any:



									"""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        _snake_case					(							self      :List[Any]							,			__A      :Dict[str, Any]						)				->				None:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								= self.output_dir.joinpath("""best_tfmr"""						)
									SCREAMING_SNAKE_CASE__								= self.step_count
									self.model.save_pretrained(__A						)
									self.tokenizer.save_pretrained(__A						)


							@staticmethod
							def        _snake_case					(							__A      :Union[str, Any]							,			__A      :Any						)				->				List[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=""""""							,			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        _snake_case					(							self      :Optional[Any]							,			__A      :Tuple							,			__A      :int						)				->				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 					UpperCamelCase_    (       pl.Callback      ):


							def        _snake_case					(							self      :Dict							,			__A      :Optional[Any]							,			__A      :Union[str, Any]						)				->				Dict:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								= trainer.lr_schedulers[0]["""scheduler"""]
									SCREAMING_SNAKE_CASE__								= {f'''lr_group_{i}''': lr for i, lr in enumerate(lr_scheduler.get_lr()						)}
									pl_module.logger.log_metrics(__A						)


							def        _snake_case					(							self      :Tuple							,			__A      :pl.Trainer							,			__A      :pl.LightningModule						)				->				Any:



									"""simple docstring"""



									rank_zero_info("""***** Validation results *****"""						)
									SCREAMING_SNAKE_CASE__								= 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      :int							,			__A      :pl.Trainer							,			__A      :pl.LightningModule						)				->				List[Any]:



									"""simple docstring"""



									rank_zero_info("""***** Test results *****"""						)
									SCREAMING_SNAKE_CASE__								= trainer.callback_metrics
									# Log and save results to file
									SCREAMING_SNAKE_CASE__								= 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 			SCREAMING_SNAKE_CASE__							(    UpperCamelCase__:							List[Any]   ,				UpperCamelCase__:							List[str]      ):
		#  To allow all pl args uncomment the following line
		#  parser = pl.Trainer.add_argparse_args(parser)
		parser.add_argument(
		    """--output_dir"""   ,				default=UpperCamelCase__   ,				type=UpperCamelCase__   ,				required=UpperCamelCase__   ,				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=UpperCamelCase__   ,				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=UpperCamelCase__      )
		parser.add_argument("""--max_grad_norm"""   ,				dest="""gradient_clip_val"""   ,				default=1.0   ,				type=UpperCamelCase__   ,				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=UpperCamelCase__   ,				default=1   ,				help="""Number of updates steps to accumulate before performing a backward/update pass."""   ,				)
		parser.add_argument("""--seed"""   ,				type=UpperCamelCase__   ,				default=42   ,				help="""random seed for initialization"""      )
		parser.add_argument(
		    """--data_dir"""   ,				default=UpperCamelCase__   ,				type=UpperCamelCase__   ,				required=UpperCamelCase__   ,				help="""The input data dir. Should contain the training files for the CoNLL-2003 NER task."""   ,				)






def 			SCREAMING_SNAKE_CASE__							(    UpperCamelCase__:							BaseTransformer   ,				UpperCamelCase__:							argparse.Namespace   ,				UpperCamelCase__:							Dict=None   ,				UpperCamelCase__:							int=True   ,				UpperCamelCase__:							Dict=None   ,				UpperCamelCase__:							List[Any]=[]   ,				UpperCamelCase__:							Optional[Any]=None   ,				UpperCamelCase__:							str=None   ,				**UpperCamelCase__:							List[Any]   ,				):
		pl.seed_everything(args.seed      )

		# init model
		SCREAMING_SNAKE_CASE__								= Path(model.hparams.output_dir      )
		odir.mkdir(exist_ok=UpperCamelCase__      )

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

		SCREAMING_SNAKE_CASE__								= {}

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

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

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

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

		if args.do_train:
				trainer.fit(UpperCamelCase__      )

		return trainer






import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


_lowerCamelCase							=     500000

_lowerCamelCase			,    _lowerCamelCase							=     os.path.split(__file__)
_lowerCamelCase							=     os.path.join(RESULTS_BASEPATH, 'results', RESULTS_FILENAME.replace('.py', '.json'))






@get_duration
def 			SCREAMING_SNAKE_CASE__							(    UpperCamelCase__:							datasets.Dataset      ):
		SCREAMING_SNAKE_CASE__								= dataset.select(range(0   ,				len(UpperCamelCase__      )   ,				2      )      )






@get_duration
def 			SCREAMING_SNAKE_CASE__							(    UpperCamelCase__:							datasets.Dataset      ):
		SCREAMING_SNAKE_CASE__								= dataset.sort("""numbers"""      )






@get_duration
def 			SCREAMING_SNAKE_CASE__							(    UpperCamelCase__:							datasets.Dataset      ):
		SCREAMING_SNAKE_CASE__								= dataset.shuffle()






@get_duration
def 			SCREAMING_SNAKE_CASE__							(    UpperCamelCase__:							datasets.Dataset      ):
		SCREAMING_SNAKE_CASE__								= dataset.train_test_split(0.1      )






@get_duration
def 			SCREAMING_SNAKE_CASE__							(    UpperCamelCase__:							datasets.Dataset   ,				UpperCamelCase__:							Tuple=10      ):
		for shard_id in range(UpperCamelCase__      ):
				SCREAMING_SNAKE_CASE__								= dataset.shard(UpperCamelCase__   ,				UpperCamelCase__      )






def 			SCREAMING_SNAKE_CASE__							(    ):
		SCREAMING_SNAKE_CASE__								= {"""num examples""": SPEED_TEST_N_EXAMPLES}
		SCREAMING_SNAKE_CASE__								= (select, sort, shuffle, train_test_split, shard)
		with tempfile.TemporaryDirectory() as tmp_dir:
				print("""generating dataset"""      )
				SCREAMING_SNAKE_CASE__								= datasets.Features({"""text""": datasets.Value("""string"""      ), """numbers""": datasets.Value("""float32"""      )}      )
				SCREAMING_SNAKE_CASE__								= generate_example_dataset(
				    os.path.join(UpperCamelCase__   ,				"""dataset.arrow"""      )   ,				UpperCamelCase__   ,				num_examples=UpperCamelCase__      )
				print("""Functions"""      )
				for func in functions:
						print(func.__name__      )
						SCREAMING_SNAKE_CASE__								= func(UpperCamelCase__      )

		with open(UpperCamelCase__   ,				"""wb"""      ) as f:
				f.write(json.dumps(UpperCamelCase__      ).encode("""utf-8"""      )      )


if __name__ == "__main__":  # useful to run the profiler
	benchmark_indices_mapping()