



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'''):
   __lowerCamelCase         =			True
   from torch.cuda.amp import autocast


__lowerCamelCase         =			logging.getLogger(__name__)


@dataclass
class   snake_case_ :




       """simple docstring"""






       _lowerCamelCase			    =				field(
           metadata={"""help""": """Path to pretrained model or model identifier from huggingface.co/models"""}  )
       _lowerCamelCase			    =				field(
           default=lowercase__ , metadata={"""help""": """Where do you want to store the pretrained models downloaded from huggingface.co"""} , )
       _lowerCamelCase			    =				field(
           default=lowercase__ , metadata={"""help""": """Whether to freeze the feature extractor layers of the model."""}  )
       _lowerCamelCase			    =				field(
           default=lowercase__ , metadata={"""help""": """Whether to log verbose messages or not."""} , )
       _lowerCamelCase			    =				field(
           default=2.0 , metadata={"""help""": """Maximum temperature for gumbel softmax."""}  )
       _lowerCamelCase			    =				field(
           default=0.5 , metadata={"""help""": """Minimum temperature for gumbel softmax."""}  )
       _lowerCamelCase			    =				field(
           default=0.99_99_95 , metadata={"""help""": """Decay of gumbel temperature during training."""}  )







def    _snake_case (       __snake_case					,  __snake_case    )     ->     int:






      '''simple docstring'''







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


@dataclass
class   snake_case_ :




       """simple docstring"""






       _lowerCamelCase			    =				field(
           default=lowercase__ , metadata={"""help""": """The name of the dataset to use (via the datasets library)."""}  )
       _lowerCamelCase			    =				field(
           default=lowercase__ , metadata={"""help""": """The configuration name of the dataset to use (via the datasets library)."""}  )
       _lowerCamelCase			    =				field(
           default="""train""" , metadata={
               """help""": """The name of the training data set split to use (via the datasets library). Defaults to 'train'"""
           } , )
       _lowerCamelCase			    =				field(
           default="""validation""" , metadata={
               """help""": (
                   """The name of the validation data set split to use (via the datasets library). Defaults to 'validation'"""
               )
           } , )
       _lowerCamelCase			    =				field(
           default="""file""" , metadata={"""help""": """Column in the dataset that contains speech file path. Defaults to 'file'"""} , )
       _lowerCamelCase			    =				field(
           default=lowercase__ , metadata={"""help""": """Overwrite the cached preprocessed datasets or not."""}  )
       _lowerCamelCase			    =				field(
           default=1 , metadata={
               """help""": """The percentage of the train set used as validation set in case there's no validation split"""
           } , )
       _lowerCamelCase			    =				field(
           default=lowercase__ , metadata={"""help""": """The number of processes to use for the preprocessing."""} , )
       _lowerCamelCase			    =				field(
           default=20.0 , metadata={"""help""": """Filter audio files that are longer than `max_duration_in_seconds` seconds"""}  )


@dataclass
class   snake_case_ :




       """simple docstring"""






       _lowerCamelCase			    =				42
       _lowerCamelCase			    =				42
       _lowerCamelCase			    =				"longest"
       _lowerCamelCase			    =				None
       _lowerCamelCase			    =				None



       def __call__(     self     ,lowercase):

             """simple docstring"""





             UpperCAmelCase_  :	int				      =		self.feature_extractor.pad(
                 lowercase     ,max_length=self.max_length     ,padding=self.padding     ,pad_to_multiple_of=self.pad_to_multiple_of     ,return_tensors="pt"     ,)
             UpperCAmelCase_  :	Union[str, Any]				      =		self.model._get_feat_extract_output_lengths(batch["input_values"].shape[-1])

             UpperCAmelCase_  :	Dict				      =		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
                   UpperCAmelCase_  :	Union[str, Any]				      =		self.model._get_feat_extract_output_lengths(batch["attention_mask"].sum(-1)).to(
                       torch.long)

                   UpperCAmelCase_  :	List[Any]				      =		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
                   UpperCAmelCase_  :	Tuple				      =		1
                   UpperCAmelCase_  :	List[str]				      =		attention_mask.flip([-1]).cumsum(-1).flip([-1]).bool()

             # sample randomly masked indices
             UpperCAmelCase_  :	str				      =		_compute_mask_indices(
                 (batch_size, mask_indices_seq_length)     ,self.model.config.mask_time_prob     ,self.model.config.mask_time_length     ,attention_mask=lowercase     ,min_masks=2     ,)

             return batch


class   snake_case_ (lowercase__  ):




       """simple docstring"""






       def __init__(     self     ,*lowercase     ,lowercase=1     ,lowercase=0     ,lowercase=1.0     ,**lowercase):

             """simple docstring"""





             super().__init__(*lowercase     ,**lowercase)
             UpperCAmelCase_  :	List[Any]				      =		0
             UpperCAmelCase_  :	Dict				      =		max_gumbel_temp
             UpperCAmelCase_  :	Union[str, Any]				      =		min_gumbel_temp
             UpperCAmelCase_  :	str				      =		gumbel_temp_decay



       def  A_   (     self     ,lowercase     ,lowercase):

             """simple docstring"""





             model.train()
             UpperCAmelCase_  :	int				      =		self._prepare_inputs(lowercase)

             if self.use_amp:
                   with autocast():
                         UpperCAmelCase_  :	List[str]				      =		self.compute_loss(lowercase     ,lowercase)
             else:
                   UpperCAmelCase_  :	List[str]				      =		self.compute_loss(lowercase     ,lowercase)

             if self.args.n_gpu > 1 or self.deepspeed:
                   if model.module.config.ctc_loss_reduction == "mean":
                         UpperCAmelCase_  :	Tuple				      =		loss.mean()
                   elif model.module.config.ctc_loss_reduction == "sum":
                         UpperCAmelCase_  :	Optional[int]				      =		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:
                   UpperCAmelCase_  :	Tuple				      =		loss / self.args.gradient_accumulation_steps

             if self.use_amp:
                   self.scaler.scale(lowercase).backward()
             elif self.use_apex:
                   with amp.scale_loss(lowercase     ,self.optimizer) as scaled_loss:
                         scaled_loss.backward()
             elif self.deepspeed:
                   self.deepspeed.backward(lowercase)
             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    _snake_case (       )     ->     Union[str, Any]:






      '''simple docstring'''







      UpperCAmelCase_  :	Optional[Any]				      =		HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)    )

      UpperCAmelCase_      ,      UpperCAmelCase_      ,      UpperCAmelCase_  :	Any				      =		parser.parse_args_into_dataclasses()
      configure_logger(__snake_case					,  __snake_case    )

      # Downloading and loading a dataset from the hub.
      UpperCAmelCase_  :	Any				      =		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"
            UpperCAmelCase_  :	Tuple				      =		DatasetDict()
            UpperCAmelCase_  :	Dict				      =		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					,  )
            UpperCAmelCase_  :	int				      =		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"
            UpperCAmelCase_  :	Optional[int]				      =		DatasetDict()
            UpperCAmelCase_  :	Any				      =		load_dataset(
                data_args.dataset_name					,  data_args.dataset_config_name					,  split="validation"					,  cache_dir=model_args.cache_dir					,  )
            UpperCAmelCase_  :	List[Any]				      =		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
      UpperCAmelCase_  :	Any				      =		WavaVecaFeatureExtractor.from_pretrained(
          model_args.model_name_or_path					,  cache_dir=model_args.cache_dir					,  do_normalize=__snake_case    )

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

      # load audio files into numpy arrays
      UpperCAmelCase_  :	Optional[int]				      =		datasets.map(
          __snake_case					,  num_proc=data_args.preprocessing_num_workers					,  remove_columns=datasets["train"].column_names    )

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

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

      # normalize and transform to `BatchFeatures`
      UpperCAmelCase_  :	str				      =		vectorized_datasets.map(
          __snake_case					,  batched=__snake_case					,  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
      UpperCAmelCase_  :	Any				      =		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'"    )

      UpperCAmelCase_  :	str				      =		WavaVecaForPreTraining(__snake_case    )

      UpperCAmelCase_  :	str				      =		DataCollatorForWavaVecaPretraining(model=__snake_case					,  feature_extractor=__snake_case    )

      UpperCAmelCase_  :	int				      =		WavaVecaPreTrainer(
          model=__snake_case					,  data_collator=__snake_case					,  args=__snake_case					,  train_dataset=vectorized_datasets["train"]					,  eval_dataset=vectorized_datasets["validation"]					,  tokenizer=__snake_case					,  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()





from abc import ABC, abstractmethod
from typing import Optional, Union

from .. import Dataset, DatasetDict, Features, IterableDataset, IterableDatasetDict, NamedSplit
from ..utils.typing import NestedDataStructureLike, PathLike


class   snake_case_ (lowercase__  ):




       """simple docstring"""






       def __init__(     self     ,lowercase = None     ,lowercase = None     ,lowercase = None     ,lowercase = None     ,lowercase = False     ,lowercase = False     ,lowercase = None     ,**lowercase     ,):

             """simple docstring"""





             UpperCAmelCase_  :	Optional[Any]				      =		path_or_paths
             UpperCAmelCase_  :	Optional[int]				      =		split if split or isinstance(lowercase     ,lowercase) else "train"
             UpperCAmelCase_  :	Optional[int]				      =		features
             UpperCAmelCase_  :	Dict				      =		cache_dir
             UpperCAmelCase_  :	int				      =		keep_in_memory
             UpperCAmelCase_  :	Tuple				      =		streaming
             UpperCAmelCase_  :	Union[str, Any]				      =		num_proc
             UpperCAmelCase_  :	Union[str, Any]				      =		kwargs



       @abstractmethod
       def  A_   (     self):

             """simple docstring"""





             pass


class   snake_case_ (lowercase__  ):




       """simple docstring"""






       def __init__(     self     ,lowercase = None     ,lowercase = None     ,lowercase = False     ,lowercase = False     ,lowercase = None     ,**lowercase     ,):

             """simple docstring"""





             UpperCAmelCase_  :	List[Any]				      =		features
             UpperCAmelCase_  :	List[Any]				      =		cache_dir
             UpperCAmelCase_  :	List[str]				      =		keep_in_memory
             UpperCAmelCase_  :	Optional[int]				      =		streaming
             UpperCAmelCase_  :	Dict				      =		num_proc
             UpperCAmelCase_  :	str				      =		kwargs



       @abstractmethod
       def  A_   (     self):

             """simple docstring"""





             pass
