




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


lowercase__			=		logging.getLogger(__name__)






@dataclass
class lowerCAmelCase__      :






			'''simple docstring'''



			lowerCamelCase__ =     field(
			    metadata={"""help""": """Path to pretrained model or model identifier from huggingface.co/models"""}		)
			lowerCamelCase__ =     field(
			    default=lowercase,      metadata={"""help""": """Where do you want to store the pretrained models downloaded from huggingface.co"""},      )
			lowerCamelCase__ =     field(
			    default=lowercase,      metadata={"""help""": """Whether to freeze the feature extractor layers of the model."""}		)
			lowerCamelCase__ =     field(
			    default=lowercase,      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.999995,      metadata={"""help""": """Decay of gumbel temperature during training."""}		)







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






@dataclass
class lowerCAmelCase__      :






			'''simple docstring'''



			lowerCamelCase__ =     field(
			    default=lowercase,      metadata={"""help""": """The name of the dataset to use (via the datasets library)."""}		)
			lowerCamelCase__ =     field(
			    default=lowercase,      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=lowercase,      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=lowercase,      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 lowerCAmelCase__      :






			'''simple docstring'''



			lowerCamelCase__ =     42
			lowerCamelCase__ =     42
			lowerCamelCase__ =     "longest"
			lowerCamelCase__ =     None
			lowerCamelCase__ =     None
			def __call__(    self    , lowercase    ):
							# reformat list to dict and set to pytorch format
							_lowerCamelCase	:     Any       =       self.feature_extractor.pad(
							    lowercase    , max_length=self.max_length    , padding=self.padding    , pad_to_multiple_of=self.pad_to_multiple_of    , return_tensors='pt'    , )
							_lowerCamelCase	:     Optional[Any]       =       self.model._get_feat_extract_output_lengths(batch['input_values'].shape[-1]    )

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

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

							# sample randomly masked indices
							_lowerCamelCase	:     Any       =       _compute_mask_indices(
							    (batch_size, mask_indices_seq_length)    , self.model.config.mask_time_prob    , self.model.config.mask_time_length    , attention_mask=lowercase    , min_masks=2    , )

							return batch






class lowerCAmelCase__      (	lowercase		):






			'''simple docstring'''



			def __init__(    self    , *lowercase    , lowercase=1    , lowercase=0    , lowercase=1.0    , **lowercase    ):
							super().__init__(*lowercase    , **lowercase    )
							_lowerCamelCase	:     int       =       0
							_lowerCamelCase	:     Union[str, Any]       =       max_gumbel_temp
							_lowerCamelCase	:     str       =       min_gumbel_temp
							_lowerCamelCase	:     int       =       gumbel_temp_decay






			def    A_       (    self    , lowercase    , lowercase    ):
							model.train()
							_lowerCamelCase	:     int       =       self._prepare_inputs(lowercase    )

							if self.use_amp:
											with autocast():
															_lowerCamelCase	:     Any       =       self.compute_loss(lowercase    , lowercase    )
							else:
											_lowerCamelCase	:     List[str]       =       self.compute_loss(lowercase    , lowercase    )

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

							if self.use_amp:
											self.scaler.scale(lowercase    ).backward()
							elif self.use_apex:
											with amp.scale_loss(lowercase    , self.optimizer    ) as scaled_loss:
															scaled_loss.backward()
							elif self.deepspeed:
											self.deepspeed.backward(lowercase    )
							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 	_snake_case     (							):
				# 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.

				_lowerCamelCase	:     Dict       =       HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)						)

				_lowerCamelCase,					_lowerCamelCase,					_lowerCamelCase	:     int       =       parser.parse_args_into_dataclasses()
				configure_logger(lowercase__				,       lowercase__						)

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

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

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

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

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

				# normalize and transform to `BatchFeatures`
				_lowerCamelCase	:     Any       =       vectorized_datasets.map(
				    lowercase__				,       batched=lowercase__				,       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
				_lowerCamelCase	:     Tuple       =       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\''						)

				_lowerCamelCase	:     Optional[int]       =       WavaVecaForPreTraining(lowercase__						)

				_lowerCamelCase	:     Dict       =       DataCollatorForWavaVecaPretraining(model=lowercase__				,       feature_extractor=lowercase__						)

				_lowerCamelCase	:     Any       =       WavaVecaPreTrainer(
				    model=lowercase__				,       data_collator=lowercase__				,       args=lowercase__				,       train_dataset=vectorized_datasets['train']				,       eval_dataset=vectorized_datasets['validation']				,       tokenizer=lowercase__				,       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 ...utils import (
    OptionalDependencyNotAvailable,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)


try:
			if not (is_transformers_available() and is_torch_available() and is_transformers_version(""">=""", """4.25.0""")):
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			from ...utils.dummy_torch_and_transformers_objects import (
			    VersatileDiffusionDualGuidedPipeline,
			    VersatileDiffusionImageVariationPipeline,
			    VersatileDiffusionPipeline,
			    VersatileDiffusionTextToImagePipeline,
			)
else:
			from .modeling_text_unet import UNetFlatConditionModel
			from .pipeline_versatile_diffusion import VersatileDiffusionPipeline
			from .pipeline_versatile_diffusion_dual_guided import VersatileDiffusionDualGuidedPipeline
			from .pipeline_versatile_diffusion_image_variation import VersatileDiffusionImageVariationPipeline
			from .pipeline_versatile_diffusion_text_to_image import VersatileDiffusionTextToImagePipeline