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


UpperCamelCase     :    Union[str, Any]       =			logging.getLogger(__name__)

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






def 							A						(    snake_case				:ModelArguments    , snake_case				:TrainingArguments		)				->     str:
 logging.basicConfig(
     format='%(asctime)s - %(levelname)s - %(name)s - %(message)s'    , datefmt='%m/%d/%Y %H:%M:%S'    , handlers=[logging.StreamHandler(sys.stdout		)]    , )
 __UpperCamelCase      =	logging.WARNING
 if model_args.verbose_logging:
  __UpperCamelCase      =	logging.DEBUG
 elif trainer_utils.is_main_process(training_args.local_rank		):
  __UpperCamelCase      =	logging.INFO
 logger.setLevel(snake_case		)

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


@dataclass
class __lowerCAmelCase		:
     lowercase						       =						42
     lowercase						       =						42
     lowercase						       =						"longest"
     lowercase						       =						None
     lowercase						       =						None





     def __call__(     self     ,					__UpperCAmelCase				):

      '''simple docstring'''




      __UpperCamelCase      =	self.feature_extractor.pad(
          __UpperCAmelCase     ,					max_length=self.max_length     ,					padding=self.padding     ,					pad_to_multiple_of=self.pad_to_multiple_of     ,					return_tensors='pt'     ,					)
      __UpperCamelCase      =	self.model._get_feat_extract_output_lengths(batch['input_values'].shape[-1]				)

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

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

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

      return batch


class __lowerCAmelCase		(				__SCREAMING_SNAKE_CASE      ):





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

      '''simple docstring'''




      super().__init__(*__UpperCAmelCase     ,					**__UpperCAmelCase				)
      __UpperCamelCase      =	0
      __UpperCamelCase      =	max_gumbel_temp
      __UpperCamelCase      =	min_gumbel_temp
      __UpperCamelCase      =	gumbel_temp_decay







     def 						UpperCAmelCase (     self     ,					__UpperCAmelCase     ,					__UpperCAmelCase				):

      '''simple docstring'''




      model.train()
      __UpperCamelCase      =	self._prepare_inputs(__UpperCAmelCase				)

      if self.use_amp:
       with autocast():
        __UpperCamelCase      =	self.compute_loss(__UpperCAmelCase     ,					__UpperCAmelCase				)
      else:
       __UpperCamelCase      =	self.compute_loss(__UpperCAmelCase     ,					__UpperCAmelCase				)

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

      if self.use_amp:
       self.scaler.scale(__UpperCAmelCase				).backward()
      elif self.use_apex:
       with amp.scale_loss(__UpperCAmelCase     ,					self.optimizer				) as scaled_loss:
        scaled_loss.backward()
      elif self.deepspeed:
       self.deepspeed.backward(__UpperCAmelCase				)
      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						(    )				->     int:
 # 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.

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

 __UpperCamelCase  ,     __UpperCamelCase  ,     __UpperCamelCase      =	parser.parse_args_into_dataclasses()
 configure_logger(snake_case    , snake_case		)

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

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

 # load audio files into numpy arrays
 __UpperCamelCase      =	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      =	vectorized_datasets.filter(
     lambda snake_case		: len(data['speech']		) < int(data_args.max_duration_in_seconds * feature_extractor.sampling_rate		)		)

 def normalize(snake_case				:Union[str, Any]		):
  return feature_extractor(batch['speech']    , sampling_rate=feature_extractor.sampling_rate		)

 # normalize and transform to `BatchFeatures`
 __UpperCamelCase      =	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      =	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      =	WavaVecaForPreTraining(snake_case		)

 __UpperCamelCase      =	DataCollatorForWavaVecaPretraining(model=snake_case    , feature_extractor=snake_case		)

 __UpperCamelCase      =	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()





"""simple docstring"""

import functools
import operator

from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCamelCase     :    str       =			logging.get_logger(__name__)

UpperCamelCase     :    Dict       =			{
    "microsoft/wavlm-base": "https://huggingface.co/microsoft/wavlm-base/resolve/main/config.json",
    # See all WavLM models at https://huggingface.co/models?filter=wavlm
}

class __lowerCAmelCase		(				__SCREAMING_SNAKE_CASE      ):
     lowercase						       =						"wavlm"





     def __init__(     self     ,					__UpperCAmelCase=32     ,					__UpperCAmelCase=768     ,					__UpperCAmelCase=12     ,					__UpperCAmelCase=12     ,					__UpperCAmelCase=3072     ,					__UpperCAmelCase="gelu"     ,					__UpperCAmelCase=0.1     ,					__UpperCAmelCase=0.1     ,					__UpperCAmelCase=0.1     ,					__UpperCAmelCase=0.0     ,					__UpperCAmelCase=0.1     ,					__UpperCAmelCase=0.1     ,					__UpperCAmelCase=0.0_2     ,					__UpperCAmelCase=1E-5     ,					__UpperCAmelCase="group"     ,					__UpperCAmelCase="gelu"     ,					__UpperCAmelCase=(512, 512, 512, 512, 512, 512, 512)     ,					__UpperCAmelCase=(5, 2, 2, 2, 2, 2, 2)     ,					__UpperCAmelCase=(10, 3, 3, 3, 3, 2, 2)     ,					__UpperCAmelCase=False     ,					__UpperCAmelCase=128     ,					__UpperCAmelCase=16     ,					__UpperCAmelCase=320     ,					__UpperCAmelCase=800     ,					__UpperCAmelCase=False     ,					__UpperCAmelCase=True     ,					__UpperCAmelCase=0.0_5     ,					__UpperCAmelCase=10     ,					__UpperCAmelCase=2     ,					__UpperCAmelCase=0.0     ,					__UpperCAmelCase=10     ,					__UpperCAmelCase=320     ,					__UpperCAmelCase=2     ,					__UpperCAmelCase=0.1     ,					__UpperCAmelCase=100     ,					__UpperCAmelCase=256     ,					__UpperCAmelCase=256     ,					__UpperCAmelCase=0.1     ,					__UpperCAmelCase="mean"     ,					__UpperCAmelCase=False     ,					__UpperCAmelCase=False     ,					__UpperCAmelCase=256     ,					__UpperCAmelCase=(512, 512, 512, 512, 1500)     ,					__UpperCAmelCase=(5, 3, 3, 1, 1)     ,					__UpperCAmelCase=(1, 2, 3, 1, 1)     ,					__UpperCAmelCase=512     ,					__UpperCAmelCase=80     ,					__UpperCAmelCase=0     ,					__UpperCAmelCase=1     ,					__UpperCAmelCase=2     ,					__UpperCAmelCase=False     ,					__UpperCAmelCase=3     ,					__UpperCAmelCase=2     ,					__UpperCAmelCase=3     ,					__UpperCAmelCase=None     ,					**__UpperCAmelCase     ,					):

      '''simple docstring'''




      super().__init__(**__UpperCAmelCase     ,					pad_token_id=__UpperCAmelCase     ,					bos_token_id=__UpperCAmelCase     ,					eos_token_id=__UpperCAmelCase				)
      __UpperCamelCase      =	hidden_size
      __UpperCamelCase      =	feat_extract_norm
      __UpperCamelCase      =	feat_extract_activation
      __UpperCamelCase      =	list(__UpperCAmelCase				)
      __UpperCamelCase      =	list(__UpperCAmelCase				)
      __UpperCamelCase      =	list(__UpperCAmelCase				)
      __UpperCamelCase      =	conv_bias
      __UpperCamelCase      =	num_buckets
      __UpperCamelCase      =	max_bucket_distance
      __UpperCamelCase      =	num_conv_pos_embeddings
      __UpperCamelCase      =	num_conv_pos_embedding_groups
      __UpperCamelCase      =	len(self.conv_dim				)
      __UpperCamelCase      =	num_hidden_layers
      __UpperCamelCase      =	intermediate_size
      __UpperCamelCase      =	hidden_act
      __UpperCamelCase      =	num_attention_heads
      __UpperCamelCase      =	hidden_dropout
      __UpperCamelCase      =	attention_dropout
      __UpperCamelCase      =	activation_dropout
      __UpperCamelCase      =	feat_proj_dropout
      __UpperCamelCase      =	final_dropout
      __UpperCamelCase      =	layerdrop
      __UpperCamelCase      =	layer_norm_eps
      __UpperCamelCase      =	initializer_range
      __UpperCamelCase      =	num_ctc_classes
      __UpperCamelCase      =	vocab_size
      __UpperCamelCase      =	do_stable_layer_norm
      __UpperCamelCase      =	use_weighted_layer_sum
      __UpperCamelCase      =	classifier_proj_size

      if (
          (len(self.conv_stride				) != self.num_feat_extract_layers)
          or (len(self.conv_kernel				) != self.num_feat_extract_layers)
          or (len(self.conv_dim				) != self.num_feat_extract_layers)
      ):
       raise ValueError(
           'Configuration for convolutional layers is incorrect. It is required that `len(config.conv_dim)` =='
           ' `len(config.conv_stride)` == `len(config.conv_kernel)`, but is `len(config.conv_dim) ='
           F' {len(self.conv_dim				)}`, `len(config.conv_stride) = {len(self.conv_stride				)}`,'
           F' `len(config.conv_kernel) = {len(self.conv_kernel				)}`.'				)

      # fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
      __UpperCamelCase      =	apply_spec_augment
      __UpperCamelCase      =	mask_time_prob
      __UpperCamelCase      =	mask_time_length
      __UpperCamelCase      =	mask_time_min_masks
      __UpperCamelCase      =	mask_feature_prob
      __UpperCamelCase      =	mask_feature_length

      # parameters for pretraining with codevector quantized representations
      __UpperCamelCase      =	num_codevectors_per_group
      __UpperCamelCase      =	num_codevector_groups
      __UpperCamelCase      =	contrastive_logits_temperature
      __UpperCamelCase      =	num_negatives
      __UpperCamelCase      =	codevector_dim
      __UpperCamelCase      =	proj_codevector_dim
      __UpperCamelCase      =	diversity_loss_weight

      # ctc loss
      __UpperCamelCase      =	ctc_loss_reduction
      __UpperCamelCase      =	ctc_zero_infinity

      # adapter
      __UpperCamelCase      =	add_adapter
      __UpperCamelCase      =	adapter_kernel_size
      __UpperCamelCase      =	adapter_stride
      __UpperCamelCase      =	num_adapter_layers
      __UpperCamelCase      =	output_hidden_size or hidden_size

      # SequenceClassification-specific parameter. Feel free to ignore for other classes.
      __UpperCamelCase      =	classifier_proj_size

      # XVector-specific parameters. Feel free to ignore for other classes.
      __UpperCamelCase      =	list(__UpperCAmelCase				)
      __UpperCamelCase      =	list(__UpperCAmelCase				)
      __UpperCamelCase      =	list(__UpperCAmelCase				)
      __UpperCamelCase      =	xvector_output_dim







     @property
     def 						UpperCAmelCase (     self				):

      '''simple docstring'''




      return functools.reduce(operator.mul     ,					self.conv_stride     ,					1				)



