




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


snake_case__          =						logging.getLogger(__name__)

require_version("""pytorch_lightning>=1.0.4""")

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






class     UpperCAmelCase_     (pl.LightningModule     ):







						"""simple docstring"""


						def __init__(     self      :       Optional[Any]							, _lowerCamelCase      :       argparse.Namespace							, _lowerCamelCase      :       Tuple=None							, _lowerCamelCase      :       Union[str, Any]="base"							, _lowerCamelCase      :       Any=None							, _lowerCamelCase      :       List[Any]=None							, _lowerCamelCase      :       Union[str, Any]=None							, **_lowerCamelCase      :       Any							, ):





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

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

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







						def       _a    (     self      :       Dict							, *_lowerCamelCase      :       Any							, **_lowerCamelCase      :       Optional[Any]    ):





											"""simple docstring"""







											A_ :     Optional[int] =							self.model_type.from_pretrained(*_lowerCamelCase							, **_lowerCamelCase    )







						def       _a    (     self      :       Tuple    ):





											"""simple docstring"""







											A_ :     Optional[int] =							arg_to_scheduler[self.hparams.lr_scheduler]
											A_ :     Optional[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      :       List[str]    ):





											"""simple docstring"""







											A_ :     Optional[int] =							self.model
											A_ :     Any =							["bias", "LayerNorm.weight"]
											A_ :     Optional[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:
																A_ :     Optional[int] =							Adafactor(
																    _lowerCamelCase							, lr=self.hparams.learning_rate							, scale_parameter=_lowerCamelCase							, relative_step=_lowerCamelCase    )

											else:
																A_ :     List[str] =							AdamW(
																    _lowerCamelCase							, lr=self.hparams.learning_rate							, eps=self.hparams.adam_epsilon    )
											A_ :     List[str] =							optimizer

											A_ :     Tuple =							self.get_lr_scheduler()

											return [optimizer], [scheduler]







						def       _a    (     self      :       Optional[Any]							, _lowerCamelCase      :       Union[str, Any]							, _lowerCamelCase      :       int    ):





											"""simple docstring"""







											return self.validation_step(_lowerCamelCase							, _lowerCamelCase    )







						def       _a    (     self      :       List[str]							, _lowerCamelCase      :       Dict    ):





											"""simple docstring"""







											return self.validation_end(_lowerCamelCase    )







						def       _a    (     self      :       Optional[Any]    ):





											"""simple docstring"""







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





											"""simple docstring"""







											if mode == "test":
																A_ :     Union[str, Any] =							len(self.test_dataloader().dataset    )
											else:
																A_ :     str =							self.get_dataloader('''train'''							, self.hparams.train_batch_size							, shuffle=_lowerCamelCase    )
																A_ :     Optional[int] =							len(self.train_dataloader().dataset    )







						def       _a    (     self      :       Optional[Any]							, _lowerCamelCase      :       str							, _lowerCamelCase      :       int							, _lowerCamelCase      :       bool = False    ):





											"""simple docstring"""







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







						def       _a    (     self      :       Optional[int]    ):





											"""simple docstring"""







											return self.train_loader







						def       _a    (     self      :       Optional[int]    ):





											"""simple docstring"""







											return self.get_dataloader('''dev'''							, self.hparams.eval_batch_size							, shuffle=_lowerCamelCase    )







						def       _a    (     self      :       Any    ):





											"""simple docstring"""







											return self.get_dataloader('''test'''							, self.hparams.eval_batch_size							, shuffle=_lowerCamelCase    )







						def       _a    (     self      :       Union[str, Any]							, _lowerCamelCase      :       Tuple    ):





											"""simple docstring"""







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







						@pl.utilities.rank_zero_only
						def       _a    (     self      :       Any							, _lowerCamelCase      :       Dict[str, Any]    ):





											"""simple docstring"""







											A_ :     List[str] =							self.output_dir.joinpath('''best_tfmr'''    )
											A_ :     List[str] =							self.step_count
											self.model.save_pretrained(_lowerCamelCase    )
											self.tokenizer.save_pretrained(_lowerCamelCase    )







						@staticmethod
						def       _a    (     _lowerCamelCase      :       Optional[int]							, _lowerCamelCase      :       Dict    ):





											"""simple docstring"""







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






class     UpperCAmelCase_     (pl.Callback     ):







						"""simple docstring"""


						def       _a    (     self      :       int							, _lowerCamelCase      :       Dict							, _lowerCamelCase      :       Optional[int]    ):





											"""simple docstring"""







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







						def       _a    (     self      :       Union[str, Any]							, _lowerCamelCase      :       pl.Trainer							, _lowerCamelCase      :       pl.LightningModule    ):





											"""simple docstring"""







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







						def       _a    (     self      :       Tuple							, _lowerCamelCase      :       pl.Trainer							, _lowerCamelCase      :       pl.LightningModule    ):





											"""simple docstring"""







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






def 						snake_case__    ( lowerCamelCase__				:   str			,		lowerCamelCase__				:   str			)   ->   None:
					#  To allow all pl args uncomment the following line
					#  parser = pl.Trainer.add_argparse_args(parser)
					parser.add_argument(
					    '''--output_dir'''			,		default=__lowerCAmelCase			,		type=__lowerCAmelCase			,		required=__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=4_2			,		help='''random seed for initialization'''			)
					parser.add_argument(
					    '''--data_dir'''			,		default=__lowerCAmelCase			,		type=__lowerCAmelCase			,		required=__lowerCAmelCase			,		help='''The input data dir. Should contain the training files for the CoNLL-2003 NER task.'''			,		)





def 						snake_case__    ( lowerCamelCase__				:   str			,		lowerCamelCase__				:   Union[str, Any]			,		lowerCamelCase__				:   List[str]=None			,		lowerCamelCase__				:   Optional[Any]=True			,		lowerCamelCase__				:   List[str]=[]			,		lowerCamelCase__				:   Optional[Any]=None			,		lowerCamelCase__				:   List[str]=None			,		**lowerCamelCase__				:   Tuple			,		)   ->   str:
					pl.seed_everything(args.seed			)

					# init model
					A_ :     List[Any] =							Path(model.hparams.output_dir			)
					odir.mkdir(exist_ok=__lowerCAmelCase			)

					# add custom checkpoints
					if checkpoint_callback is None:
										A_ :     List[Any] =							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_ :     Tuple =							{}

					# TODO: remove with PyTorch 1.6 since pl uses native amp
					if args.fpaa:
										A_ :     Optional[int] =							1_6
										A_ :     List[str] =							args.fpaa_opt_level

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

					A_ :     int =							args.accumulate_grad_batches
					A_ :     List[Any] =							extra_train_kwargs.get('''accelerator'''			,		__lowerCAmelCase			)
					A_ :     Union[str, Any] =							extra_train_kwargs.get('''profiler'''			,		__lowerCAmelCase			)

					A_ :     Optional[int] =							pl.Trainer.from_argparse_args(
					    __lowerCAmelCase			,		weights_summary=__lowerCAmelCase			,		callbacks=[logging_callback] + extra_callbacks			,		logger=__lowerCAmelCase			,		checkpoint_callback=__lowerCAmelCase			,		**__lowerCAmelCase			,		)

					if args.do_train:
										trainer.fit(__lowerCAmelCase			)

					return trainer









'''simple docstring'''




from ...configuration_utils import PretrainedConfig
from ...utils import logging


snake_case__          =						logging.get_logger(__name__)

snake_case__          =						{
    """bert-base-uncased""": """https://huggingface.co/bert-base-uncased/resolve/main/config.json""",
    # QDQBERT models can be loaded from any BERT checkpoint, available at https://huggingface.co/models?filter=bert
}






class     UpperCamelCase_     (a__     ):







						"""simple docstring"""


						_lowerCAmelCase    						=					'qdqbert'







						def __init__(     self      :       Union[str, Any]							, _lowerCamelCase      :       List[Any]=30522							, _lowerCamelCase      :       str=768							, _lowerCamelCase      :       Optional[Any]=12							, _lowerCamelCase      :       str=12							, _lowerCamelCase      :       List[Any]=3072							, _lowerCamelCase      :       Union[str, Any]="gelu"							, _lowerCamelCase      :       Dict=0.1							, _lowerCamelCase      :       Union[str, Any]=0.1							, _lowerCamelCase      :       Optional[int]=512							, _lowerCamelCase      :       Dict=2							, _lowerCamelCase      :       Optional[Any]=0.02							, _lowerCamelCase      :       List[Any]=1E-12							, _lowerCamelCase      :       Tuple=True							, _lowerCamelCase      :       Tuple=1							, _lowerCamelCase      :       Optional[Any]=0							, _lowerCamelCase      :       Optional[int]=2							, **_lowerCamelCase      :       List[str]							, ):





											"""simple docstring"""







											super().__init__(pad_token_id=_lowerCamelCase							, bos_token_id=_lowerCamelCase							, eos_token_id=_lowerCamelCase							, **_lowerCamelCase    )

											A_ :     Optional[Any] =							vocab_size
											A_ :     str =							max_position_embeddings
											A_ :     List[str] =							hidden_size
											A_ :     str =							num_hidden_layers
											A_ :     List[str] =							num_attention_heads
											A_ :     List[str] =							intermediate_size
											A_ :     int =							hidden_act
											A_ :     Tuple =							hidden_dropout_prob
											A_ :     Optional[Any] =							attention_probs_dropout_prob
											A_ :     Dict =							initializer_range
											A_ :     Optional[int] =							type_vocab_size
											A_ :     Any =							layer_norm_eps
											A_ :     Optional[int] =							use_cache



