


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


__snake_case											=					logging.getLogger(__name__)
@dataclass
class        lowercase     :

		"""simple docstring"""

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



def       a				(      __a			, __a       )  ->     str:



					'''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       )]			, )
					UpperCamelCase__       :int							 =			logging.WARNING
					if model_args.verbose_logging:
										UpperCamelCase__       :Dict							 =			logging.DEBUG
					elif trainer_utils.is_main_process(training_args.local_rank       ):
										UpperCamelCase__       :Optional[int]							 =			logging.INFO
					logger.setLevel(__a       )
@dataclass
class        lowercase     :

		"""simple docstring"""

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



@dataclass
class        lowercase     :

		"""simple docstring"""

		_a   			=    42
		_a   			=    42
		_a   			=    "longest"
		_a   			=    None
		_a   			=    None




		def __call__(		self   ,       UpperCamelCase_							):

							'''simple docstring'''


							UpperCamelCase__       :Any							 =			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__       :List[Any]							 =			self.model._get_feat_extract_output_lengths(batch['''input_values'''].shape[-1]							)

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

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

							# sample randomly masked indices
							UpperCamelCase__       :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        lowercase     (							A__ ):

		"""simple docstring"""

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

							'''simple docstring'''


							super().__init__(*UpperCamelCase_   ,       **UpperCamelCase_							)
							UpperCamelCase__       :Dict							 =			0
							UpperCamelCase__       :Tuple							 =			max_gumbel_temp
							UpperCamelCase__       :Optional[Any]							 =			min_gumbel_temp
							UpperCamelCase__       :Optional[Any]							 =			gumbel_temp_decay




		def    lowerCAmelCase__    (		self   ,       UpperCamelCase_   ,       UpperCamelCase_							):

							'''simple docstring'''


							model.train()
							UpperCamelCase__       :List[str]							 =			self._prepare_inputs(UpperCamelCase_							)

							if self.use_amp:
												with autocast():
																	UpperCamelCase__       :Any							 =			self.compute_loss(UpperCamelCase_   ,       UpperCamelCase_							)
							else:
												UpperCamelCase__       :Union[str, Any]							 =			self.compute_loss(UpperCamelCase_   ,       UpperCamelCase_							)

							if self.args.n_gpu > 1 or self.deepspeed:
												if model.module.config.ctc_loss_reduction == "mean":
																	UpperCamelCase__       :Any							 =			loss.mean()
												elif model.module.config.ctc_loss_reduction == "sum":
																	UpperCamelCase__       :List[str]							 =			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__       :Optional[int]							 =			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				(      )  ->     List[str]:



					'''simple docstring'''






					UpperCamelCase__       :int							 =			HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)       )

					UpperCamelCase__ ,    UpperCamelCase__ ,    UpperCamelCase__       :List[Any]							 =			parser.parse_args_into_dataclasses()
					configure_logger(__a			, __a       )

					# Downloading and loading a dataset from the hub.
					UpperCamelCase__       :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"
										UpperCamelCase__       :List[str]							 =			DatasetDict()
										UpperCamelCase__       :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			, )
										UpperCamelCase__       :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			, )
					else:
										# make sure only "validation" and "train" keys remain"
										UpperCamelCase__       :Tuple							 =			DatasetDict()
										UpperCamelCase__       :List[Any]							 =			load_dataset(
										    data_args.dataset_name			, data_args.dataset_config_name			, split='''validation'''			, cache_dir=model_args.cache_dir			, )
										UpperCamelCase__       :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
					UpperCamelCase__       :str							 =			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
										UpperCamelCase__ ,    UpperCamelCase__       :Union[str, Any]							 =			librosa.load(batch[data_args.speech_file_column]			, sr=feature_extractor.sampling_rate       )
										return batch

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

					# filter audio files that are too long
					UpperCamelCase__       :Tuple							 =			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`
					UpperCamelCase__       :Optional[int]							 =			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__       :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\''''       )

					UpperCamelCase__       :str							 =			WavaVecaForPreTraining(__a       )

					UpperCamelCase__       :List[str]							 =			DataCollatorForWavaVecaPretraining(model=__a			, feature_extractor=__a       )

					UpperCamelCase__       :Any							 =			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'''
# tests directory-specific settings - this file is run automatically
# by pytest before any tests are run

import sys
import warnings
from os.path import abspath, dirname, join


# allow having multiple repository checkouts and not needing to remember to rerun
# 'pip install -e .[dev]' when switching between checkouts and running tests.
__snake_case											=					abspath(join(dirname(dirname(__file__)), '''src'''))
sys.path.insert(1, git_repo_path)

# silence FutureWarning warnings in tests since often we can't act on them until
# they become normal warnings - i.e. the tests still need to test the current functionality
warnings.simplefilter(action='''ignore''', category=FutureWarning)


def       a				(      __a       )  ->     Optional[int]:



					'''simple docstring'''






					from diffusers.utils.testing_utils import pytest_addoption_shared

					pytest_addoption_shared(__a       )







def       a				(      __a       )  ->     str:



					'''simple docstring'''






					from diffusers.utils.testing_utils import pytest_terminal_summary_main

					UpperCamelCase__       :Union[str, Any]							 =			terminalreporter.config.getoption('''--make-reports'''       )
					if make_reports:
										pytest_terminal_summary_main(__a			, id=__a       )