


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


_a  :  int=				logging.getLogger(__name__)

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







def 						__UpperCAmelCase       (					UpperCAmelCase_					:       ModelArguments    ,							UpperCAmelCase_					:       TrainingArguments						)							->    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						)]    ,							)
	__snake_case    :							Union[str, Any]   	=				logging.WARNING
	if model_args.verbose_logging:
		__snake_case    :							Optional[int]   	=				logging.DEBUG
	elif trainer_utils.is_main_process(training_args.local_rank						):
		__snake_case    :							List[Any]   	=				logging.INFO
	logger.setLevel(UpperCAmelCase_						)

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







@dataclass
class 	UpperCamelCase    :
				UpperCAmelCase       :   WavaVecaForPreTraining
				UpperCAmelCase       :   WavaVecaFeatureExtractor
				UpperCAmelCase       :   Union[bool, str]					=	"longest"
				UpperCAmelCase       :   Optional[int]					=	None
				UpperCAmelCase       :   Optional[int]					=	None



				def __call__(self       :   int			,						_A       :   List[Dict[str, Union[List[int], torch.Tensor]]])					->			Dict[str, torch.Tensor]:
					# reformat list to dict and set to pytorch format
					__snake_case    :							Dict   	=				self.feature_extractor.pad(
					    _A			,						max_length=self.max_length			,						padding=self.padding			,						pad_to_multiple_of=self.pad_to_multiple_of			,						return_tensors='pt'			,						)
					__snake_case    :							Tuple   	=				self.model._get_feat_extract_output_lengths(batch['input_values'].shape[-1])

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

						__snake_case    :							Optional[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
						__snake_case    :							Union[str, Any]   	=				1
						__snake_case    :							Tuple   	=				attention_mask.flip([-1]).cumsum(-1).flip([-1]).bool()

					# sample randomly masked indices
					__snake_case    :							Tuple   	=				_compute_mask_indices(
					    (batch_size, mask_indices_seq_length)			,						self.model.config.mask_time_prob			,						self.model.config.mask_time_length			,						attention_mask=_A			,						min_masks=2			,						)

					return batch







class 	UpperCamelCase    (					lowercase		):



				def __init__(self       :   Union[str, Any]			,						*_A       :   Dict			,						_A       :   Any=1			,						_A       :   int=0			,						_A       :   Optional[int]=1.0			,						**_A       :   Union[str, Any])					->			int:
					super().__init__(*_A			,						**_A)
					__snake_case    :							Tuple   	=				0
					__snake_case    :							Dict   	=				max_gumbel_temp
					__snake_case    :							List[Any]   	=				min_gumbel_temp
					__snake_case    :							List[Any]   	=				gumbel_temp_decay






				def       _lowercase  (self       :   Optional[int]			,						_A       :   nn.Module			,						_A       :   Dict[str, Union[torch.Tensor, Any]])					->			torch.Tensor:
					model.train()
					__snake_case    :							Dict   	=				self._prepare_inputs(_A)

					if self.use_amp:
						with autocast():
							__snake_case    :							Union[str, Any]   	=				self.compute_loss(_A			,						_A)
					else:
						__snake_case    :							Dict   	=				self.compute_loss(_A			,						_A)

					if self.args.n_gpu > 1 or self.deepspeed:
						if model.module.config.ctc_loss_reduction == "mean":
							__snake_case    :							List[str]   	=				loss.mean()
						elif model.module.config.ctc_loss_reduction == "sum":
							__snake_case    :							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:
						__snake_case    :							Union[str, Any]   	=				loss / self.args.gradient_accumulation_steps

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

	'''simple docstring'''





	__snake_case    :							List[str]   	=				HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)						)

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

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

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

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

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

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

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

	__snake_case    :							Optional[Any]   	=				WavaVecaForPreTraining(UpperCAmelCase_						)

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

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

import tensorflow as tf
from packaging import version
def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Union[str, Any]						)							->    Any:

	'''simple docstring'''





	__snake_case    :							List[str]   	=				tf.convert_to_tensor(UpperCAmelCase_						)
	__snake_case    :							List[str]   	=				0.5 * (1.0 + tf.math.erf(x / tf.cast(tf.sqrt(2.0						)    ,							x.dtype						)						))

	return x * cdf
def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Tuple						)							->    Optional[Any]:

	'''simple docstring'''





	__snake_case    :							Optional[int]   	=				tf.convert_to_tensor(UpperCAmelCase_						)
	__snake_case    :							Any   	=				tf.cast(math.pi    ,							x.dtype						)
	__snake_case    :							List[Any]   	=				tf.cast(0.044_715    ,							x.dtype						)
	__snake_case    :							Union[str, Any]   	=				0.5 * (1.0 + tf.tanh(tf.sqrt(2.0 / pi						) * (x + coeff * tf.pow(UpperCAmelCase_    ,							3						))						))

	return x * cdf
def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Tuple						)							->    List[Any]:

	'''simple docstring'''





	__snake_case    :							Optional[int]   	=				tf.convert_to_tensor(UpperCAmelCase_						)

	return x * tf.tanh(tf.math.softplus(UpperCAmelCase_						)						)
def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Tuple						)							->    List[str]:

	'''simple docstring'''





	__snake_case    :							Union[str, Any]   	=				tf.convert_to_tensor(UpperCAmelCase_						)
	__snake_case    :							Dict   	=				tf.cast(0.044_715    ,							x.dtype						)
	__snake_case    :							int   	=				tf.cast(0.7_978_845_608    ,							x.dtype						)

	return 0.5 * x * (1.0 + tf.tanh(x * coeffa * (1.0 + coeffa * x * x)						))
def 						__UpperCAmelCase       (					UpperCAmelCase_					:       int						)							->    Dict:

	'''simple docstring'''





	__snake_case    :							Any   	=				tf.convert_to_tensor(UpperCAmelCase_						)
	__snake_case    :							Dict   	=				tf.cast(1.702    ,							x.dtype						)
	return x * tf.math.sigmoid(coeff * x						)
def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Union[str, Any]						)							->    List[Any]:

	'''simple docstring'''





	return tf.clip_by_value(_gelu(UpperCAmelCase_						)    ,							-10    ,							10						)
def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Union[str, Any]    ,							UpperCAmelCase_					:       Any=-1						)							->    int:

	'''simple docstring'''





	__snake_case	,       __snake_case    :							str   	=				tf.split(UpperCAmelCase_    ,							2    ,							axis=UpperCAmelCase_						)
	return a * tf.math.sigmoid(UpperCAmelCase_						)


if version.parse(tf.version.VERSION) >= version.parse("2.4"):
	def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Tuple						)							->    int:

		'''simple docstring'''





		return tf.keras.activations.gelu(UpperCAmelCase_    ,							approximate=UpperCAmelCase_						)

	_a  :  Tuple=				tf.keras.activations.gelu

	_a  :  Tuple=				approximate_gelu_wrap
else:
	_a  :  Any=				_gelu
	_a  :  str=				_gelu_new


_a  :  List[str]=				{
    "gelu": gelu,
    "gelu_10": gelu_aa,
    "gelu_fast": gelu_fast,
    "gelu_new": gelu_new,
    "glu": glu,
    "mish": mish,
    "quick_gelu": quick_gelu,
    "relu": tf.keras.activations.relu,
    "sigmoid": tf.keras.activations.sigmoid,
    "silu": tf.keras.activations.swish,
    "swish": tf.keras.activations.swish,
    "tanh": tf.keras.activations.tanh,
}

def 						__UpperCAmelCase       (					UpperCAmelCase_					:       Optional[Any]						)							->    List[Any]:

	'''simple docstring'''





	if activation_string in ACTaFN:
		return ACTaFN[activation_string]
	else:
		raise KeyError(F"function {activation_string} not found in ACT2FN mapping {list(ACTaFN.keys()						)}"						)





