






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


UpperCamelCase__	=     logging.getLogger(__name__)

require_version('pytorch_lightning>=1.0.4')

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


class     _UpperCAmelCase  (						pl.LightningModule     ):




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

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

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




    def 			lowerCAmelCase__			(     self					:    List[str]					,       *a					:    Optional[int]					,       **a					:    str     ):






       '''simple docstring'''






       lowercase_		:      Union[str, Any]          =       self.model_type.from_pretrained(*a					,       **a     )




    def 			lowerCAmelCase__			(     self					:    Tuple     ):






       '''simple docstring'''






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




    def 			lowerCAmelCase__			(     self					:    Union[str, Any]     ):






       '''simple docstring'''






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

       else:
          lowercase_		:      List[Any]          =       AdamW(
              a					,       lr=self.hparams.learning_rate					,       eps=self.hparams.adam_epsilon     )
       lowercase_		:      str          =       optimizer

       lowercase_		:      List[str]          =       self.get_lr_scheduler()

       return [optimizer], [scheduler]




    def 			lowerCAmelCase__			(     self					:    Dict					,       a					:    Optional[int]					,       a					:    int     ):






       '''simple docstring'''






       return self.validation_step(a					,       a     )




    def 			lowerCAmelCase__			(     self					:    int					,       a					:    Union[str, Any]     ):






       '''simple docstring'''






       return self.validation_end(a     )




    def 			lowerCAmelCase__			(     self					:    List[str]     ):






       '''simple docstring'''






       lowercase_		:      Tuple          =       max(1					,       self.hparams.gpus     )  # TODO: consider num_tpu_cores
       lowercase_		:      Optional[int]          =       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					:    Optional[int]					,       a					:    Optional[Any]     ):






       '''simple docstring'''






       if stage == "test":
          lowercase_		:      Optional[int]          =       len(self.test_dataloader().dataset     )
       else:
          lowercase_		:      Union[str, Any]          =       self.get_dataloader("train"					,       self.hparams.train_batch_size					,       shuffle=a     )
          lowercase_		:      Any          =       len(self.train_dataloader().dataset     )




    def 			lowerCAmelCase__			(     self					:    Dict					,       a					:    str					,       a					:    int					,       a					:    bool = False     ):






       '''simple docstring'''






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




    def 			lowerCAmelCase__			(     self					:    Dict     ):






       '''simple docstring'''






       return self.train_loader




    def 			lowerCAmelCase__			(     self					:    Any     ):






       '''simple docstring'''






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




    def 			lowerCAmelCase__			(     self					:    str     ):






       '''simple docstring'''






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




    def 			lowerCAmelCase__			(     self					:    Dict					,       a					:    List[Any]     ):






       '''simple docstring'''






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




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






       '''simple docstring'''






       lowercase_		:      str          =       self.output_dir.joinpath("best_tfmr"     )
       lowercase_		:      List[Any]          =       self.step_count
       self.model.save_pretrained(a     )
       self.tokenizer.save_pretrained(a     )





    @staticmethod
    def 			lowerCAmelCase__			(     a					:    str					,       a					:    Union[str, Any]     ):






       '''simple docstring'''






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


class     _UpperCAmelCase  (						pl.Callback     ):




    def 			lowerCAmelCase__			(     self					:    Optional[Any]					,       a					:    int					,       a					:    Optional[int]     ):






       '''simple docstring'''






       if (
           trainer.is_global_zero and trainer.global_rank == 0
       ):  # we initialize the retriever only on master worker with RAY. In new pytorch-lightning accelorators are removed.
          pl_module.model.rag.retriever.init_retrieval()  # better to use hook functions.


class     _UpperCAmelCase  (						pl.Callback     ):




    def 			lowerCAmelCase__			(     self					:    List[str]					,       a					:    Optional[Any]					,       a					:    List[str]     ):






       '''simple docstring'''






       lowercase_		:      List[str]          =       trainer.lr_schedulers[0]["scheduler"]
       lowercase_		:      Union[str, Any]          =       {f"""lr_group_{i}""": lr for i, lr in enumerate(lr_scheduler.get_lr()     )}
       pl_module.logger.log_metrics(a     )




    def 			lowerCAmelCase__			(     self					:    Dict					,       a					:    pl.Trainer					,       a					:    pl.LightningModule     ):






       '''simple docstring'''






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





    def 			lowerCAmelCase__			(     self					:    Optional[Any]					,       a					:    pl.Trainer					,       a					:    pl.LightningModule     ):






       '''simple docstring'''






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




def 						__SCREAMING_SNAKE_CASE   (					_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 						__SCREAMING_SNAKE_CASE   (					_UpperCamelCase  ,						_UpperCamelCase  ,						_UpperCamelCase=None  ,						_UpperCamelCase=True  ,						_UpperCamelCase=None  ,						_UpperCamelCase=[]  ,						_UpperCamelCase=None  ,						_UpperCamelCase=None  ,						**_UpperCamelCase  ,						):


   """simple docstring"""




   pl.seed_everything(args.seed							)

   # init model
   lowercase_		:      str          =       Path(model.hparams.output_dir							)
   odir.mkdir(exist_ok=_UpperCamelCase							)

   # add custom checkpoints
   if checkpoint_callback is None:
      lowercase_		:      Optional[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(_UpperCamelCase							)
   if logging_callback is None:
      lowercase_		:      Dict          =       LoggingCallback()

   lowercase_		:      Optional[Any]          =       {}

   # TODO: remove with PyTorch 1.6 since pl uses native amp
   if args.fpaa:
      lowercase_		:      Optional[int]          =       16
      # train_params["amp_level"] = args.fp16_opt_level

   if args.gpus > 1:
      lowercase_		:      str          =       "auto"  # "ddp"
      lowercase_		:      str          =       "ddp"

   lowercase_		:      Tuple          =       args.accumulate_grad_batches
   lowercase_		:      int          =       None  # extra_train_kwargs.get("profiler", None) #get unwanted logs
   lowercase_		:      Union[str, Any]          =       "auto"

   lowercase_		:      Optional[Any]          =       pl.Trainer.from_argparse_args(
       _UpperCamelCase  ,						weights_summary=_UpperCamelCase  ,						callbacks=[logging_callback] + extra_callbacks + [checkpoint_callback] + [InitCallback()]  ,						logger=_UpperCamelCase  ,						**_UpperCamelCase  ,						)

   if args.do_train:
      trainer.fit(_UpperCamelCase							)

   return trainer












'''simple docstring'''
from typing import TYPE_CHECKING

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


UpperCamelCase__	=     {
    'configuration_unispeech_sat': ['UNISPEECH_SAT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'UniSpeechSatConfig'],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 UpperCamelCase__	=     [
     'UNISPEECH_SAT_PRETRAINED_MODEL_ARCHIVE_LIST',
     'UniSpeechSatForAudioFrameClassification',
     'UniSpeechSatForCTC',
     'UniSpeechSatForPreTraining',
     'UniSpeechSatForSequenceClassification',
     'UniSpeechSatForXVector',
     'UniSpeechSatModel',
     'UniSpeechSatPreTrainedModel',
 ]

if TYPE_CHECKING:
 from .configuration_unispeech_sat import UNISPEECH_SAT_PRETRAINED_CONFIG_ARCHIVE_MAP, UniSpeechSatConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_unispeech_sat import (
      UNISPEECH_SAT_PRETRAINED_MODEL_ARCHIVE_LIST,
      UniSpeechSatForAudioFrameClassification,
      UniSpeechSatForCTC,
      UniSpeechSatForPreTraining,
      UniSpeechSatForSequenceClassification,
      UniSpeechSatForXVector,
      UniSpeechSatModel,
      UniSpeechSatPreTrainedModel,
  )

else:
 import sys

 UpperCamelCase__	=     _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




