

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

require_version('''pytorch_lightning>=1.0.4''')

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





class     snake_case  (	pl.LightningModule							):


   '''simple docstring'''







   def __init__(					self     :	Tuple						,    lowerCamelCase_     :	argparse.Namespace						,    lowerCamelCase_     :	int=None						,    lowerCamelCase_     :	Optional[Any]="base"						,    lowerCamelCase_     :	Union[str, Any]=None						,    lowerCamelCase_     :	Union[str, Any]=None						,    lowerCamelCase_     :	List[str]=None						,    **lowerCamelCase_     :	List[Any]						,    )					->Optional[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_			)
     UpperCAmelCase__    						=   0
     UpperCAmelCase__    						=   Path(self.hparams.output_dir			)
     UpperCAmelCase__    						=   self.hparams.cache_dir if self.hparams.cache_dir else None
     if config is None:
       UpperCAmelCase__    						=   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:
       UpperCAmelCase__    						=   config

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




   def 	UpperCAmelCase       (					self     :	Union[str, Any]						,    *lowerCamelCase_     :	List[str]						,    **lowerCamelCase_     :	List[str]			)					->List[Any]:
     '''simple docstring'''







     UpperCAmelCase__    						=   self.model_type.from_pretrained(*lowerCamelCase_						,    **lowerCamelCase_			)




   def 	UpperCAmelCase       (					self     :	Union[str, Any]			)					->Optional[Any]:
     '''simple docstring'''







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




   def 	UpperCAmelCase       (					self     :	Dict			)					->Dict:
     '''simple docstring'''







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

     else:
       UpperCAmelCase__    						=   AdamW(
           lowerCamelCase_						,    lr=self.hparams.learning_rate						,    eps=self.hparams.adam_epsilon			)
     UpperCAmelCase__    						=   optimizer

     UpperCAmelCase__    						=   self.get_lr_scheduler()

     return [optimizer], [scheduler]




   def 	UpperCAmelCase       (					self     :	Optional[int]						,    lowerCamelCase_     :	Optional[int]						,    lowerCamelCase_     :	int			)					->Optional[int]:
     '''simple docstring'''







     return self.validation_step(lowerCamelCase_						,    lowerCamelCase_			)




   def 	UpperCAmelCase       (					self     :	Any						,    lowerCamelCase_     :	Optional[Any]			)					->List[str]:
     '''simple docstring'''







     return self.validation_end(lowerCamelCase_			)




   def 	UpperCAmelCase       (					self     :	List[Any]			)					->int:
     '''simple docstring'''







     UpperCAmelCase__    						=   max(1						,    self.hparams.gpus			)  # TODO: consider num_tpu_cores
     UpperCAmelCase__    						=   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     :	int						,    lowerCamelCase_     :	Dict			)					->Union[str, Any]:
     '''simple docstring'''







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




   def 	UpperCAmelCase       (					self     :	Dict						,    lowerCamelCase_     :	str						,    lowerCamelCase_     :	int						,    lowerCamelCase_     :	bool = False			)					->Optional[Any]:
     '''simple docstring'''







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




   def 	UpperCAmelCase       (					self     :	Optional[Any]			)					->List[str]:
     '''simple docstring'''







     return self.train_loader




   def 	UpperCAmelCase       (					self     :	Any			)					->Any:
     '''simple docstring'''







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




   def 	UpperCAmelCase       (					self     :	List[Any]			)					->List[Any]:
     '''simple docstring'''







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




   def 	UpperCAmelCase       (					self     :	Dict						,    lowerCamelCase_     :	List[Any]			)					->Dict:
     '''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 	UpperCAmelCase       (					self     :	Any						,    lowerCamelCase_     :	Dict[str, Any]			)					->None:
     '''simple docstring'''







     UpperCAmelCase__    						=   self.output_dir.joinpath("""best_tfmr"""			)
     UpperCAmelCase__    						=   self.step_count
     self.model.save_pretrained(lowerCamelCase_			)
     self.tokenizer.save_pretrained(lowerCamelCase_			)







   @staticmethod
   def 	UpperCAmelCase       (					lowerCamelCase_     :	Optional[Any]						,    lowerCamelCase_     :	List[Any]			)					->List[str]:
     '''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     snake_case  (	pl.Callback							):


   '''simple docstring'''







   def 	UpperCAmelCase       (					self     :	Dict						,    lowerCamelCase_     :	Optional[Any]						,    lowerCamelCase_     :	int			)					->Any:
     '''simple docstring'''







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




   def 	UpperCAmelCase       (					self     :	Any						,    lowerCamelCase_     :	pl.Trainer						,    lowerCamelCase_     :	pl.LightningModule			)					->Union[str, Any]:
     '''simple docstring'''







     rank_zero_info("""***** Validation results *****"""			)
     UpperCAmelCase__    						=   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 	UpperCAmelCase       (					self     :	List[Any]						,    lowerCamelCase_     :	pl.Trainer						,    lowerCamelCase_     :	pl.LightningModule			)					->List[str]:
     '''simple docstring'''







     rank_zero_info("""***** Test results *****"""			)
     UpperCAmelCase__    						=   trainer.callback_metrics
     # Log and save results to file
     UpperCAmelCase__    						=   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 						lowercase		(     _lowerCAmelCase   ,     _lowerCAmelCase      ):
  #  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=42   ,     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 						lowercase		(     _lowerCAmelCase   ,     _lowerCAmelCase   ,     _lowerCAmelCase=None   ,     _lowerCAmelCase=True   ,     _lowerCAmelCase=[]   ,     _lowerCAmelCase=None   ,     _lowerCAmelCase=None   ,     **_lowerCAmelCase   ,     ):
  pl.seed_everything(args.seed      )

  # init model
  UpperCAmelCase__    						=   Path(model.hparams.output_dir      )
  odir.mkdir(exist_ok=_lowerCAmelCase      )

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

  UpperCAmelCase__    						=   {}

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

  if args.gpus > 1:
    UpperCAmelCase__    						=   """ddp"""

  UpperCAmelCase__    						=   args.accumulate_grad_batches
  UpperCAmelCase__    						=   extra_train_kwargs.get("""accelerator"""   ,     _lowerCAmelCase      )
  UpperCAmelCase__    						=   extra_train_kwargs.get("""profiler"""   ,     _lowerCAmelCase      )

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










# Copyright 2022 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pathlib import Path

from .config_args import default_config_file, load_config_from_file
from .config_utils import SubcommandHelpFormatter


snake_case__ :						str     =			'''Update an existing config file with the latest defaults while maintaining the old configuration.'''


def 						lowercase		(     _lowerCAmelCase      ):
  UpperCAmelCase__    						=   args.config_file
  if config_file is None and Path(_lowerCAmelCase      ).exists():
    UpperCAmelCase__    						=   default_config_file
  elif not Path(_lowerCAmelCase      ).exists():
    raise ValueError(F'''The passed config file located at {config_file} doesn\'t exist.'''      )
  UpperCAmelCase__    						=   load_config_from_file(_lowerCAmelCase      )

  if config_file.endswith(""".json"""      ):
    config.to_json_file(_lowerCAmelCase      )
  else:
    config.to_yaml_file(_lowerCAmelCase      )
  return config_file


def 						lowercase		(     _lowerCAmelCase   ,     _lowerCAmelCase      ):
  UpperCAmelCase__    						=   parser.add_parser("""update"""   ,     parents=_lowerCAmelCase   ,     help=_lowerCAmelCase   ,     formatter_class=_lowerCAmelCase      )
  parser.add_argument(
      """--config_file"""   ,     default=_lowerCAmelCase   ,     help=(
          """The path to the config file to update. Will default to a file named default_config.yaml in the cache """
          """location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have """
          """such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed """
          """with 'huggingface'."""
      )   ,     )

  parser.set_defaults(func=_lowerCAmelCase      )
  return parser





def 						lowercase		(     _lowerCAmelCase      ):
  UpperCAmelCase__    						=   update_config(_lowerCAmelCase      )
  print(F'''Sucessfully updated the configuration file at {config_file}.'''      )




