

"""simple docstring"""


UpperCAmelCase													="\n# Transformers installation\n! pip install transformers datasets\n# To install from source instead of the last release, comment the command above and uncomment the following one.\n# ! pip install git+https://github.com/huggingface/transformers.git\n"

UpperCAmelCase													=[{"type": "code", "content": INSTALL_CONTENT}]
UpperCAmelCase													={
    "{processor_class}": "FakeProcessorClass",
    "{model_class}": "FakeModelClass",
    "{object_class}": "FakeObjectClass",
}







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


UpperCAmelCase													=logging.getLogger(__name__)





@dataclass
class lowerCamelCase__   :



    '''simple docstring'''




    _lowerCamelCase					=   field(
        metadata={'''help''': '''Path to pretrained model or model identifier from huggingface.co/models'''}    )
    _lowerCamelCase					=   field(
        default=SCREAMING_SNAKE_CASE     ,   metadata={'''help''': '''Where do you want to store the pretrained models downloaded from huggingface.co'''}     ,   )
    _lowerCamelCase					=   field(
        default=SCREAMING_SNAKE_CASE     ,   metadata={'''help''': '''Whether to freeze the feature extractor layers of the model.'''}    )
    _lowerCamelCase					=   field(
        default=SCREAMING_SNAKE_CASE     ,   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.9_9_9_9_9_5     ,   metadata={'''help''': '''Decay of gumbel temperature during training.'''}    )




def   _A					(	_a	:						ModelArguments     ,   _a	:						TrainingArguments				):



      """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				)]     ,   )
      A				  =		logging.WARNING
      if model_args.verbose_logging:
            A				  =		logging.DEBUG
      elif trainer_utils.is_main_process(training_args.local_rank				):
            A				  =		logging.INFO
      logger.setLevel(_a				)





@dataclass
class lowerCamelCase__   :



    '''simple docstring'''




    _lowerCamelCase					=   field(
        default=SCREAMING_SNAKE_CASE     ,   metadata={'''help''': '''The name of the dataset to use (via the datasets library).'''}    )
    _lowerCamelCase					=   field(
        default=SCREAMING_SNAKE_CASE     ,   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=SCREAMING_SNAKE_CASE     ,   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=SCREAMING_SNAKE_CASE     ,   metadata={'''help''': '''The number of processes to use for the preprocessing.'''}     ,   )
    _lowerCamelCase					=   field(
        default=2_0.0     ,   metadata={'''help''': '''Filter audio files that are longer than `max_duration_in_seconds` seconds'''}    )





@dataclass
class lowerCamelCase__   :



    '''simple docstring'''




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




    def __call__( self ,lowerCamelCase_    )    ->      Dict[str, torch.Tensor]:
          # reformat list to dict and set to pytorch format
          A				  =		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""" ,)
          A				  =		self.model._get_feat_extract_output_lengths(batch["""input_values"""].shape[-1]    )

          A				  =		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
                A				  =		self.model._get_feat_extract_output_lengths(batch["""attention_mask"""].sum(-1    )    ).to(
                    torch.long    )

                A				  =		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
                A				  =		1
                A				  =		attention_mask.flip([-1]    ).cumsum(-1    ).flip([-1]    ).bool()

          # sample randomly masked indices
          A				  =		_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 lowerCamelCase__   (				SCREAMING_SNAKE_CASE    ):



    '''simple docstring'''




    def __init__( self ,*lowerCamelCase_ ,lowerCamelCase_=1 ,lowerCamelCase_=0 ,lowerCamelCase_=1.0 ,**lowerCamelCase_    )    ->      Union[str, Any]:
          super().__init__(*lowerCamelCase_ ,**lowerCamelCase_    )
          A				  =		0
          A				  =		max_gumbel_temp
          A				  =		min_gumbel_temp
          A				  =		gumbel_temp_decay




    def UpperCamelCase__		( self ,lowerCamelCase_ ,lowerCamelCase_    )    ->      torch.Tensor:
          model.train()
          A				  =		self._prepare_inputs(lowerCamelCase_    )

          if self.use_amp:
                with autocast():
                      A				  =		self.compute_loss(lowerCamelCase_ ,lowerCamelCase_    )
          else:
                A				  =		self.compute_loss(lowerCamelCase_ ,lowerCamelCase_    )

          if self.args.n_gpu > 1 or self.deepspeed:
                if model.module.config.ctc_loss_reduction == "mean":
                      A				  =		loss.mean()
                elif model.module.config.ctc_loss_reduction == "sum":
                      A				  =		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:
                A				  =		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   _A					(	):



      """simple docstring"""






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

      A  ,			A  ,			A				  =		parser.parse_args_into_dataclasses()
      configure_logger(_a     ,   _a				)

      # Downloading and loading a dataset from the hub.
      A				  =		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"
            A				  =		DatasetDict()
            A				  =		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     ,   )
            A				  =		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"
            A				  =		DatasetDict()
            A				  =		load_dataset(
                data_args.dataset_name     ,   data_args.dataset_config_name     ,   split="""validation"""     ,   cache_dir=model_args.cache_dir     ,   )
            A				  =		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
      A				  =		WavaVecaFeatureExtractor.from_pretrained(
          model_args.model_name_or_path     ,   cache_dir=model_args.cache_dir     ,   do_normalize=_a				)

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

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

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

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

      # normalize and transform to `BatchFeatures`
      A				  =		vectorized_datasets.map(
          _a     ,   batched=_a     ,   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
      A				  =		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'"""				)

      A				  =		WavaVecaForPreTraining(_a				)

      A				  =		DataCollatorForWavaVecaPretraining(model=_a     ,   feature_extractor=_a				)

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




