






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


snake_case   :   Optional[int]    		=		logging.getLogger(__name__)

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

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



class 					UpperCamelCase__   (							pl.LightningModule):







    """simple docstring"""







    def __init__(		self   :  Tuple							,						UpperCamelCase_   :  argparse.Namespace							,						UpperCamelCase_   :  int=None							,						UpperCamelCase_   :  Union[str, Any]="base"							,						UpperCamelCase_   :  Union[str, Any]=None							,						UpperCamelCase_   :  Any=None							,						UpperCamelCase_   :  str=None							,						**UpperCamelCase_   :  Tuple							,						):





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

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







    def 							a__ (		self   :  List[str]							,						*UpperCamelCase_   :  Any							,						**UpperCamelCase_   :  Any   ):





           '''simple docstring'''

           __magic_name__										=						self.model_type.from_pretrained(*UpperCamelCase_							,						**UpperCamelCase_   )







    def 							a__ (		self   :  Dict   ):





           '''simple docstring'''

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







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





           '''simple docstring'''

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

           else:
                  __magic_name__										=						AdamW(
                      UpperCamelCase_							,						lr=self.hparams.learning_rate							,						eps=self.hparams.adam_epsilon   )
           __magic_name__										=						optimizer

           __magic_name__										=						self.get_lr_scheduler()

           return [optimizer], [scheduler]







    def 							a__ (		self   :  str							,						UpperCamelCase_   :  Union[str, Any]							,						UpperCamelCase_   :  Tuple   ):





           '''simple docstring'''

           return self.validation_step(UpperCamelCase_							,						UpperCamelCase_   )







    def 							a__ (		self   :  int							,						UpperCamelCase_   :  Optional[int]   ):





           '''simple docstring'''

           return self.validation_end(UpperCamelCase_   )







    def 							a__ (		self   :  Optional[int]   ):





           '''simple docstring'''

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







    def 							a__ (		self   :  List[str]							,						UpperCamelCase_   :  List[str]   ):





           '''simple docstring'''

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







    def 							a__ (		self   :  List[str]							,						UpperCamelCase_   :  str							,						UpperCamelCase_   :  int							,						UpperCamelCase_   :  bool = False   ):





           '''simple docstring'''

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







    def 							a__ (		self   :  int   ):





           '''simple docstring'''

           return self.train_loader







    def 							a__ (		self   :  int   ):





           '''simple docstring'''

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







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





           '''simple docstring'''

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







    def 							a__ (		self   :  str							,						UpperCamelCase_   :  Optional[int]   ):





           '''simple docstring'''

           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 							a__ (		self   :  Tuple							,						UpperCamelCase_   :  Dict[str, Any]   ):





           '''simple docstring'''

           __magic_name__										=						self.output_dir.joinpath('best_tfmr'   )
           __magic_name__										=						self.step_count
           self.model.save_pretrained(UpperCamelCase_   )
           self.tokenizer.save_pretrained(UpperCamelCase_   )







    @staticmethod
    def 							a__ (		UpperCamelCase_   :  List[str]							,						UpperCamelCase_   :  Union[str, Any]   ):





           '''simple docstring'''

           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=3_2							,						type=UpperCamelCase_   )
           parser.add_argument('--eval_batch_size'							,						default=3_2							,						type=UpperCamelCase_   )
           parser.add_argument('--adafactor'							,						action='store_true'   )



class 					UpperCamelCase__   (							pl.Callback):







    """simple docstring"""







    def 							a__ (		self   :  Any							,						UpperCamelCase_   :  str							,						UpperCamelCase_   :  str   ):





           '''simple docstring'''

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







    def 							a__ (		self   :  Optional[int]							,						UpperCamelCase_   :  pl.Trainer							,						UpperCamelCase_   :  pl.LightningModule   ):





           '''simple docstring'''

           rank_zero_info('***** Validation results *****'   )
           __magic_name__										=						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 							a__ (		self   :  List[str]							,						UpperCamelCase_   :  pl.Trainer							,						UpperCamelCase_   :  pl.LightningModule   ):





           '''simple docstring'''

           rank_zero_info('***** Test results *****'   )
           __magic_name__										=						trainer.callback_metrics
           # Log and save results to file
           __magic_name__										=						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        A	(							__snake_case:   Tuple   ,    __snake_case:   Tuple     )						->		None:

       """simple docstring"""





       parser.add_argument(
           '--output_dir'   ,    default=__snake_case   ,    type=__snake_case   ,    required=__snake_case   ,    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=__snake_case   ,    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=__snake_case     )
       parser.add_argument('--max_grad_norm'   ,    dest='gradient_clip_val'   ,    default=1.0   ,    type=__snake_case   ,    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=__snake_case   ,    default=1   ,    help='Number of updates steps to accumulate before performing a backward/update pass.'   ,    )
       parser.add_argument('--seed'   ,    type=__snake_case   ,    default=4_2   ,    help='random seed for initialization'     )
       parser.add_argument(
           '--data_dir'   ,    default=__snake_case   ,    type=__snake_case   ,    required=__snake_case   ,    help='The input data dir. Should contain the training files for the CoNLL-2003 NER task.'   ,    )






def        A	(							__snake_case:   BaseTransformer   ,    __snake_case:   argparse.Namespace   ,    __snake_case:   str=None   ,    __snake_case:   str=True   ,    __snake_case:   Optional[int]=[]   ,    __snake_case:   Tuple=None   ,    __snake_case:   Dict=None   ,    **__snake_case:   Dict   ,    )						->		Any:

       """simple docstring"""





       pl.seed_everything(args.seed     )

       # init model
       __magic_name__										=						Path(model.hparams.output_dir     )
       odir.mkdir(exist_ok=__snake_case     )

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

       __magic_name__										=						{}

       # TODO: remove with PyTorch 1.6 since pl uses native amp
       if args.fpaa:
              __magic_name__										=						1_6
              __magic_name__										=						args.fpaa_opt_level

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

       __magic_name__										=						args.accumulate_grad_batches
       __magic_name__										=						extra_train_kwargs.get('accelerator'   ,    __snake_case     )
       __magic_name__										=						extra_train_kwargs.get('profiler'   ,    __snake_case     )

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

       if args.do_train:
              trainer.fit(__snake_case     )

       return trainer







"""simple docstring"""
from typing import List, Optional, Union

import numpy as np

from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
from ...feature_extraction_utils import BatchFeature
from ...utils import PaddingStrategy, TensorType, logging


snake_case   :   Tuple    		=		logging.get_logger(__name__)



class 					UpperCamelCase__   (							a_):







    """simple docstring"""
    __UpperCAmelCase											=   ["""input_values""", """attention_mask"""]







    def __init__(		self   :  Union[str, Any]							,						UpperCamelCase_   :  int=1							,						UpperCamelCase_   :  Optional[int]=1_6_0_0_0							,						UpperCamelCase_   :  List[str]=0.0							,						UpperCamelCase_   :  Optional[int]=False							,						UpperCamelCase_   :  Tuple=True							,						**UpperCamelCase_   :  Tuple							,						):





           '''simple docstring'''

           super().__init__(feature_size=UpperCamelCase_							,						sampling_rate=UpperCamelCase_							,						padding_value=UpperCamelCase_							,						**UpperCamelCase_   )
           __magic_name__										=						return_attention_mask
           __magic_name__										=						do_normalize







    @staticmethod
    def 							a__ (		UpperCamelCase_   :  List[np.ndarray]							,						UpperCamelCase_   :  List[np.ndarray]							,						UpperCamelCase_   :  float = 0.0   ):





           '''simple docstring'''

           if attention_mask is not None:
                  __magic_name__										=						np.array(UpperCamelCase_							,						np.intaa   )
                  __magic_name__										=						[]

                  for vector, length in zip(UpperCamelCase_							,						attention_mask.sum(-1   )   ):
                         __magic_name__										=						(vector - vector[:length].mean()) / np.sqrt(vector[:length].var() + 1e-7   )
                         if length < normed_slice.shape[0]:
                                __magic_name__										=						padding_value

                         normed_input_values.append(UpperCamelCase_   )
           else:
                  __magic_name__										=						[(x - x.mean()) / np.sqrt(x.var() + 1e-7   ) for x in input_values]

           return normed_input_values







    def __call__(		self   :  Any							,						UpperCamelCase_   :  Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]							,						UpperCamelCase_   :  Union[bool, str, PaddingStrategy] = False							,						UpperCamelCase_   :  Optional[int] = None							,						UpperCamelCase_   :  bool = False							,						UpperCamelCase_   :  Optional[int] = None							,						UpperCamelCase_   :  Optional[bool] = None							,						UpperCamelCase_   :  Optional[Union[str, TensorType]] = None							,						UpperCamelCase_   :  Optional[int] = None							,						**UpperCamelCase_   :  Any							,						):





           '''simple docstring'''

           if sampling_rate is not None:
                  if sampling_rate != self.sampling_rate:
                         raise ValueError(
                             f"""The model corresponding to this feature extractor: {self} was trained using a sampling rate of"""
                             f""" {self.sampling_rate}. Please make sure that the provided `raw_speech` input was sampled with"""
                             f""" {self.sampling_rate} and not {sampling_rate}."""   )
           else:
                  logger.warning(
                      'It is strongly recommended to pass the ``sampling_rate`` argument to this function. '
                      'Failing to do so can result in silent errors that might be hard to debug.'   )

           __magic_name__										=						isinstance(UpperCamelCase_							,						np.ndarray   ) and len(raw_speech.shape   ) > 1
           if is_batched_numpy and len(raw_speech.shape   ) > 2:
                  raise ValueError(f"""Only mono-channel audio is supported for input to {self}"""   )
           __magic_name__										=						is_batched_numpy or (
               isinstance(UpperCamelCase_							,						(list, tuple)   ) and (isinstance(raw_speech[0]							,						(np.ndarray, tuple, list)   ))
           )

           # always return batch
           if not is_batched:
                  __magic_name__										=						[raw_speech]

           # convert into correct format for padding
           __magic_name__										=						BatchFeature({'input_values': raw_speech}   )

           __magic_name__										=						self.pad(
               UpperCamelCase_							,						padding=UpperCamelCase_							,						max_length=UpperCamelCase_							,						truncation=UpperCamelCase_							,						pad_to_multiple_of=UpperCamelCase_							,						return_attention_mask=UpperCamelCase_							,						)

           # convert input values to correct format
           __magic_name__										=						padded_inputs['input_values']
           if not isinstance(input_values[0]							,						np.ndarray   ):
                  __magic_name__										=						[np.asarray(UpperCamelCase_							,						dtype=np.floataa   ) for array in input_values]
           elif (
               not isinstance(UpperCamelCase_							,						np.ndarray   )
               and isinstance(input_values[0]							,						np.ndarray   )
               and input_values[0].dtype is np.dtype(np.floataa   )
           ):
                  __magic_name__										=						[array.astype(np.floataa   ) for array in input_values]
           elif isinstance(UpperCamelCase_							,						np.ndarray   ) and input_values.dtype is np.dtype(np.floataa   ):
                  __magic_name__										=						input_values.astype(np.floataa   )

           # convert attention_mask to correct format
           __magic_name__										=						padded_inputs.get('attention_mask'   )
           if attention_mask is not None:
                  __magic_name__										=						[np.asarray(UpperCamelCase_							,						dtype=np.intaa   ) for array in attention_mask]

           # zero-mean and unit-variance normalization
           if self.do_normalize:
                  __magic_name__										=						(
                      attention_mask
                      if self._get_padding_strategies(UpperCamelCase_							,						max_length=UpperCamelCase_   ) is not PaddingStrategy.DO_NOT_PAD
                      else None
                  )
                  __magic_name__										=						self.zero_mean_unit_var_norm(
                      padded_inputs['input_values']							,						attention_mask=UpperCamelCase_							,						padding_value=self.padding_value   )

           if return_tensors is not None:
                  __magic_name__										=						padded_inputs.convert_to_tensors(UpperCamelCase_   )

           return padded_inputs