
'''simple docstring'''





import logging
import sys
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union

import librosa
import torch
from datasets import DatasetDict, load_dataset
from packaging import version
from torch import nn

from transformers import (
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaForPreTraining,
    is_apex_available,
    trainer_utils,
)
from transformers.models.wavaveca.modeling_wavaveca import _compute_mask_indices


if is_apex_available():
     from apex import amp

if version.parse(version.parse(torch.__version__).base_version) >= version.parse('''1.6'''):
     __magic_name__            =  True
     from torch.cuda.amp import autocast


__magic_name__            =  logging.getLogger(__name__)




@dataclass
class 					_SCREAMING_SNAKE_CASE  :
      _A						:		str    	=     field(
          metadata={'help': 'Path to pretrained model or model identifier from huggingface.co/models'} )
      _A						:		Optional[str]    	=     field(
          default=__UpperCamelCase       ,					metadata={'help': 'Where do you want to store the pretrained models downloaded from huggingface.co'}       ,					)
      _A						:		Optional[bool]    	=     field(
          default=__UpperCamelCase       ,					metadata={'help': 'Whether to freeze the feature extractor layers of the model.'} )
      _A						:		Optional[bool]    	=     field(
          default=__UpperCamelCase       ,					metadata={'help': 'Whether to log verbose messages or not.'}       ,					)
      _A						:		Optional[float]    	=     field(
          default=2.0       ,					metadata={'help': 'Maximum temperature for gumbel softmax.'} )
      _A						:		Optional[float]    	=     field(
          default=0.5       ,					metadata={'help': 'Minimum temperature for gumbel softmax.'} )
      _A						:		Optional[float]    	=     field(
          default=0.9_9_9_9_9_5       ,					metadata={'help': 'Decay of gumbel temperature during training.'} )





def 	SCREAMING_SNAKE_CASE__    (      __lowerCAmelCase    ,		__lowerCAmelCase							):
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s"    ,		datefmt="%m/%d/%Y %H:%M:%S"    ,		handlers=[logging.StreamHandler(sys.stdout							)]    ,		)
    snake_case__						=				logging.WARNING
    if model_args.verbose_logging:
        snake_case__						=				logging.DEBUG
    elif trainer_utils.is_main_process(training_args.local_rank							):
        snake_case__						=				logging.INFO
    logger.setLevel(__lowerCAmelCase							)




@dataclass
class 					_SCREAMING_SNAKE_CASE  :
      _A						:		str    	=     field(
          default=__UpperCamelCase       ,					metadata={'help': 'The name of the dataset to use (via the datasets library).'} )
      _A						:		Optional[str]    	=     field(
          default=__UpperCamelCase       ,					metadata={'help': 'The configuration name of the dataset to use (via the datasets library).'} )
      _A						:		Optional[str]    	=     field(
          default='train'       ,					metadata={
              'help': 'The name of the training data set split to use (via the datasets library). Defaults to \'train\''
          }       ,					)
      _A						:		Optional[str]    	=     field(
          default='validation'       ,					metadata={
              'help': (
                  'The name of the validation data set split to use (via the datasets library). Defaults to \'validation\''
              )
          }       ,					)
      _A						:		Optional[str]    	=     field(
          default='file'       ,					metadata={'help': 'Column in the dataset that contains speech file path. Defaults to \'file\''}       ,					)
      _A						:		bool    	=     field(
          default=__UpperCamelCase       ,					metadata={'help': 'Overwrite the cached preprocessed datasets or not.'} )
      _A						:		Optional[int]    	=     field(
          default=1       ,					metadata={
              'help': 'The percentage of the train set used as validation set in case there\'s no validation split'
          }       ,					)
      _A						:		Optional[int]    	=     field(
          default=__UpperCamelCase       ,					metadata={'help': 'The number of processes to use for the preprocessing.'}       ,					)
      _A						:		Optional[float]    	=     field(
          default=20.0       ,					metadata={'help': 'Filter audio files that are longer than `max_duration_in_seconds` seconds'} )





@dataclass
class 					_SCREAMING_SNAKE_CASE  :
      _A						:		WavaVecaForPreTraining
      _A						:		WavaVecaFeatureExtractor
      _A						:		Union[bool, str]    	=     "longest"
      _A						:		Optional[int]    	=     None
      _A						:		Optional[int]    	=     None
      def __call__(			self		,	lowerCamelCase ):
          # reformat list to dict and set to pytorch format
          snake_case__						=				self.feature_extractor.pad(
              lowerCamelCase		,	max_length=self.max_length		,	padding=self.padding		,	pad_to_multiple_of=self.pad_to_multiple_of		,	return_tensors="pt"		,	)
          snake_case__						=				self.model._get_feat_extract_output_lengths(batch["input_values"].shape[-1] )

          snake_case__						=				batch["input_values"].shape[0]

          # make sure that no loss is computed on padded inputs
          if batch["attention_mask"] is not None:
              # compute real output lengths according to convolution formula
              snake_case__						=				self.model._get_feat_extract_output_lengths(batch["attention_mask"].sum(-1 ) ).to(
                  torch.long )

              snake_case__						=				torch.zeros(
                  (batch_size, mask_indices_seq_length)		,	dtype=torch.long		,	device=batch["input_values"].device )

              # these two operations makes sure that all values
              # before the output lengths indices are attended to
              snake_case__						=				1
              snake_case__						=				attention_mask.flip([-1] ).cumsum(-1 ).flip([-1] ).bool()

          # sample randomly masked indices
          snake_case__						=				_compute_mask_indices(
              (batch_size, mask_indices_seq_length)		,	self.model.config.mask_time_prob		,	self.model.config.mask_time_length		,	attention_mask=lowerCamelCase		,	min_masks=2		,	)

          return batch





class 					_SCREAMING_SNAKE_CASE  (      __UpperCamelCase ):
      def __init__(			self		,	*lowerCamelCase		,	lowerCamelCase=1		,	lowerCamelCase=0		,	lowerCamelCase=1.0		,	**lowerCamelCase ):
          super().__init__(*lowerCamelCase		,	**lowerCamelCase )
          snake_case__						=				0
          snake_case__						=				max_gumbel_temp
          snake_case__						=				min_gumbel_temp
          snake_case__						=				gumbel_temp_decay



      def A_			(			self		,	lowerCamelCase		,	lowerCamelCase ):
          model.train()
          snake_case__						=				self._prepare_inputs(lowerCamelCase )

          if self.use_amp:
              with autocast():
                  snake_case__						=				self.compute_loss(lowerCamelCase		,	lowerCamelCase )
          else:
              snake_case__						=				self.compute_loss(lowerCamelCase		,	lowerCamelCase )

          if self.args.n_gpu > 1 or self.deepspeed:
              if model.module.config.ctc_loss_reduction == "mean":
                  snake_case__						=				loss.mean()
              elif model.module.config.ctc_loss_reduction == "sum":
                  snake_case__						=				loss.sum() / (inputs["mask_time_indices"]).sum()
              else:
                  raise ValueError(F"""{model.config.ctc_loss_reduction} is not valid. Choose one of ['mean', 'sum']""" )

          if self.args.gradient_accumulation_steps > 1:
              snake_case__						=				loss / self.args.gradient_accumulation_steps

          if self.use_amp:
              self.scaler.scale(lowerCamelCase ).backward()
          elif self.use_apex:
              with amp.scale_loss(lowerCamelCase		,	self.optimizer ) as scaled_loss:
                  scaled_loss.backward()
          elif self.deepspeed:
              self.deepspeed.backward(lowerCamelCase )
          else:
              loss.backward()

          self.num_update_step += 1
          # make sure gumbel softmax temperature is decayed
          if self.args.n_gpu > 1 or self.deepspeed:
              model.module.set_gumbel_temperature(
                  max(self.max_gumbel_temp * self.gumbel_temp_decay**self.num_update_step		,	self.min_gumbel_temp ) )
          else:
              model.set_gumbel_temperature(
                  max(self.max_gumbel_temp * self.gumbel_temp_decay**self.num_update_step		,	self.min_gumbel_temp ) )

          return loss.detach()





def 	SCREAMING_SNAKE_CASE__    (      ):
    # See all possible arguments in src/transformers/training_args.py
    # or by passing the --help flag to this script.
    # We now keep distinct sets of args, for a cleaner separation of concerns.

    snake_case__						=				HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)							)

    snake_case__						,					snake_case__						,					snake_case__						=				parser.parse_args_into_dataclasses()
    configure_logger(__lowerCAmelCase    ,		__lowerCAmelCase							)

    # Downloading and loading a dataset from the hub.
    snake_case__						=				load_dataset(data_args.dataset_name    ,		data_args.dataset_config_name    ,		cache_dir=model_args.cache_dir							)

    if "validation" not in datasets.keys():
        # make sure only "validation" and "train" keys remain"
        snake_case__						=				DatasetDict()
        snake_case__						=				load_dataset(
            data_args.dataset_name    ,		data_args.dataset_config_name    ,		split=F"""{data_args.train_split_name}[:{data_args.validation_split_percentage}%]"""    ,		cache_dir=model_args.cache_dir    ,		)
        snake_case__						=				load_dataset(
            data_args.dataset_name    ,		data_args.dataset_config_name    ,		split=F"""{data_args.train_split_name}[{data_args.validation_split_percentage}%:]"""    ,		cache_dir=model_args.cache_dir    ,		)
    else:
        # make sure only "validation" and "train" keys remain"
        snake_case__						=				DatasetDict()
        snake_case__						=				load_dataset(
            data_args.dataset_name    ,		data_args.dataset_config_name    ,		split="validation"    ,		cache_dir=model_args.cache_dir    ,		)
        snake_case__						=				load_dataset(
            data_args.dataset_name    ,		data_args.dataset_config_name    ,		split=F"""{data_args.train_split_name}"""    ,		cache_dir=model_args.cache_dir    ,		)

    # only normalized-inputs-training is supported
    snake_case__						=				WavaVecaFeatureExtractor.from_pretrained(
        model_args.model_name_or_path    ,		cache_dir=model_args.cache_dir    ,		do_normalize=__lowerCAmelCase							)

    def prepare_dataset(__lowerCAmelCase							):
        # check that all files have the correct sampling rate
        snake_case__						,					snake_case__						=				librosa.load(batch[data_args.speech_file_column]    ,		sr=feature_extractor.sampling_rate							)
        return batch

    # load audio files into numpy arrays
    snake_case__						=				datasets.map(
        __lowerCAmelCase    ,		num_proc=data_args.preprocessing_num_workers    ,		remove_columns=datasets["train"].column_names							)

    # filter audio files that are too long
    snake_case__						=				vectorized_datasets.filter(
        lambda __lowerCAmelCase							: len(data["speech"]							) < int(data_args.max_duration_in_seconds * feature_extractor.sampling_rate							)							)

    def normalize(__lowerCAmelCase							):
        return feature_extractor(batch["speech"]    ,		sampling_rate=feature_extractor.sampling_rate							)

    # normalize and transform to `BatchFeatures`
    snake_case__						=				vectorized_datasets.map(
        __lowerCAmelCase    ,		batched=__lowerCAmelCase    ,		num_proc=data_args.preprocessing_num_workers    ,		load_from_cache_file=not data_args.overwrite_cache    ,		remove_columns=vectorized_datasets["train"].column_names    ,		)

    # pretraining is only supported for "newer" stable layer norm architecture
    # apply_spec_augment has to be True, mask_feature_prob has to be 0.0
    snake_case__						=				WavaVecaConfig.from_pretrained(
        model_args.model_name_or_path    ,		cache_dir=model_args.cache_dir    ,		gradient_checkpointing=training_args.gradient_checkpointing    ,		)

    if not config.do_stable_layer_norm or config.feat_extract_norm != "layer":
        raise ValueError(
            "PreTraining is only supported for ``config.do_stable_layer_norm=True`` and"
            " ``config.feat_extract_norm='layer'"							)

    snake_case__						=				WavaVecaForPreTraining(__lowerCAmelCase							)

    snake_case__						=				DataCollatorForWavaVecaPretraining(model=__lowerCAmelCase    ,		feature_extractor=__lowerCAmelCase							)

    snake_case__						=				WavaVecaPreTrainer(
        model=__lowerCAmelCase    ,		data_collator=__lowerCAmelCase    ,		args=__lowerCAmelCase    ,		train_dataset=vectorized_datasets["train"]    ,		eval_dataset=vectorized_datasets["validation"]    ,		tokenizer=__lowerCAmelCase    ,		max_gumbel_temp=model_args.max_gumbel_temperature    ,		min_gumbel_temp=model_args.min_gumbel_temperature    ,		gumbel_temp_decay=model_args.gumbel_temperature_decay    ,		)
    trainer.train()


if __name__ == "__main__":
     main()





import os
from shutil import copyfile
from typing import List, Optional, Tuple

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
     from .tokenization_pegasus import PegasusTokenizer
else:
     __magic_name__            =  None


__magic_name__            =  logging.get_logger(__name__)


__magic_name__            =  '''▁'''

__magic_name__            =  {'''vocab_file''': '''spiece.model''', '''tokenizer_file''': '''tokenizer.json'''}

__magic_name__            =  {
    '''vocab_file''': {'''google/pegasus-xsum''': '''https://huggingface.co/google/pegasus-xsum/resolve/main/spiece.model'''},
    '''tokenizer_file''': {
        '''google/pegasus-xsum''': '''https://huggingface.co/google/pegasus-xsum/resolve/main/tokenizer.json'''
    },
}

__magic_name__            =  {
    '''google/pegasus-xsum''': 512,
}




class 					_SCREAMING_SNAKE_CASE  (      __UpperCamelCase ):
      _A						:		int    	=     VOCAB_FILES_NAMES
      _A						:		Dict    	=     PRETRAINED_VOCAB_FILES_MAP
      _A						:		Any    	=     PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      _A						:		Tuple    	=     PegasusTokenizer
      _A						:		str    	=     ['input_ids', 'attention_mask']
      def __init__(			self		,	lowerCamelCase=None		,	lowerCamelCase=None		,	lowerCamelCase="<pad>"		,	lowerCamelCase="</s>"		,	lowerCamelCase="<unk>"		,	lowerCamelCase="<mask_2>"		,	lowerCamelCase="<mask_1>"		,	lowerCamelCase=None		,	lowerCamelCase=1_03		,	**lowerCamelCase		,	):
          snake_case__						=				offset

          if additional_special_tokens is not None:
              if not isinstance(lowerCamelCase		,	lowerCamelCase ):
                  raise TypeError(
                      F"""additional_special_tokens should be of type {type(lowerCamelCase )}, but is"""
                      F""" {type(lowerCamelCase )}""" )

              snake_case__						=				(
                  ([mask_token_sent] + additional_special_tokens)
                  if mask_token_sent not in additional_special_tokens and mask_token_sent is not None
                  else additional_special_tokens
              )
              # fill additional tokens with ..., <unk_token_102> in case not all additional tokens are already taken
              additional_special_tokens_extended += [
                  F"""<unk_{i}>""" for i in range(len(lowerCamelCase )		,	self.offset - 1 )
              ]

              if len(set(lowerCamelCase ) ) != len(lowerCamelCase ):
                  raise ValueError(
                      "Please make sure that the provided additional_special_tokens do not contain an incorrectly"
                      F""" shifted list of <unk_x> tokens. Found {additional_special_tokens_extended}.""" )
              snake_case__						=				additional_special_tokens_extended
          else:
              snake_case__						=				[mask_token_sent] if mask_token_sent is not None else []
              additional_special_tokens += [F"""<unk_{i}>""" for i in range(2		,	self.offset )]

          super().__init__(
              lowerCamelCase		,	tokenizer_file=lowerCamelCase		,	pad_token=lowerCamelCase		,	eos_token=lowerCamelCase		,	unk_token=lowerCamelCase		,	mask_token=lowerCamelCase		,	mask_token_sent=lowerCamelCase		,	offset=lowerCamelCase		,	additional_special_tokens=lowerCamelCase		,	**lowerCamelCase		,	)
          snake_case__						=				vocab_file
          snake_case__						=				False if not self.vocab_file else True
      def A_			(			self		,	lowerCamelCase ):
          snake_case__						=				set(self.all_special_ids )  # call it once instead of inside list comp
          all_special_ids.remove(self.unk_token_id )  # <unk> is only sometimes special

          if all_special_ids != set(range(len(self.additional_special_tokens ) + 3 ) ):
              raise ValueError(
                  "There should be 3 special tokens: mask_token, pad_token, and eos_token +"
                  F""" {len(self.additional_special_tokens )} additional_special_tokens, but got {all_special_ids}""" )

          return [1 if x in all_special_ids else 0 for x in seq]
      def A_			(			self		,	lowerCamelCase		,	lowerCamelCase = None		,	lowerCamelCase = False ):
          if already_has_special_tokens:
              return self._special_token_mask(lowerCamelCase )
          elif token_ids_a is None:
              return self._special_token_mask(lowerCamelCase ) + [1]
          else:
              return self._special_token_mask(token_ids_a + token_ids_a ) + [1]
      def A_			(			self		,	lowerCamelCase		,	lowerCamelCase=None ):
          if token_ids_a is None:
              return token_ids_a + [self.eos_token_id]
          # We don't expect to process pairs, but leave the pair logic for API consistency
          return token_ids_a + token_ids_a + [self.eos_token_id]



      def A_			(			self		,	lowerCamelCase		,	lowerCamelCase = None ):
          if not self.can_save_slow_tokenizer:
              raise ValueError(
                  "Your fast tokenizer does not have the necessary information to save the vocabulary for a slow "
                  "tokenizer." )

          if not os.path.isdir(lowerCamelCase ):
              logger.error(F"""Vocabulary path ({save_directory}) should be a directory""" )
              return
          snake_case__						=				os.path.join(
              lowerCamelCase		,	(filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"] )

          if os.path.abspath(self.vocab_file ) != os.path.abspath(lowerCamelCase ):
              copyfile(self.vocab_file		,	lowerCamelCase )

          return (out_vocab_file,)



