



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


SCREAMING_SNAKE_CASE_				= 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=a_     ,metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"}     ,)
							__lowerCAmelCase    					=					field(
							    default=a_     ,metadata={"help": "Whether to freeze the feature extractor layers of the model."}				)
							__lowerCAmelCase    					=					field(
							    default=a_     ,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.999_995     ,metadata={"help": "Decay of gumbel temperature during training."}				)



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




@dataclass
class 			snake_case_						:

							'''simple docstring'''





							__lowerCAmelCase    					=					field(
							    default=a_     ,metadata={"help": "The name of the dataset to use (via the datasets library)."}				)
							__lowerCAmelCase    					=					field(
							    default=a_     ,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=a_     ,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=a_     ,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			,     a_      ):
									a_     :      str 						=     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"			,     )
									a_     :      Optional[int] 						=     self.model._get_feat_extract_output_lengths(batch["input_values"].shape[-1]      )

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

											a_     :      Dict 						=     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_     :      Optional[Any] 						=     1
											a_     :      List[Any] 						=     attention_mask.flip([-1]      ).cumsum(-1      ).flip([-1]      ).bool()

									# sample randomly masked indices
									a_     :      Union[str, Any] 						=     _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 			snake_case_						(						a_				):

							'''simple docstring'''





							def __init__(   self			,     *a_			,     a_=1			,     a_=0			,     a_=1.0			,     **a_      ):
									super().__init__(*UpperCAmelCase__			,     **UpperCAmelCase__      )
									a_     :      Any 						=     0
									a_     :      List[str] 						=     max_gumbel_temp
									a_     :      Optional[Any] 						=     min_gumbel_temp
									a_     :      Tuple 						=     gumbel_temp_decay






							def snake_case_  (   self			,     a_			,     a_      ):
									model.train()
									a_     :      Optional[Any] 						=     self._prepare_inputs(UpperCAmelCase__      )

									if self.use_amp:
											with autocast():
													a_     :      Dict 						=     self.compute_loss(UpperCAmelCase__			,     UpperCAmelCase__      )
									else:
											a_     :      Tuple 						=     self.compute_loss(UpperCAmelCase__			,     UpperCAmelCase__      )

									if self.args.n_gpu > 1 or self.deepspeed:
											if model.module.config.ctc_loss_reduction == "mean":
													a_     :      Tuple 						=     loss.mean()
											elif model.module.config.ctc_loss_reduction == "sum":
													a_     :      Optional[Any] 						=     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_     :      Optional[Any] 						=     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        lowerCAmelCase_						(							)			->      Optional[Any]:
		a_     :      Optional[int] 						=     HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)	)

		a_     :      str 						=     parser.parse_args_into_dataclasses()
		configure_logger(SCREAMING_SNAKE_CASE__, SCREAMING_SNAKE_CASE__	)

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

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

		# load audio files into numpy arrays
		a_     :      Dict 						=     datasets.map(
		    SCREAMING_SNAKE_CASE__, num_proc=data_args.preprocessing_num_workers, remove_columns=datasets["train"].column_names	)

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

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

		# normalize and transform to `BatchFeatures`
		a_     :      Dict 						=     vectorized_datasets.map(
		    SCREAMING_SNAKE_CASE__, batched=SCREAMING_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
		a_     :      List[str] 						=     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_     :      Optional[int] 						=     WavaVecaForPreTraining(SCREAMING_SNAKE_CASE__	)

		a_     :      Optional[int] 						=     DataCollatorForWavaVecaPretraining(model=SCREAMING_SNAKE_CASE__, feature_extractor=SCREAMING_SNAKE_CASE__	)

		a_     :      int 						=     WavaVecaPreTrainer(
		    model=SCREAMING_SNAKE_CASE__, data_collator=SCREAMING_SNAKE_CASE__, args=SCREAMING_SNAKE_CASE__, train_dataset=vectorized_datasets["train"], eval_dataset=vectorized_datasets["validation"], tokenizer=SCREAMING_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"""
from typing import TYPE_CHECKING

from ..utils import _LazyModule


SCREAMING_SNAKE_CASE_				= {
    """config""": [
        """EXTERNAL_DATA_FORMAT_SIZE_LIMIT""",
        """OnnxConfig""",
        """OnnxConfigWithPast""",
        """OnnxSeq2SeqConfigWithPast""",
        """PatchingSpec""",
    ],
    """convert""": ["""export""", """validate_model_outputs"""],
    """features""": ["""FeaturesManager"""],
    """utils""": ["""ParameterFormat""", """compute_serialized_parameters_size"""],
}


if TYPE_CHECKING:
					from .config import (
					    EXTERNAL_DATA_FORMAT_SIZE_LIMIT,
					    OnnxConfig,
					    OnnxConfigWithPast,
					    OnnxSeqaSeqConfigWithPast,
					    PatchingSpec,
					)
					from .convert import export, validate_model_outputs
					from .features import FeaturesManager
					from .utils import ParameterFormat, compute_serialized_parameters_size

else:
					import sys

					SCREAMING_SNAKE_CASE_				= _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)