




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


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



					"""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(SCREAMING_SNAKE_CASE							)
					snake_case						:  List[Any]        =			0
					snake_case						:  List[str]        =			Path(self.hparams.output_dir							)
					snake_case						:  str        =			self.hparams.cache_dir if self.hparams.cache_dir else None
					if config is None:
							snake_case						:  List[Any]        =			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=SCREAMING_SNAKE_CASE			,	**SCREAMING_SNAKE_CASE			,	)
					else:
							snake_case						:  PretrainedConfig        =			config

					snake_case						:  Optional[int]        =			("encoder_layerdrop", "decoder_layerdrop", "dropout", "attention_dropout")
					for p in extra_model_params:
							if getattr(self.hparams			,	SCREAMING_SNAKE_CASE			,	SCREAMING_SNAKE_CASE							):
									assert hasattr(self.config			,	SCREAMING_SNAKE_CASE							), F'''model config doesn\'t have a `{p}` attribute'''
									setattr(self.config			,	SCREAMING_SNAKE_CASE			,	getattr(self.hparams			,	SCREAMING_SNAKE_CASE							)							)

					if tokenizer is None:
							snake_case						:  Optional[int]        =			AutoTokenizer.from_pretrained(
							    self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path			,	cache_dir=SCREAMING_SNAKE_CASE			,	)
					else:
							snake_case						:  PreTrainedTokenizer        =			tokenizer
					snake_case						:  Optional[int]        =			MODEL_MODES[mode]
					if model is None:
							snake_case						:  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=SCREAMING_SNAKE_CASE			,	)
					else:
							snake_case						:  Any        =			model


			def 							lowerCamelCase_			(       self			,	*SCREAMING_SNAKE_CASE			,	**SCREAMING_SNAKE_CASE							):



					"""simple docstring"""







					snake_case						:  List[Any]        =			self.model_type.from_pretrained(*SCREAMING_SNAKE_CASE			,	**SCREAMING_SNAKE_CASE							)


			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







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


			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







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

					else:
							snake_case						:  Tuple        =			AdamW(
							    SCREAMING_SNAKE_CASE			,	lr=self.hparams.learning_rate			,	eps=self.hparams.adam_epsilon							)
					snake_case						:  Optional[int]        =			optimizer

					snake_case						:  Dict        =			self.get_lr_scheduler()

					return [optimizer], [scheduler]


			def 							lowerCamelCase_			(       self			,	SCREAMING_SNAKE_CASE			,	SCREAMING_SNAKE_CASE							):



					"""simple docstring"""







					return self.validation_step(SCREAMING_SNAKE_CASE			,	SCREAMING_SNAKE_CASE							)


			def 							lowerCamelCase_			(       self			,	SCREAMING_SNAKE_CASE							):



					"""simple docstring"""







					return self.validation_end(SCREAMING_SNAKE_CASE							)


			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







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



					"""simple docstring"""







					if mode == "test":
							snake_case						:  Optional[int]        =			len(self.test_dataloader().dataset							)
					else:
							snake_case						:  Any        =			self.get_dataloader("train"			,	self.hparams.train_batch_size			,	shuffle=SCREAMING_SNAKE_CASE							)
							snake_case						:  Any        =			len(self.train_dataloader().dataset							)


			def 							lowerCamelCase_			(       self			,	SCREAMING_SNAKE_CASE			,	SCREAMING_SNAKE_CASE			,	SCREAMING_SNAKE_CASE = False							):



					"""simple docstring"""







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


			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







					return self.train_loader


			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







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


			def 							lowerCamelCase_			(       self							):



					"""simple docstring"""







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


			def 							lowerCamelCase_			(       self			,	SCREAMING_SNAKE_CASE							):



					"""simple docstring"""







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


			@pl.utilities.rank_zero_only
			def 							lowerCamelCase_			(       self			,	SCREAMING_SNAKE_CASE							):



					"""simple docstring"""







					snake_case						:  str        =			self.output_dir.joinpath("best_tfmr"							)
					snake_case						:  Optional[int]        =			self.step_count
					self.model.save_pretrained(SCREAMING_SNAKE_CASE							)
					self.tokenizer.save_pretrained(SCREAMING_SNAKE_CASE							)



			@staticmethod
			def 							lowerCamelCase_			(       SCREAMING_SNAKE_CASE			,	SCREAMING_SNAKE_CASE							):



					"""simple docstring"""







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







class      lowerCamelCase__      (					pl.Callback     ):


			def 							lowerCamelCase_			(       self			,	SCREAMING_SNAKE_CASE			,	SCREAMING_SNAKE_CASE							):



					"""simple docstring"""







					snake_case						:  Union[str, Any]        =			trainer.lr_schedulers[0]["scheduler"]
					snake_case						:  Optional[Any]        =			{F'''lr_group_{i}''': lr for i, lr in enumerate(lr_scheduler.get_lr()							)}
					pl_module.logger.log_metrics(SCREAMING_SNAKE_CASE							)


			def 							lowerCamelCase_			(       self			,	SCREAMING_SNAKE_CASE			,	SCREAMING_SNAKE_CASE							):



					"""simple docstring"""







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



			def 							lowerCamelCase_			(       self			,	SCREAMING_SNAKE_CASE			,	SCREAMING_SNAKE_CASE							):



					"""simple docstring"""







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







def 					UpperCamelCase__       (						lowercase__					:       Optional[Any]    ,       lowercase__					:       List[str]				):
		#  To allow all pl args uncomment the following line
		#  parser = pl.Trainer.add_argparse_args(parser)
		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__					:       BaseTransformer    ,       lowercase__					:       argparse.Namespace    ,       lowercase__					:       Tuple=None    ,       lowercase__					:       List[str]=True    ,       lowercase__					:       List[str]=[]    ,       lowercase__					:       Any=None    ,       lowercase__					:       List[str]=None    ,       **lowercase__					:       List[str]    ,       ):
		pl.seed_everything(args.seed				)

		# init model
		snake_case						:  Union[str, Any]        =			Path(model.hparams.output_dir				)
		odir.mkdir(exist_ok=lowercase__				)

		# add custom checkpoints
		if checkpoint_callback is None:
				snake_case						:  str        =			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:
				snake_case						:  Optional[int]        =			LoggingCallback()

		snake_case						:  List[str]        =			{}

		# TODO: remove with PyTorch 1.6 since pl uses native amp
		if args.fpaa:
				snake_case						:  Tuple        =			16
				snake_case						:  List[str]        =			args.fpaa_opt_level

		if args.gpus > 1:
				snake_case						:  Any        =			"ddp"

		snake_case						:  Union[str, Any]        =			args.accumulate_grad_batches
		snake_case						:  Optional[int]        =			extra_train_kwargs.get("accelerator"    ,       lowercase__				)
		snake_case						:  Optional[Any]        =			extra_train_kwargs.get("profiler"    ,       lowercase__				)

		snake_case						:  Dict        =			pl.Trainer.from_argparse_args(
		    lowercase__    ,       weights_summary=lowercase__    ,       callbacks=[logging_callback] + extra_callbacks    ,       logger=lowercase__    ,       checkpoint_callback=lowercase__    ,       **lowercase__    ,       )

		if args.do_train:
				trainer.fit(lowercase__				)

		return trainer







"""simple docstring"""

from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
)


__A							=     {"configuration_gptj": ["GPTJ_PRETRAINED_CONFIG_ARCHIVE_MAP", "GPTJConfig", "GPTJOnnxConfig"]}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							__A							=     [
							    "GPTJ_PRETRAINED_MODEL_ARCHIVE_LIST",
							    "GPTJForCausalLM",
							    "GPTJForQuestionAnswering",
							    "GPTJForSequenceClassification",
							    "GPTJModel",
							    "GPTJPreTrainedModel",
							]

try:
							if not is_tf_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							__A							=     [
							    "TFGPTJForCausalLM",
							    "TFGPTJForQuestionAnswering",
							    "TFGPTJForSequenceClassification",
							    "TFGPTJModel",
							    "TFGPTJPreTrainedModel",
							]

try:
							if not is_flax_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							__A							=     [
							    "FlaxGPTJForCausalLM",
							    "FlaxGPTJModel",
							    "FlaxGPTJPreTrainedModel",
							]


if TYPE_CHECKING:
							from .configuration_gptj import GPTJ_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTJConfig, GPTJOnnxConfig

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_gptj import (
														    GPTJ_PRETRAINED_MODEL_ARCHIVE_LIST,
														    GPTJForCausalLM,
														    GPTJForQuestionAnswering,
														    GPTJForSequenceClassification,
														    GPTJModel,
														    GPTJPreTrainedModel,
														)

							try:
														if not is_tf_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_tf_gptj import (
														    TFGPTJForCausalLM,
														    TFGPTJForQuestionAnswering,
														    TFGPTJForSequenceClassification,
														    TFGPTJModel,
														    TFGPTJPreTrainedModel,
														)

							try:
														if not is_flax_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_flax_gptj import FlaxGPTJForCausalLM, FlaxGPTJModel, FlaxGPTJPreTrainedModel

else:
							import sys

							__A							=     _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)

