



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


lowerCamelCase	 =       logging.getLogger(__name__)





@dataclass
class _UpperCamelCase		:





				'''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.99_99_95      ,       metadata={"""help""": """Decay of gumbel temperature during training."""}		)



def   _A       (   _lowerCAmelCase  ,   _lowerCAmelCase					):



							"""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(_lowerCAmelCase					)





@dataclass
class _UpperCamelCase		:





				'''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 _UpperCamelCase		:





				'''simple docstring'''





				lowerCAmelCase__					    =    42
				lowerCAmelCase__					    =    42
				lowerCAmelCase__					    =    "longest"
				lowerCAmelCase__					    =    None
				lowerCAmelCase__					    =    None

				def __call__(							self  :  int  ,       _lowerCAmelCase  :  List[Dict[str, Union[List[int], torch.Tensor]]]):







											'''simple docstring'''

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

											return batch





class _UpperCamelCase		(     A		):





				'''simple docstring'''





				def __init__(							self  :  str  ,       *_lowerCAmelCase  :  Any  ,       _lowerCAmelCase  :  Optional[Any]=1  ,       _lowerCAmelCase  :  List[Any]=0  ,       _lowerCAmelCase  :  Union[str, Any]=1.0  ,       **_lowerCAmelCase  :  Optional[int]):







											'''simple docstring'''

											super().__init__(*_lowerCAmelCase  ,       **_lowerCAmelCase)
											__lowercase       	=0
											__lowercase       	=max_gumbel_temp
											__lowercase       	=min_gumbel_temp
											__lowercase       	=gumbel_temp_decay

				def 				__lowerCamelCase     (							self  :  List[Any]  ,       _lowerCAmelCase  :  nn.Module  ,       _lowerCAmelCase  :  Dict[str, Union[torch.Tensor, Any]]):







											'''simple docstring'''

											model.train()
											__lowercase       	=self._prepare_inputs(_lowerCAmelCase)

											if self.use_amp:
																		with autocast():
																									__lowercase       	=self.compute_loss(_lowerCAmelCase  ,       _lowerCAmelCase)
											else:
																		__lowercase       	=self.compute_loss(_lowerCAmelCase  ,       _lowerCAmelCase)

											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(_lowerCAmelCase).backward()
											elif self.use_apex:
																		with amp.scale_loss(_lowerCAmelCase  ,       self.optimizer) as scaled_loss:
																									scaled_loss.backward()
											elif self.deepspeed:
																		self.deepspeed.backward(_lowerCAmelCase)
											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	,				__lowercase	,				__lowercase       	=parser.parse_args_into_dataclasses()
							configure_logger(_lowerCAmelCase  ,   _lowerCAmelCase					)

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

							def prepare_dataset(_lowerCAmelCase					):
														# check that all files have the correct sampling rate
														__lowercase	,				__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(
							    _lowerCAmelCase  ,   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 _lowerCAmelCase					: len(data['speech']					) < int(data_args.max_duration_in_seconds * feature_extractor.sampling_rate					)					)

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

							# normalize and transform to `BatchFeatures`
							__lowercase       	=vectorized_datasets.map(
							    _lowerCAmelCase  ,   batched=_lowerCAmelCase  ,   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(_lowerCAmelCase					)

							__lowercase       	=DataCollatorForWavaVecaPretraining(model=_lowerCAmelCase  ,   feature_extractor=_lowerCAmelCase					)

							__lowercase       	=WavaVecaPreTrainer(
							    model=_lowerCAmelCase  ,   data_collator=_lowerCAmelCase  ,   args=_lowerCAmelCase  ,   train_dataset=vectorized_datasets['train']  ,   eval_dataset=vectorized_datasets['validation']  ,   tokenizer=_lowerCAmelCase  ,   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 warnings

from ...utils import logging
from .image_processing_deit import DeiTImageProcessor


lowerCamelCase	 =       logging.get_logger(__name__)





class _UpperCamelCase		(     A		):





				'''simple docstring'''





				def __init__(							self  :  Optional[int]  ,       *_lowerCAmelCase  :  Union[str, Any]  ,       **_lowerCAmelCase  :  Union[str, Any]):







											'''simple docstring'''

											warnings.warn(
											    'The class DeiTFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please'
											    ' use DeiTImageProcessor instead.'  ,       _lowerCAmelCase  ,       )
											super().__init__(*_lowerCAmelCase  ,       **_lowerCAmelCase)



