





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

require_version('pytorch_lightning>=1.0.4')

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





class lowerCamelCase    (pl.LightningModule       ):







  def __init__(		self  ,  lowercase__  ,  lowercase__=None  ,  lowercase__="base"  ,  lowercase__=None  ,  lowercase__=None  ,  lowercase__=None  ,  **lowercase__  ,  )		->    str:






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

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







  def 					UpperCAmelCase_					(		self  ,  *lowercase__  ,  **lowercase__		)		->    List[str]:






        """simple docstring"""





        _snake_case		:						Optional[Any]         =  self.model_type.from_pretrained(*lowercase__  ,  **lowercase__		)







  def 					UpperCAmelCase_					(		self		)		->    str:






        """simple docstring"""





        _snake_case		:						Tuple         =  arg_to_scheduler[self.hparams.lr_scheduler]
        _snake_case		:						List[Any]         =  get_schedule_func(
            self.opt  ,  num_warmup_steps=self.hparams.warmup_steps  ,  num_training_steps=self.total_steps()		)
        _snake_case		:						Tuple         =  {'''scheduler''': scheduler, '''interval''': '''step''', '''frequency''': 1}
        return scheduler







  def 					UpperCAmelCase_					(		self		)		->    List[Any]:






        """simple docstring"""





        _snake_case		:						List[Any]         =  self.model
        _snake_case		:						Optional[int]         =  ['''bias''', '''LayerNorm.weight''']
        _snake_case		:						Union[str, 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:
              _snake_case		:						Union[str, Any]         =  Adafactor(
                  lowercase__  ,  lr=self.hparams.learning_rate  ,  scale_parameter=lowercase__  ,  relative_step=lowercase__		)

        else:
              _snake_case		:						str         =  AdamW(
                  lowercase__  ,  lr=self.hparams.learning_rate  ,  eps=self.hparams.adam_epsilon		)
        _snake_case		:						Any         =  optimizer

        _snake_case		:						Any         =  self.get_lr_scheduler()

        return [optimizer], [scheduler]







  def 					UpperCAmelCase_					(		self  ,  lowercase__  ,  lowercase__		)		->    List[str]:






        """simple docstring"""





        return self.validation_step(lowercase__  ,  lowercase__		)







  def 					UpperCAmelCase_					(		self  ,  lowercase__		)		->    str:






        """simple docstring"""





        return self.validation_end(lowercase__		)







  def 					UpperCAmelCase_					(		self		)		->    int:






        """simple docstring"""





        _snake_case		:						Any         =  max(1  ,  self.hparams.gpus		)  # TODO: consider num_tpu_cores
        _snake_case		:						List[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 					UpperCAmelCase_					(		self  ,  lowercase__		)		->    List[Any]:






        """simple docstring"""





        if stage == "test":
              _snake_case		:						List[str]         =  len(self.test_dataloader().dataset		)
        else:
              _snake_case		:						Union[str, Any]         =  self.get_dataloader('''train'''  ,  self.hparams.train_batch_size  ,  shuffle=lowercase__		)
              _snake_case		:						Tuple         =  len(self.train_dataloader().dataset		)







  def 					UpperCAmelCase_					(		self  ,  lowercase__  ,  lowercase__  ,  lowercase__ = False		)		->    List[Any]:






        """simple docstring"""





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







  def 					UpperCAmelCase_					(		self		)		->    Any:






        """simple docstring"""





        return self.train_loader







  def 					UpperCAmelCase_					(		self		)		->    str:






        """simple docstring"""





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







  def 					UpperCAmelCase_					(		self		)		->    int:






        """simple docstring"""





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







  def 					UpperCAmelCase_					(		self  ,  lowercase__		)		->    List[Any]:






        """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  ,  lowercase__		)		->    None:






        """simple docstring"""





        _snake_case		:						Any         =  self.output_dir.joinpath('''best_tfmr'''		)
        _snake_case		:						Optional[Any]         =  self.step_count
        self.model.save_pretrained(lowercase__		)
        self.tokenizer.save_pretrained(lowercase__		)







  @staticmethod
  def 					UpperCAmelCase_					(		lowercase__  ,  lowercase__		)		->    Dict:






        """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 lowerCamelCase    (pl.Callback       ):







  def 					UpperCAmelCase_					(		self  ,  lowercase__  ,  lowercase__		)		->    Any:






        """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 lowerCamelCase    (pl.Callback       ):







  def 					UpperCAmelCase_					(		self  ,  lowercase__  ,  lowercase__		)		->    List[Any]:






        """simple docstring"""





        _snake_case		:						Any         =  trainer.lr_schedulers[0]['''scheduler''']
        _snake_case		:						Optional[int]         =  {F'''lr_group_{i}''': lr for i, lr in enumerate(lr_scheduler.get_lr()		)}
        pl_module.logger.log_metrics(lowercase__		)







  def 					UpperCAmelCase_					(		self  ,  lowercase__  ,  lowercase__		)		->    int:






        """simple docstring"""





        rank_zero_info('''***** Validation results *****'''		)
        _snake_case		:						str         =  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  ,  lowercase__  ,  lowercase__		)		->    str:






        """simple docstring"""





        rank_zero_info('''***** Test results *****'''		)
        _snake_case		:						Optional[int]         =  trainer.callback_metrics
        # Log and save results to file
        _snake_case		:						Optional[int]         =  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  _a     (				lowerCAmelCase_      ,   lowerCAmelCase_						):





      """simple docstring"""



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





      """simple docstring"""



      pl.seed_everything(args.seed						)

      # init model
      _snake_case		:						Any         =  Path(model.hparams.output_dir						)
      odir.mkdir(exist_ok=lowerCAmelCase_						)

      # add custom checkpoints
      if checkpoint_callback is None:
            _snake_case		:						int         =  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:
            _snake_case		:						Union[str, Any]         =  LoggingCallback()

      _snake_case		:						List[Any]         =  {}

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

      if args.gpus > 1:
            _snake_case		:						Dict         =  '''auto'''  # "ddp"
            _snake_case		:						int         =  '''ddp'''

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

      _snake_case		:						Tuple         =  pl.Trainer.from_argparse_args(
          lowerCAmelCase_      ,   weights_summary=lowerCAmelCase_      ,   callbacks=[logging_callback] + extra_callbacks + [checkpoint_callback] + [InitCallback()]      ,   logger=lowerCAmelCase_      ,   **lowerCAmelCase_      ,   )

      if args.do_train:
            trainer.fit(lowerCAmelCase_						)

      return trainer









'''simple docstring'''

from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


UpperCAmelCase       :		Dict		=				{'configuration_hubert': ['HUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'HubertConfig']}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase       :		Union[str, Any]		=				[
       'HUBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
       'HubertForCTC',
       'HubertForSequenceClassification',
       'HubertModel',
       'HubertPreTrainedModel',
   ]


try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase       :		Optional[int]		=				[
       'TF_HUBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
       'TFHubertForCTC',
       'TFHubertModel',
       'TFHubertPreTrainedModel',
   ]

if TYPE_CHECKING:
   from .configuration_hubert import HUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, HubertConfig

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_hubert import (
          HUBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
          HubertForCTC,
          HubertForSequenceClassification,
          HubertModel,
          HubertPreTrainedModel,
      )

   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_hubert import (
          TF_HUBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
          TFHubertForCTC,
          TFHubertModel,
          TFHubertPreTrainedModel,
      )


else:
   import sys

   UpperCAmelCase       :		Optional[int]		=				_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)


