


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"""






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



def        UpperCamelCase_(				_A   :ModelArguments				,     _A   :TrainingArguments  )->		Tuple:
			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(_A  )







@dataclass
class        lowerCamelCase__	:






			"""simple docstring"""






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







@dataclass
class        lowerCamelCase__	:






			"""simple docstring"""






			_UpperCamelCase		:					WavaVecaForPreTraining
			_UpperCamelCase		:					WavaVecaFeatureExtractor
			_UpperCamelCase		:					Union[bool, str]    =		"longest"
			_UpperCamelCase		:					Optional[int]    =		None
			_UpperCamelCase		:					Optional[int]    =		None




			def __call__(							self      ,						snake_case     ):

						'''simple docstring'''


						UpperCamelCase__					=       self.feature_extractor.pad(
						    snake_case      ,						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=snake_case      ,						min_masks=2      ,						)

						return batch







class        lowerCamelCase__	(							UpperCAmelCase      ):






			"""simple docstring"""






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

						'''simple docstring'''


						super().__init__(*snake_case      ,						**snake_case     )
						UpperCamelCase__					=       0
						UpperCamelCase__					=       max_gumbel_temp
						UpperCamelCase__					=       min_gumbel_temp
						UpperCamelCase__					=       gumbel_temp_decay







			def 				snake_case__				(							self      ,						snake_case      ,						snake_case     ):

						'''simple docstring'''


						model.train()
						UpperCamelCase__					=       self._prepare_inputs(snake_case     )

						if self.use_amp:
									with autocast():
												UpperCamelCase__					=       self.compute_loss(snake_case      ,						snake_case     )
						else:
									UpperCamelCase__					=       self.compute_loss(snake_case      ,						snake_case     )

						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(snake_case     ).backward()
						elif self.use_apex:
									with amp.scale_loss(snake_case      ,						self.optimizer     ) as scaled_loss:
												scaled_loss.backward()
						elif self.deepspeed:
									self.deepspeed.backward(snake_case     )
						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_(				)->		Any:
			# 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(_A				,     _A  )

			# 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=_A  )

			def prepare_dataset(_A   :Optional[Any]  ):
						# 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(
			    _A				,     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 _A  : len(data["speech"]  ) < int(data_args.max_duration_in_seconds * feature_extractor.sampling_rate  )  )

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

			# normalize and transform to `BatchFeatures`
			UpperCamelCase__					=       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
			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(_A  )

			UpperCamelCase__					=       DataCollatorForWavaVecaPretraining(model=_A				,     feature_extractor=_A  )

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






import itertools
import os
import random
import tempfile
import unittest

import numpy as np
from datasets import load_dataset

from transformers import is_speech_available
from transformers.testing_utils import check_json_file_has_correct_format, require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
						from transformers import WhisperFeatureExtractor

if is_torch_available():
						import torch

__UpperCamelCase    		=       random.Random()
def        UpperCamelCase_(				_A   :Tuple				,     _A   :str=1.0				,     _A   :int=None				,     _A   :Dict=None  )->		str:
			if rng is None:
						UpperCamelCase__					=       global_rng

			UpperCamelCase__					=       []
			for batch_idx in range(shape[0]  ):
						values.append([]  )
						for _ in range(shape[1]  ):
									values[-1].append(rng.random() * scale  )

			return values







@require_torch
@require_torchaudio
class        lowerCamelCase__	(							unittest.TestCase      ):






			"""simple docstring"""






			def __init__(							self      ,						snake_case      ,						snake_case=7      ,						snake_case=400      ,						snake_case=2000      ,						snake_case=10      ,						snake_case=160      ,						snake_case=8      ,						snake_case=0.0      ,						snake_case=4000      ,						snake_case=False      ,						snake_case=True      ,						):

						'''simple docstring'''


						UpperCamelCase__					=       parent
						UpperCamelCase__					=       batch_size
						UpperCamelCase__					=       min_seq_length
						UpperCamelCase__					=       max_seq_length
						UpperCamelCase__					=       (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
						UpperCamelCase__					=       padding_value
						UpperCamelCase__					=       sampling_rate
						UpperCamelCase__					=       return_attention_mask
						UpperCamelCase__					=       do_normalize
						UpperCamelCase__					=       feature_size
						UpperCamelCase__					=       chunk_length
						UpperCamelCase__					=       hop_length




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						return {
						    "feature_size": self.feature_size,
						    "hop_length": self.hop_length,
						    "chunk_length": self.chunk_length,
						    "padding_value": self.padding_value,
						    "sampling_rate": self.sampling_rate,
						    "return_attention_mask": self.return_attention_mask,
						    "do_normalize": self.do_normalize,
						}







			def 				snake_case__				(							self      ,						snake_case=False      ,						snake_case=False     ):

						'''simple docstring'''


						def _flatten(snake_case     ):
									return list(itertools.chain(*snake_case     )     )

						if equal_length:
									UpperCamelCase__					=       [floats_list((self.max_seq_length, self.feature_size)     ) for _ in range(self.batch_size     )]
						else:
									# make sure that inputs increase in size
									UpperCamelCase__					=       [
									    floats_list((x, self.feature_size)     )
									    for x in range(self.min_seq_length      ,						self.max_seq_length      ,						self.seq_length_diff     )
									]
						if numpify:
									UpperCamelCase__					=       [np.asarray(snake_case     ) for x in speech_inputs]
						return speech_inputs







@require_torch
@require_torchaudio
class        lowerCamelCase__	(							UpperCAmelCase					,	unittest.TestCase      ):






			"""simple docstring"""






			_UpperCamelCase		:					Any    =		WhisperFeatureExtractor if is_speech_available() else None




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       WhisperFeatureExtractionTester(self     )




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       self.feature_extraction_class(**self.feat_extract_dict     )

						with tempfile.TemporaryDirectory() as tmpdirname:
									UpperCamelCase__					=       feat_extract_first.save_pretrained(snake_case     )[0]
									check_json_file_has_correct_format(snake_case     )
									UpperCamelCase__					=       self.feature_extraction_class.from_pretrained(snake_case     )

						UpperCamelCase__					=       feat_extract_first.to_dict()
						UpperCamelCase__					=       feat_extract_second.to_dict()
						UpperCamelCase__					=       feat_extract_first.mel_filters
						UpperCamelCase__					=       feat_extract_second.mel_filters
						self.assertTrue(np.allclose(snake_case      ,						snake_case     )     )
						self.assertEqual(snake_case      ,						snake_case     )




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       self.feature_extraction_class(**self.feat_extract_dict     )

						with tempfile.TemporaryDirectory() as tmpdirname:
									UpperCamelCase__					=       os.path.join(snake_case      ,						"feat_extract.json"     )
									feat_extract_first.to_json_file(snake_case     )
									UpperCamelCase__					=       self.feature_extraction_class.from_json_file(snake_case     )

						UpperCamelCase__					=       feat_extract_first.to_dict()
						UpperCamelCase__					=       feat_extract_second.to_dict()
						UpperCamelCase__					=       feat_extract_first.mel_filters
						UpperCamelCase__					=       feat_extract_second.mel_filters
						self.assertTrue(np.allclose(snake_case      ,						snake_case     )     )
						self.assertEqual(snake_case      ,						snake_case     )




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()     )
						# create three inputs of length 800, 1000, and 1200
						UpperCamelCase__					=       [floats_list((1, x)     )[0] for x in range(800      ,						1400      ,						200     )]
						UpperCamelCase__					=       [np.asarray(snake_case     ) for speech_input in speech_inputs]

						# Test feature size
						UpperCamelCase__					=       feature_extractor(snake_case      ,						padding="max_length"      ,						return_tensors="np"     ).input_features
						self.assertTrue(input_features.ndim == 3     )
						self.assertTrue(input_features.shape[-1] == feature_extractor.nb_max_frames     )
						self.assertTrue(input_features.shape[-2] == feature_extractor.feature_size     )

						# Test not batched input
						UpperCamelCase__					=       feature_extractor(speech_inputs[0]      ,						return_tensors="np"     ).input_features
						UpperCamelCase__					=       feature_extractor(np_speech_inputs[0]      ,						return_tensors="np"     ).input_features
						self.assertTrue(np.allclose(snake_case      ,						snake_case      ,						atol=1E-3     )     )

						# Test batched
						UpperCamelCase__					=       feature_extractor(snake_case      ,						return_tensors="np"     ).input_features
						UpperCamelCase__					=       feature_extractor(snake_case      ,						return_tensors="np"     ).input_features
						for enc_seq_a, enc_seq_a in zip(snake_case      ,						snake_case     ):
									self.assertTrue(np.allclose(snake_case      ,						snake_case      ,						atol=1E-3     )     )

						# Test 2-D numpy arrays are batched.
						UpperCamelCase__					=       [floats_list((1, x)     )[0] for x in (800, 800, 800)]
						UpperCamelCase__					=       np.asarray(snake_case     )
						UpperCamelCase__					=       feature_extractor(snake_case      ,						return_tensors="np"     ).input_features
						UpperCamelCase__					=       feature_extractor(snake_case      ,						return_tensors="np"     ).input_features
						for enc_seq_a, enc_seq_a in zip(snake_case      ,						snake_case     ):
									self.assertTrue(np.allclose(snake_case      ,						snake_case      ,						atol=1E-3     )     )

						# Test truncation required
						UpperCamelCase__					=       [floats_list((1, x)     )[0] for x in range(200      ,						(feature_extractor.n_samples + 500)      ,						200     )]
						UpperCamelCase__					=       [np.asarray(snake_case     ) for speech_input in speech_inputs]

						UpperCamelCase__					=       [x[: feature_extractor.n_samples] for x in speech_inputs]
						UpperCamelCase__					=       [np.asarray(snake_case     ) for speech_input in speech_inputs_truncated]

						UpperCamelCase__					=       feature_extractor(snake_case      ,						return_tensors="np"     ).input_features
						UpperCamelCase__					=       feature_extractor(snake_case      ,						return_tensors="np"     ).input_features
						for enc_seq_a, enc_seq_a in zip(snake_case      ,						snake_case     ):
									self.assertTrue(np.allclose(snake_case      ,						snake_case      ,						atol=1E-3     )     )




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						import torch

						UpperCamelCase__					=       self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()     )
						UpperCamelCase__					=       np.random.rand(100      ,						32     ).astype(np.floataa     )
						UpperCamelCase__					=       np_speech_inputs.tolist()

						for inputs in [py_speech_inputs, np_speech_inputs]:
									UpperCamelCase__					=       feature_extractor.pad([{"input_features": inputs}]      ,						return_tensors="np"     )
									self.assertTrue(np_processed.input_features.dtype == np.floataa     )
									UpperCamelCase__					=       feature_extractor.pad([{"input_features": inputs}]      ,						return_tensors="pt"     )
									self.assertTrue(pt_processed.input_features.dtype == torch.floataa     )




			def 				snake_case__				(							self      ,						snake_case     ):

						'''simple docstring'''


						UpperCamelCase__					=       load_dataset("hf-internal-testing/librispeech_asr_dummy"      ,						"clean"      ,						split="validation"     )
						# automatic decoding with librispeech
						UpperCamelCase__					=       ds.sort("id"     ).select(range(snake_case     )     )[:num_samples]["audio"]

						return [x["array"] for x in speech_samples]




			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       torch.tensor(
						    [
						        0.1193, -0.0946, -0.1098, -0.0196, 0.0225, -0.0690, -0.1736, 0.0951,
						        0.0971, -0.0817, -0.0702, 0.0162, 0.0260, 0.0017, -0.0192, -0.1678,
						        0.0709, -0.1867, -0.0655, -0.0274, -0.0234, -0.1884, -0.0516, -0.0554,
						        -0.0274, -0.1425, -0.1423, 0.0837, 0.0377, -0.0854
						    ]     )
						# fmt: on

						UpperCamelCase__					=       self._load_datasamples(1     )
						UpperCamelCase__					=       WhisperFeatureExtractor()
						UpperCamelCase__					=       feature_extractor(snake_case      ,						return_tensors="pt"     ).input_features
						self.assertEqual(input_features.shape      ,						(1, 80, 3000)     )
						self.assertTrue(torch.allclose(input_features[0, 0, :30]      ,						snake_case      ,						atol=1E-4     )     )







			def 				snake_case__				(							self     ):

						'''simple docstring'''


						UpperCamelCase__					=       self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()     )
						UpperCamelCase__					=       self._load_datasamples(1     )[0]
						UpperCamelCase__					=       ((audio - audio.min()) / (audio.max() - audio.min())) * 65535  # Rescale to [0, 65535] to show issue
						UpperCamelCase__					=       feat_extract.zero_mean_unit_var_norm([audio]      ,						attention_mask=snake_case     )[0]

						self.assertTrue(np.all(np.mean(snake_case     ) < 1E-3     )     )
						self.assertTrue(np.all(np.abs(np.var(snake_case     ) - 1     ) < 1E-3     )     )


