






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     :     Dict       =  logging.getLogger(__name__)

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

__snake_case     :     int       =  {
    '''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     :     Tuple       =  {
    '''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     :     Union[str, Any]       =  sorted(arg_to_scheduler.keys())
__snake_case     :     Any       =  '''{''' + ''', '''.join(arg_to_scheduler_choices) + '''}'''


class      __SCREAMING_SNAKE_CASE       ( pl.LightningModule):






						def __init__(						self  ,    _UpperCamelCase  ,    _UpperCamelCase=None  ,    _UpperCamelCase="base"  ,    _UpperCamelCase=None  ,    _UpperCamelCase=None  ,    _UpperCamelCase=None  ,    **_UpperCamelCase  ,    ):






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

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






						def    UpperCamelCase__   (						self  ,    *_UpperCamelCase  ,    **_UpperCamelCase ):






												"""simple docstring"""
												lowerCAmelCase__         =    self.model_type.from_pretrained(*_lowercase  ,    **_lowercase )






						def    UpperCamelCase__   (						self ):






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






						def    UpperCamelCase__   (						self ):






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

												else:
																		lowerCAmelCase__         =    AdamW(
																		    _lowercase  ,    lr=self.hparams.learning_rate  ,    eps=self.hparams.adam_epsilon )
												lowerCAmelCase__         =    optimizer

												lowerCAmelCase__         =    self.get_lr_scheduler()

												return [optimizer], [scheduler]






						def    UpperCamelCase__   (						self  ,    _UpperCamelCase  ,    _UpperCamelCase ):






												"""simple docstring"""
												return self.validation_step(_lowercase  ,    _lowercase )






						def    UpperCamelCase__   (						self  ,    _UpperCamelCase ):






												"""simple docstring"""
												return self.validation_end(_lowercase )






						def    UpperCamelCase__   (						self ):






												"""simple docstring"""
												lowerCAmelCase__         =    max(1  ,    self.hparams.gpus )  # TODO: consider num_tpu_cores
												lowerCAmelCase__         =    self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
												return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs






						def    UpperCamelCase__   (						self  ,    _UpperCamelCase ):






												"""simple docstring"""






												if mode == "test":
																		lowerCAmelCase__         =    len(self.test_dataloader().dataset )
												else:
																		lowerCAmelCase__         =    self.get_dataloader('train'  ,    self.hparams.train_batch_size  ,    shuffle=_lowercase )
																		lowerCAmelCase__         =    len(self.train_dataloader().dataset )






						def    UpperCamelCase__   (						self  ,    _UpperCamelCase  ,    _UpperCamelCase  ,    _UpperCamelCase = False ):






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






						def    UpperCamelCase__   (						self ):






												"""simple docstring"""
												return self.train_loader






						def    UpperCamelCase__   (						self ):






												"""simple docstring"""
												return self.get_dataloader('dev'  ,    self.hparams.eval_batch_size  ,    shuffle=_lowercase )






						def    UpperCamelCase__   (						self ):






												"""simple docstring"""
												return self.get_dataloader('test'  ,    self.hparams.eval_batch_size  ,    shuffle=_lowercase )






						def    UpperCamelCase__   (						self  ,    _UpperCamelCase ):






												"""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    UpperCamelCase__   (						self  ,    _UpperCamelCase ):






												"""simple docstring"""
												lowerCAmelCase__         =    self.output_dir.joinpath('best_tfmr' )
												lowerCAmelCase__         =    self.step_count
												self.model.save_pretrained(_lowercase )
												self.tokenizer.save_pretrained(_lowercase )







						@staticmethod
						def    UpperCamelCase__   (						_UpperCamelCase  ,    _UpperCamelCase ):






												"""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    UpperCamelCase__   (						self  ,    _UpperCamelCase  ,    _UpperCamelCase ):






												"""simple docstring"""
												lowerCAmelCase__         =    trainer.lr_schedulers[0]['scheduler']
												lowerCAmelCase__         =    {F"lr_group_{i}": lr for i, lr in enumerate(lr_scheduler.get_lr() )}
												pl_module.logger.log_metrics(_lowercase )






						def    UpperCamelCase__   (						self  ,    _UpperCamelCase  ,    _UpperCamelCase ):






												"""simple docstring"""
												rank_zero_info('***** Validation results *****' )
												lowerCAmelCase__         =    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    UpperCamelCase__   (						self  ,    _UpperCamelCase  ,    _UpperCamelCase ):






												"""simple docstring"""
												rank_zero_info('***** Test results *****' )
												lowerCAmelCase__         =    trainer.callback_metrics
												# Log and save results to file
												lowerCAmelCase__         =    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 (     UpperCamelCase_							:       Optional[int]     ,							UpperCamelCase_							:       Optional[Any]		)	-> None:

						"""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_							:       BaseTransformer     ,							UpperCamelCase_							:       argparse.Namespace     ,							UpperCamelCase_							:       str=None     ,							UpperCamelCase_							:       List[str]=True     ,							UpperCamelCase_							:       str=[]     ,							UpperCamelCase_							:       str=None     ,							UpperCamelCase_							:       str=None     ,							**UpperCamelCase_							:       Tuple     ,							)	-> List[Any]:

						"""simple docstring"""



						pl.seed_everything(args.seed		)

						# init model
						lowerCAmelCase__         =    Path(model.hparams.output_dir		)
						odir.mkdir(exist_ok=__UpperCamelCase		)

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

						lowerCAmelCase__         =    {}

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

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

						lowerCAmelCase__         =    args.accumulate_grad_batches
						lowerCAmelCase__         =    extra_train_kwargs.get('accelerator'     ,							__UpperCamelCase		)
						lowerCAmelCase__         =    extra_train_kwargs.get('profiler'     ,							__UpperCamelCase		)

						lowerCAmelCase__         =    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






from hashlib import mda, shaaaa

__lowerCamelCase     :					Dict				  =						(shaaaa, mda)







class __snake_case							:




   def __init__(       self     :							str			,      _lowercase     :							int = 8 ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= 0B0
     SCREAMING_SNAKE_CASE__     				= size




   def       __a		(       self     :							List[str]			,      _lowercase     :							str ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= self.hash_(_lowercase )
     self.bitarray |= h




   def       __a		(       self     :							int			,      _lowercase     :							str ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= self.hash_(_lowercase )
     return (h & self.bitarray) == h




   def __contains__(       self     :							Union[str, Any]			,      _lowercase     :							str ):



     """simple docstring"""






     return self.exists(_lowercase )




   def       __a		(       self     :							Dict			,      _lowercase     :							int ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= bin(_lowercase )[2:]
     return res.zfill(self.size )




   @property
   def       __a		(       self     :							Tuple ):



     """simple docstring"""






     return self.format_bin(self.bitarray )




   def       __a		(       self     :							Optional[int]			,      _lowercase     :							str ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= 0B0
     for func in HASH_FUNCTIONS:
       SCREAMING_SNAKE_CASE__     				= (
           int.from_bytes(func(value.encode() ).digest()			,      """little""" ) % self.size
       )
       res |= 2**position
     return res




   def       __a		(       self     :							str			,      _lowercase     :							str ):



     """simple docstring"""






     return self.format_bin(self.hash_(_lowercase ) )




   @property
   def       __a		(       self     :							Union[str, Any] ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__     				= bin(self.bitarray ).count("""1""" )
     return (n_ones / self.size) ** len(_lowercase )
