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


lowerCAmelCase_:					Optional[int] 	=					logging.getLogger(__name__)






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







def 					__a   (		A    ,				A      ):






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






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






@dataclass
class    a__    :
   snake_case_							=      42
   snake_case_							=      42
   snake_case_							=      "longest"
   snake_case_							=      None
   snake_case_							=      None


   def __call__(     self,				_UpperCAmelCase  ):
        '''simple docstring'''



        lowercase__				 =							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",				)
        lowercase__				 =							self.model._get_feat_extract_output_lengths(batch["input_values"].shape[-1]  )

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

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

        # sample randomly masked indices
        lowercase__				 =							_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    a__    (   _a	):


   def __init__(     self,				*_UpperCAmelCase,				_UpperCAmelCase=1,				_UpperCAmelCase=0,				_UpperCAmelCase=1.0,				**_UpperCAmelCase  ):
        '''simple docstring'''



        super().__init__(*UpperCAmelCase__,				**UpperCAmelCase__  )
        lowercase__				 =							0
        lowercase__				 =							max_gumbel_temp
        lowercase__				 =							min_gumbel_temp
        lowercase__				 =							gumbel_temp_decay


   def 	snake_case__	(     self,				_UpperCAmelCase,				_UpperCAmelCase  ):
        '''simple docstring'''



        model.train()
        lowercase__				 =							self._prepare_inputs(UpperCAmelCase__  )

        if self.use_amp:
             with autocast():
                  lowercase__				 =							self.compute_loss(UpperCAmelCase__,				UpperCAmelCase__  )
        else:
             lowercase__				 =							self.compute_loss(UpperCAmelCase__,				UpperCAmelCase__  )

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






     '''simple docstring'''



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

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

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

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

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

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

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

     # normalize and transform to `BatchFeatures`
     lowercase__				 =							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
     lowercase__				 =							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\'"      )

     lowercase__				 =							WavaVecaForPreTraining(A      )

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

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





"""simple docstring"""


lowerCAmelCase_:					Union[str, Any]     =					[
    9_9_9,
    8_0_0,
    7_9_9,
    6_0_0,
    5_9_9,
    5_0_0,
    4_0_0,
    3_9_9,
    3_7_7,
    3_5_5,
    3_3_3,
    3_1_1,
    2_8_8,
    2_6_6,
    2_4_4,
    2_2_2,
    2_0_0,
    1_9_9,
    1_7_7,
    1_5_5,
    1_3_3,
    1_1_1,
    8_8,
    6_6,
    4_4,
    2_2,
    0,
]

lowerCAmelCase_:					List[str]     =					[
    9_9_9,
    9_7_6,
    9_5_2,
    9_2_8,
    9_0_5,
    8_8_2,
    8_5_8,
    8_5_7,
    8_1_0,
    7_6_2,
    7_1_5,
    7_1_4,
    5_7_2,
    4_2_9,
    4_2_8,
    2_8_6,
    2_8_5,
    2_3_8,
    1_9_0,
    1_4_3,
    1_4_2,
    1_1_8,
    9_5,
    7_1,
    4_7,
    2_4,
    0,
]

lowerCAmelCase_:					List[str]     =					[
    9_9_9,
    9_8_8,
    9_7_7,
    9_6_6,
    9_5_5,
    9_4_4,
    9_3_3,
    9_2_2,
    9_1_1,
    9_0_0,
    8_9_9,
    8_7_9,
    8_5_9,
    8_4_0,
    8_2_0,
    8_0_0,
    7_9_9,
    7_6_6,
    7_3_3,
    7_0_0,
    6_9_9,
    6_5_0,
    6_0_0,
    5_9_9,
    5_0_0,
    4_9_9,
    4_0_0,
    3_9_9,
    3_5_0,
    3_0_0,
    2_9_9,
    2_6_6,
    2_3_3,
    2_0_0,
    1_9_9,
    1_7_9,
    1_5_9,
    1_4_0,
    1_2_0,
    1_0_0,
    9_9,
    8_8,
    7_7,
    6_6,
    5_5,
    4_4,
    3_3,
    2_2,
    1_1,
    0,
]

lowerCAmelCase_:					Dict     =					[
    9_9_9,
    9_9_5,
    9_9_2,
    9_8_9,
    9_8_5,
    9_8_1,
    9_7_8,
    9_7_5,
    9_7_1,
    9_6_7,
    9_6_4,
    9_6_1,
    9_5_7,
    9_5_6,
    9_5_1,
    9_4_7,
    9_4_2,
    9_3_7,
    9_3_3,
    9_2_8,
    9_2_3,
    9_1_9,
    9_1_4,
    9_1_3,
    9_0_8,
    9_0_3,
    8_9_7,
    8_9_2,
    8_8_7,
    8_8_1,
    8_7_6,
    8_7_1,
    8_7_0,
    8_6_4,
    8_5_8,
    8_5_2,
    8_4_6,
    8_4_0,
    8_3_4,
    8_2_8,
    8_2_7,
    8_2_0,
    8_1_3,
    8_0_6,
    7_9_9,
    7_9_2,
    7_8_5,
    7_8_4,
    7_7_7,
    7_7_0,
    7_6_3,
    7_5_6,
    7_4_9,
    7_4_2,
    7_4_1,
    7_3_3,
    7_2_4,
    7_1_6,
    7_0_7,
    6_9_9,
    6_9_8,
    6_8_8,
    6_7_7,
    6_6_6,
    6_5_6,
    6_5_5,
    6_4_5,
    6_3_4,
    6_2_3,
    6_1_3,
    6_1_2,
    5_9_8,
    5_8_4,
    5_7_0,
    5_6_9,
    5_5_5,
    5_4_1,
    5_2_7,
    5_2_6,
    5_0_5,
    4_8_4,
    4_8_3,
    4_6_2,
    4_4_0,
    4_3_9,
    3_9_6,
    3_9_5,
    3_5_2,
    3_5_1,
    3_0_8,
    3_0_7,
    2_6_4,
    2_6_3,
    2_2_0,
    2_1_9,
    1_7_6,
    1_3_2,
    8_8,
    4_4,
    0,
]

lowerCAmelCase_:					Optional[int]     =					[
    9_9_9,
    9_9_7,
    9_9_5,
    9_9_2,
    9_9_0,
    9_8_8,
    9_8_6,
    9_8_4,
    9_8_1,
    9_7_9,
    9_7_7,
    9_7_5,
    9_7_2,
    9_7_0,
    9_6_8,
    9_6_6,
    9_6_4,
    9_6_1,
    9_5_9,
    9_5_7,
    9_5_6,
    9_5_4,
    9_5_1,
    9_4_9,
    9_4_6,
    9_4_4,
    9_4_1,
    9_3_9,
    9_3_6,
    9_3_4,
    9_3_1,
    9_2_9,
    9_2_6,
    9_2_4,
    9_2_1,
    9_1_9,
    9_1_6,
    9_1_4,
    9_1_3,
    9_1_0,
    9_0_7,
    9_0_5,
    9_0_2,
    8_9_9,
    8_9_6,
    8_9_3,
    8_9_1,
    8_8_8,
    8_8_5,
    8_8_2,
    8_7_9,
    8_7_7,
    8_7_4,
    8_7_1,
    8_7_0,
    8_6_7,
    8_6_4,
    8_6_1,
    8_5_8,
    8_5_5,
    8_5_2,
    8_4_9,
    8_4_6,
    8_4_3,
    8_4_0,
    8_3_7,
    8_3_4,
    8_3_1,
    8_2_8,
    8_2_7,
    8_2_4,
    8_2_1,
    8_1_7,
    8_1_4,
    8_1_1,
    8_0_8,
    8_0_4,
    8_0_1,
    7_9_8,
    7_9_5,
    7_9_1,
    7_8_8,
    7_8_5,
    7_8_4,
    7_8_0,
    7_7_7,
    7_7_4,
    7_7_0,
    7_6_6,
    7_6_3,
    7_6_0,
    7_5_6,
    7_5_2,
    7_4_9,
    7_4_6,
    7_4_2,
    7_4_1,
    7_3_7,
    7_3_3,
    7_3_0,
    7_2_6,
    7_2_2,
    7_1_8,
    7_1_4,
    7_1_0,
    7_0_7,
    7_0_3,
    6_9_9,
    6_9_8,
    6_9_4,
    6_9_0,
    6_8_5,
    6_8_1,
    6_7_7,
    6_7_3,
    6_6_9,
    6_6_4,
    6_6_0,
    6_5_6,
    6_5_5,
    6_5_0,
    6_4_6,
    6_4_1,
    6_3_6,
    6_3_2,
    6_2_7,
    6_2_2,
    6_1_8,
    6_1_3,
    6_1_2,
    6_0_7,
    6_0_2,
    5_9_6,
    5_9_1,
    5_8_6,
    5_8_0,
    5_7_5,
    5_7_0,
    5_6_9,
    5_6_3,
    5_5_7,
    5_5_1,
    5_4_5,
    5_3_9,
    5_3_3,
    5_2_7,
    5_2_6,
    5_1_9,
    5_1_2,
    5_0_5,
    4_9_8,
    4_9_1,
    4_8_4,
    4_8_3,
    4_7_4,
    4_6_6,
    4_5_7,
    4_4_9,
    4_4_0,
    4_3_9,
    4_2_8,
    4_1_8,
    4_0_7,
    3_9_6,
    3_9_5,
    3_8_1,
    3_6_6,
    3_5_2,
    3_5_1,
    3_3_0,
    3_0_8,
    3_0_7,
    2_8_6,
    2_6_4,
    2_6_3,
    2_4_2,
    2_2_0,
    2_1_9,
    1_7_6,
    1_7_5,
    1_3_2,
    1_3_1,
    8_8,
    4_4,
    0,
]

lowerCAmelCase_:					Tuple     =					[
    9_9_9,
    9_9_1,
    9_8_2,
    9_7_4,
    9_6_6,
    9_5_8,
    9_5_0,
    9_4_1,
    9_3_3,
    9_2_5,
    9_1_6,
    9_0_8,
    9_0_0,
    8_9_9,
    8_7_4,
    8_5_0,
    8_2_5,
    8_0_0,
    7_9_9,
    7_0_0,
    6_0_0,
    5_0_0,
    4_0_0,
    3_0_0,
    2_0_0,
    1_0_0,
    0,
]

lowerCAmelCase_:					str     =					[
    9_9_9,
    9_9_2,
    9_8_5,
    9_7_8,
    9_7_1,
    9_6_4,
    9_5_7,
    9_4_9,
    9_4_2,
    9_3_5,
    9_2_8,
    9_2_1,
    9_1_4,
    9_0_7,
    9_0_0,
    8_9_9,
    8_7_9,
    8_5_9,
    8_4_0,
    8_2_0,
    8_0_0,
    7_9_9,
    7_6_6,
    7_3_3,
    7_0_0,
    6_9_9,
    6_5_0,
    6_0_0,
    5_9_9,
    5_0_0,
    4_9_9,
    4_0_0,
    3_9_9,
    3_0_0,
    2_9_9,
    2_0_0,
    1_9_9,
    1_0_0,
    9_9,
    0,
]

lowerCAmelCase_:					int     =					[
    9_9_9,
    9_9_6,
    9_9_2,
    9_8_9,
    9_8_5,
    9_8_2,
    9_7_9,
    9_7_5,
    9_7_2,
    9_6_8,
    9_6_5,
    9_6_1,
    9_5_8,
    9_5_5,
    9_5_1,
    9_4_8,
    9_4_4,
    9_4_1,
    9_3_8,
    9_3_4,
    9_3_1,
    9_2_7,
    9_2_4,
    9_2_0,
    9_1_7,
    9_1_4,
    9_1_0,
    9_0_7,
    9_0_3,
    9_0_0,
    8_9_9,
    8_9_1,
    8_8_4,
    8_7_6,
    8_6_9,
    8_6_1,
    8_5_3,
    8_4_6,
    8_3_8,
    8_3_0,
    8_2_3,
    8_1_5,
    8_0_8,
    8_0_0,
    7_9_9,
    7_8_8,
    7_7_7,
    7_6_6,
    7_5_5,
    7_4_4,
    7_3_3,
    7_2_2,
    7_1_1,
    7_0_0,
    6_9_9,
    6_8_8,
    6_7_7,
    6_6_6,
    6_5_5,
    6_4_4,
    6_3_3,
    6_2_2,
    6_1_1,
    6_0_0,
    5_9_9,
    5_8_5,
    5_7_1,
    5_5_7,
    5_4_2,
    5_2_8,
    5_1_4,
    5_0_0,
    4_9_9,
    4_8_5,
    4_7_1,
    4_5_7,
    4_4_2,
    4_2_8,
    4_1_4,
    4_0_0,
    3_9_9,
    3_7_9,
    3_5_9,
    3_4_0,
    3_2_0,
    3_0_0,
    2_9_9,
    2_7_9,
    2_5_9,
    2_4_0,
    2_2_0,
    2_0_0,
    1_9_9,
    1_6_6,
    1_3_3,
    1_0_0,
    9_9,
    6_6,
    3_3,
    0,
]




