"""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									=	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		,__UpperCAmelCase		,__UpperCAmelCase=None		,__UpperCAmelCase="base"		,__UpperCAmelCase=None		,__UpperCAmelCase=None		,__UpperCAmelCase=None		,**__UpperCAmelCase		,)       ->					Union[str, Any]:
													super().__init__()
													# TODO: move to self.save_hyperparameters()
													# self.save_hyperparameters()
													# can also expand arguments into trainer signature for easier reading

													self.save_hyperparameters(__UpperCAmelCase  )
													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=__UpperCAmelCase		,**__UpperCAmelCase		,)
													else:
																				A__						  =				config

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

													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=__UpperCAmelCase		,)
													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=__UpperCAmelCase		,)
													else:
																				A__						  =				model





						def 					snake_case__						(							self		,*__UpperCAmelCase		,**__UpperCAmelCase  )       ->					Union[str, Any]:
													A__						  =				self.model_type.from_pretrained(*__UpperCAmelCase		,**__UpperCAmelCase  )





						def 					snake_case__						(							self  )       ->					List[str]:
													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 					snake_case__						(							self  )       ->					Tuple:
													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(
																				    __UpperCAmelCase		,lr=self.hparams.learning_rate		,scale_parameter=__UpperCAmelCase		,relative_step=__UpperCAmelCase  )

													else:
																				A__						  =				AdamW(
																				    __UpperCAmelCase		,lr=self.hparams.learning_rate		,eps=self.hparams.adam_epsilon  )
													A__						  =				optimizer

													A__						  =				self.get_lr_scheduler()

													return [optimizer], [scheduler]





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase  )       ->					str:
													return self.validation_step(__UpperCAmelCase		,__UpperCAmelCase  )





						def 					snake_case__						(							self		,__UpperCAmelCase  )       ->					Union[str, Any]:
													return self.validation_end(__UpperCAmelCase  )





						def 					snake_case__						(							self  )       ->					int:
													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 					snake_case__						(							self		,__UpperCAmelCase  )       ->					Optional[Any]:
													if mode == "test":
																				A__						  =				len(self.test_dataloader().dataset  )
													else:
																				A__						  =				self.get_dataloader('train'		,self.hparams.train_batch_size		,shuffle=__UpperCAmelCase  )
																				A__						  =				len(self.train_dataloader().dataset  )





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase = False  )       ->					List[str]:
													raise NotImplementedError('You must implement this for your task'  )





						def 					snake_case__						(							self  )       ->					List[Any]:
													return self.train_loader





						def 					snake_case__						(							self  )       ->					str:
													return self.get_dataloader('dev'		,self.hparams.eval_batch_size		,shuffle=__UpperCAmelCase  )





						def 					snake_case__						(							self  )       ->					Tuple:
													return self.get_dataloader('test'		,self.hparams.eval_batch_size		,shuffle=__UpperCAmelCase  )





						def 					snake_case__						(							self		,__UpperCAmelCase  )       ->					Tuple:
													return os.path.join(
													    self.hparams.data_dir		,'cached_{}_{}_{}'.format(
													        __UpperCAmelCase		,list(filter(__UpperCAmelCase		,self.hparams.model_name_or_path.split('/'  )  )  ).pop()		,str(self.hparams.max_seq_length  )		,)		,)





						@pl.utilities.rank_zero_only
						def 					snake_case__						(							self		,__UpperCAmelCase  )       ->					None:
													A__						  =				self.output_dir.joinpath('best_tfmr'  )
													A__						  =				self.step_count
													self.model.save_pretrained(__UpperCAmelCase  )
													self.tokenizer.save_pretrained(__UpperCAmelCase  )





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






class      UpperCamelCase__(      pl.Callback			):





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase  )       ->					Optional[Any]:
													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(__UpperCAmelCase  )





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase  )       ->					Union[str, Any]:
													rank_zero_info('***** Validation results *****'  )
													A__						  =				trainer.callback_metrics
													# Log results
													for key in sorted(__UpperCAmelCase  ):
																				if key not in ["log", "progress_bar"]:
																											rank_zero_info('{} = {}\n'.format(__UpperCAmelCase		,str(metrics[key]  )  )  )





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase  )       ->					Tuple:
													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(__UpperCAmelCase		,'w'  ) as writer:
																				for key in sorted(__UpperCAmelCase  ):
																											if key not in ["log", "progress_bar"]:
																																		rank_zero_info('{} = {}\n'.format(__UpperCAmelCase		,str(metrics[key]  )  )  )
																																		writer.write('{} = {}\n'.format(__UpperCAmelCase		,str(metrics[key]  )  )  )






def 	UpperCAmelCase       (						UpperCamelCase__							,     UpperCamelCase__					):






							"""simple docstring"""





							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 	UpperCAmelCase       (						UpperCamelCase__							,     UpperCamelCase__							,     UpperCamelCase__=None							,     UpperCamelCase__=True							,     UpperCamelCase__=[]							,     UpperCamelCase__=None							,     UpperCamelCase__=None							,     **UpperCamelCase__							,     ):






							"""simple docstring"""





							pl.seed_everything(args.seed					)

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

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

							A__						  =				{}

							# TODO: remove with PyTorch 1.6 since pl uses native amp
							if args.fpaa:
														A__						  =				16
														A__						  =				args.fpaa_opt_level

							if args.gpus > 1:
														A__						  =				'ddp'

							A__						  =				args.accumulate_grad_batches
							A__						  =				extra_train_kwargs.get('accelerator'							,     UpperCamelCase__					)
							A__						  =				extra_train_kwargs.get('profiler'							,     UpperCamelCase__					)

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

							if args.do_train:
														trainer.fit(UpperCamelCase__					)

							return trainer






"""simple docstring"""




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


__lowerCamelCase									=	logging.get_logger(__name__)

__lowerCamelCase									=	{
    "dandelin/vilt-b32-mlm": "https://huggingface.co/dandelin/vilt-b32-mlm/blob/main/config.json"
}
class      UpperCamelCase__(      __A			):
						lowerCAmelCase__	:    str    				=		'vilt'





						def __init__(							self		,__UpperCAmelCase=3_05_22		,__UpperCAmelCase=2		,__UpperCAmelCase=2		,__UpperCAmelCase=40		,__UpperCAmelCase=7_68		,__UpperCAmelCase=12		,__UpperCAmelCase=12		,__UpperCAmelCase=30_72		,__UpperCAmelCase="gelu"		,__UpperCAmelCase=0.0		,__UpperCAmelCase=0.0		,__UpperCAmelCase=0.0_2		,__UpperCAmelCase=1e-12		,__UpperCAmelCase=3_84		,__UpperCAmelCase=32		,__UpperCAmelCase=3		,__UpperCAmelCase=True		,__UpperCAmelCase=-1		,__UpperCAmelCase=False		,__UpperCAmelCase=-1		,**__UpperCAmelCase		,)       ->					Tuple:
													super().__init__(tie_word_embeddings=__UpperCAmelCase		,**__UpperCAmelCase  )

													A__						  =				vocab_size
													A__						  =				type_vocab_size
													A__						  =				modality_type_vocab_size
													A__						  =				max_position_embeddings

													A__						  =				hidden_size
													A__						  =				num_hidden_layers
													A__						  =				num_attention_heads
													A__						  =				intermediate_size
													A__						  =				hidden_act
													A__						  =				hidden_dropout_prob
													A__						  =				attention_probs_dropout_prob
													A__						  =				initializer_range
													A__						  =				layer_norm_eps

													A__						  =				image_size
													A__						  =				patch_size
													A__						  =				num_channels
													A__						  =				qkv_bias
													A__						  =				max_image_length
													A__						  =				num_images





