






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


__lowerCAmelCase					   =       logging.getLogger(__name__)

require_version("pytorch_lightning>=1.0.4")

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




class      __SCREAMING_SNAKE_CASE					(pl.LightningModule					):


  """simple docstring"""


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

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







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






         """simple docstring"""


         a_	    = self.model_type.from_pretrained(*UpperCamelCase__      ,  **UpperCamelCase__  )







  def       _a					( self  ):






         """simple docstring"""


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







  def       _a					( self  ):






         """simple docstring"""


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

         else:
                a_	    = AdamW(
                    UpperCamelCase__      ,  lr=self.hparams.learning_rate      ,  eps=self.hparams.adam_epsilon  )
         a_	    = optimizer

         a_	    = self.get_lr_scheduler()

         return [optimizer], [scheduler]







  def       _a					( self      ,  UpperCamelCase__      ,  UpperCamelCase__  ):






         """simple docstring"""


         return self.validation_step(UpperCamelCase__      ,  UpperCamelCase__  )







  def       _a					( self      ,  UpperCamelCase__  ):






         """simple docstring"""


         return self.validation_end(UpperCamelCase__  )







  def       _a					( self  ):






         """simple docstring"""


         a_	    = max(1      ,  self.hparams.gpus  )  # TODO: consider num_tpu_cores
         a_	    = 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      ,  UpperCamelCase__  ):






         """simple docstring"""


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







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






         """simple docstring"""


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







  def       _a					( self  ):






         """simple docstring"""


         return self.train_loader







  def       _a					( self  ):






         """simple docstring"""


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







  def       _a					( self  ):






         """simple docstring"""


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







  def       _a					( self      ,  UpperCamelCase__  ):






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






         """simple docstring"""


         a_	    = self.output_dir.joinpath('best_tfmr'  )
         a_	    = self.step_count
         self.model.save_pretrained(UpperCamelCase__  )
         self.tokenizer.save_pretrained(UpperCamelCase__  )







  @staticmethod
  def       _a					( UpperCamelCase__      ,  UpperCamelCase__  ):






         """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=32      ,  type=UpperCamelCase__  )
         parser.add_argument('--eval_batch_size'      ,  default=32      ,  type=UpperCamelCase__  )
         parser.add_argument('--adafactor'      ,  action='store_true'  )





class      __SCREAMING_SNAKE_CASE					(pl.Callback					):


  """simple docstring"""


  def       _a					( self      ,  UpperCamelCase__      ,  UpperCamelCase__  ):






         """simple docstring"""


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







  def       _a					( self      ,  UpperCamelCase__      ,  UpperCamelCase__  ):






         """simple docstring"""


         rank_zero_info('***** Validation results *****'  )
         a_	    = 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      ,  UpperCamelCase__      ,  UpperCamelCase__  ):






         """simple docstring"""


         rank_zero_info('***** Test results *****'  )
         a_	    = trainer.callback_metrics
         # Log and save results to file
         a_	    = 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 			__UpperCamelCase				(			lowercase_					:     str      ,					lowercase_					:     int      ):
       """simple docstring"""




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


def 			__UpperCamelCase				(			lowercase_					:     BaseTransformer      ,					lowercase_					:     argparse.Namespace      ,					lowercase_					:     Tuple=None      ,					lowercase_					:     Union[str, Any]=True      ,					lowercase_					:     Optional[Any]=[]      ,					lowercase_					:     int=None      ,					lowercase_					:     Tuple=None      ,					**lowercase_					:     Any      ,					):
       """simple docstring"""




       pl.seed_everything(args.seed      )

       # init model
       a_	    = Path(model.hparams.output_dir      )
       odir.mkdir(exist_ok=lowercase_      )

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

       a_	    = {}

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

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

       a_	    = args.accumulate_grad_batches
       a_	    = extra_train_kwargs.get('accelerator'      ,					lowercase_      )
       a_	    = extra_train_kwargs.get('profiler'      ,					lowercase_      )

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

       if args.do_train:
              trainer.fit(lowercase_      )

       return trainer












'''simple docstring'''







import json
import os
import shutil
import tempfile
import unittest

from transformers.models.wavaveca import WavaVecaCTCTokenizer, WavaVecaFeatureExtractor, WavaVecaProcessor
from transformers.models.wavaveca.tokenization_wavaveca import VOCAB_FILES_NAMES
from transformers.utils import FEATURE_EXTRACTOR_NAME

from .test_feature_extraction_wavaveca import floats_list




class      __SCREAMING_SNAKE_CASE					(unittest.TestCase					):


  """simple docstring"""


  def       _a					( self  ):






         """simple docstring"""


         a_	    = '<pad> <s> </s> <unk> | E T A O N I H S R D L U M W C F G Y P B V K \' X J Q Z'.split(' '  )
         a_	    = dict(zip(UpperCamelCase__      ,  range(len(UpperCamelCase__  )  )  )  )

         a_	    = {
             'pad_token': '<pad>',
             'unk_token': '<unk>',
             'bos_token': '<s>',
             'eos_token': '</s>',
         }
         a_	    = {
             'feature_size': 1,
             'padding_value': 0.0,
             'sampling_rate': 16_000,
             'return_attention_mask': False,
             'do_normalize': True,
         }

         a_	    = tempfile.mkdtemp()
         a_	    = os.path.join(self.tmpdirname      ,  VOCAB_FILES_NAMES['vocab_file']  )
         a_	    = os.path.join(self.tmpdirname      ,  UpperCamelCase__  )
         with open(self.vocab_file      ,  'w'      ,  encoding='utf-8'  ) as fp:
                fp.write(json.dumps(UpperCamelCase__  ) + '\n'  )

         with open(self.feature_extraction_file      ,  'w'      ,  encoding='utf-8'  ) as fp:
                fp.write(json.dumps(UpperCamelCase__  ) + '\n'  )







  def       _a					( self      ,  **UpperCamelCase__  ):






         """simple docstring"""


         a_	    = self.add_kwargs_tokens_map.copy()
         kwargs.update(UpperCamelCase__  )
         return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname      ,  **UpperCamelCase__  )







  def       _a					( self      ,  **UpperCamelCase__  ):






         """simple docstring"""


         return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname      ,  **UpperCamelCase__  )







  def       _a					( self  ):






         """simple docstring"""


         shutil.rmtree(self.tmpdirname  )







  def       _a					( self  ):






         """simple docstring"""


         a_	    = self.get_tokenizer()
         a_	    = self.get_feature_extractor()

         a_	    = WavaVecaProcessor(tokenizer=UpperCamelCase__      ,  feature_extractor=UpperCamelCase__  )

         processor.save_pretrained(self.tmpdirname  )
         a_	    = WavaVecaProcessor.from_pretrained(self.tmpdirname  )

         self.assertEqual(processor.tokenizer.get_vocab()      ,  tokenizer.get_vocab()  )
         self.assertIsInstance(processor.tokenizer      ,  UpperCamelCase__  )

         self.assertEqual(processor.feature_extractor.to_json_string()      ,  feature_extractor.to_json_string()  )
         self.assertIsInstance(processor.feature_extractor      ,  UpperCamelCase__  )







  def       _a					( self  ):






         """simple docstring"""


         a_	    = WavaVecaProcessor(tokenizer=self.get_tokenizer()      ,  feature_extractor=self.get_feature_extractor()  )
         processor.save_pretrained(self.tmpdirname  )

         a_	    = self.get_tokenizer(bos_token='(BOS)'      ,  eos_token='(EOS)'  )
         a_	    = self.get_feature_extractor(do_normalize=UpperCamelCase__      ,  padding_value=1.0  )

         a_	    = WavaVecaProcessor.from_pretrained(
             self.tmpdirname      ,  bos_token='(BOS)'      ,  eos_token='(EOS)'      ,  do_normalize=UpperCamelCase__      ,  padding_value=1.0  )

         self.assertEqual(processor.tokenizer.get_vocab()      ,  tokenizer_add_kwargs.get_vocab()  )
         self.assertIsInstance(processor.tokenizer      ,  UpperCamelCase__  )

         self.assertEqual(processor.feature_extractor.to_json_string()      ,  feature_extractor_add_kwargs.to_json_string()  )
         self.assertIsInstance(processor.feature_extractor      ,  UpperCamelCase__  )







  def       _a					( self  ):






         """simple docstring"""


         a_	    = self.get_feature_extractor()
         a_	    = self.get_tokenizer()

         a_	    = WavaVecaProcessor(tokenizer=UpperCamelCase__      ,  feature_extractor=UpperCamelCase__  )

         a_	    = floats_list((3, 1_000)  )

         a_	    = feature_extractor(UpperCamelCase__      ,  return_tensors='np'  )
         a_	    = processor(UpperCamelCase__      ,  return_tensors='np'  )

         for key in input_feat_extract.keys():
                self.assertAlmostEqual(input_feat_extract[key].sum()      ,  input_processor[key].sum()      ,  delta=1e-2  )







  def       _a					( self  ):






         """simple docstring"""


         a_	    = self.get_feature_extractor()
         a_	    = self.get_tokenizer()

         a_	    = WavaVecaProcessor(tokenizer=UpperCamelCase__      ,  feature_extractor=UpperCamelCase__  )

         a_	    = 'This is a test string'

         a_	    = processor(text=UpperCamelCase__  )

         a_	    = tokenizer(UpperCamelCase__  )

         for key in encoded_tok.keys():
                self.assertListEqual(encoded_tok[key]      ,  encoded_processor[key]  )







  def       _a					( self  ):






         """simple docstring"""


         a_	    = self.get_feature_extractor()
         a_	    = self.get_tokenizer()

         a_	    = WavaVecaProcessor(tokenizer=UpperCamelCase__      ,  feature_extractor=UpperCamelCase__  )

         a_	    = [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

         a_	    = processor.batch_decode(UpperCamelCase__  )
         a_	    = tokenizer.batch_decode(UpperCamelCase__  )

         self.assertListEqual(UpperCamelCase__      ,  UpperCamelCase__  )







  def       _a					( self  ):






         """simple docstring"""


         a_	    = self.get_feature_extractor()
         a_	    = self.get_tokenizer()

         a_	    = WavaVecaProcessor(tokenizer=UpperCamelCase__      ,  feature_extractor=UpperCamelCase__  )

         self.assertListEqual(
             processor.model_input_names      ,  feature_extractor.model_input_names      ,  msg='`processor` and `feature_extractor` model input names do not match'      ,  )




