




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


_SCREAMING_SNAKE_CASE							:   Dict      		= logging.getLogger(__name__)




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


def      UpperCamelCase_(	snake_case							:       Tuple							,				snake_case							:       Any     ):

							'''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_						=      logging.WARNING
							if model_args.verbose_logging:
														snake_case_						=      logging.DEBUG
							elif trainer_utils.is_main_process(training_args.local_rank     ):
														snake_case_						=      logging.INFO
							logger.setLevel(__A     )




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




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




			def __call__(					self				,							a__							)				->				Dict[str, torch.Tensor]:

										'''simple docstring'''




										snake_case_						=      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_						=      self.model._get_feat_extract_output_lengths(batch["input_values"].shape[-1]							)

										snake_case_						=      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_						=      self.model._get_feat_extract_output_lengths(batch["attention_mask"].sum(-1							)							).to(
																	    torch.long							)

																	snake_case_						=      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_						=      1
																	snake_case_						=      attention_mask.flip([-1]							).cumsum(-1							).flip([-1]							).bool()

										# sample randomly masked indices
										snake_case_						=      _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  _snake_case		(	lowercase_							):




			def __init__(					self				,							*a__				,							a__=1				,							a__=0				,							a__=1.0				,							**a__							)				->				int:

										'''simple docstring'''




										super().__init__(*a__				,							**a__							)
										snake_case_						=      0
										snake_case_						=      max_gumbel_temp
										snake_case_						=      min_gumbel_temp
										snake_case_						=      gumbel_temp_decay




			def    lowerCAmelCase__				(					self				,							a__				,							a__							)				->				torch.Tensor:

										'''simple docstring'''




										model.train()
										snake_case_						=      self._prepare_inputs(a__							)

										if self.use_amp:
																	with autocast():
																								snake_case_						=      self.compute_loss(a__				,							a__							)
										else:
																	snake_case_						=      self.compute_loss(a__				,							a__							)

										if self.args.n_gpu > 1 or self.deepspeed:
																	if model.module.config.ctc_loss_reduction == "mean":
																								snake_case_						=      loss.mean()
																	elif model.module.config.ctc_loss_reduction == "sum":
																								snake_case_						=      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_						=      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_(	):

							'''simple docstring'''






							snake_case_						=      HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)     )

							snake_case_					,      snake_case_					,      snake_case_						=      parser.parse_args_into_dataclasses()
							configure_logger(__A							,				__A     )

							# Downloading and loading a dataset from the hub.
							snake_case_						=      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_						=      DatasetDict()
														snake_case_						=      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_						=      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_						=      DatasetDict()
														snake_case_						=      load_dataset(
														    data_args.dataset_name							,				data_args.dataset_config_name							,				split="validation"							,				cache_dir=model_args.cache_dir							,				)
														snake_case_						=      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_						=      WavaVecaFeatureExtractor.from_pretrained(
							    model_args.model_name_or_path							,				cache_dir=model_args.cache_dir							,				do_normalize=__A     )

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

							# load audio files into numpy arrays
							snake_case_						=      datasets.map(
							    __A							,				num_proc=data_args.preprocessing_num_workers							,				remove_columns=datasets["train"].column_names     )

							# filter audio files that are too long
							snake_case_						=      vectorized_datasets.filter(
							    lambda snake_case     : len(data["speech"]     ) < int(data_args.max_duration_in_seconds * feature_extractor.sampling_rate     )     )

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

							# normalize and transform to `BatchFeatures`
							snake_case_						=      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
							snake_case_						=      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_						=      WavaVecaForPreTraining(__A     )

							snake_case_						=      DataCollatorForWavaVecaPretraining(model=__A							,				feature_extractor=__A     )

							snake_case_						=      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'''





# DISCLAIMER: This file is strongly influenced by https://github.com/ermongroup/ddim

from dataclasses import dataclass
from typing import Optional, Tuple, Union

import flax
import jax
import jax.numpy as jnp

from ..configuration_utils import ConfigMixin, register_to_config
from .scheduling_utils_flax import (
    CommonSchedulerState,
    FlaxKarrasDiffusionSchedulers,
    FlaxSchedulerMixin,
    FlaxSchedulerOutput,
    add_noise_common,
    get_velocity_common,
)






@flax.struct.dataclass
class __A			:



 '''simple docstring'''







 __lowerCamelCase    :       CommonSchedulerState

 # setable values
 __lowerCamelCase    :       jnp.ndarray
 __lowerCamelCase    :       jnp.ndarray
 __lowerCamelCase    :       Optional[int]          =       None



 @classmethod
 def 					a__  (cls ,   A ,   A ,   A       )    ->     str:


      """simple docstring"""



      return cls(common=A ,   init_noise_sigma=A ,   timesteps=A       )






@dataclass
class __A			(  A						):



 '''simple docstring'''







 __lowerCamelCase    :       DDPMSchedulerState






class __A			(  A   ,   A						):



 '''simple docstring'''







 __lowerCamelCase    :       Dict          =       [e.name for e in FlaxKarrasDiffusionSchedulers]

 __lowerCamelCase    :       jnp.dtype



 @property
 def 					a__  (self       )    ->     List[str]:


      """simple docstring"""



      return True



 @register_to_config
 def __init__(self ,   A = 1_000 ,   A = 0.0001 ,   A = 0.02 ,   A = "linear" ,   A = None ,   A = "fixed_small" ,   A = True ,   A = "epsilon" ,   A = jnp.floataa ,   )    ->     Union[str, Any]:


      """simple docstring"""



      _a      =  dtype



 def 					a__  (self ,   A = None       )    ->     DDPMSchedulerState:


      """simple docstring"""



      if common is None:
           _a      =  CommonSchedulerState.create(self       )

      # standard deviation of the initial noise distribution
      _a      =  jnp.array(1.0 ,   dtype=self.dtype       )

      _a      =  jnp.arange(0 ,   self.config.num_train_timesteps       ).round()[::-1]

      return DDPMSchedulerState.create(
          common=A ,   init_noise_sigma=A ,   timesteps=A ,   )



 def 					a__  (self ,   A ,   A ,   A = None       )    ->     jnp.ndarray:


      """simple docstring"""



      return sample



 def 					a__  (self ,   A ,   A ,   A = ()       )    ->     DDPMSchedulerState:


      """simple docstring"""



      _a      =  self.config.num_train_timesteps // num_inference_steps
      # creates integer timesteps by multiplying by ratio
      # rounding to avoid issues when num_inference_step is power of 3
      _a      =  (jnp.arange(0 ,   A       ) * step_ratio).round()[::-1]

      return state.replace(
          num_inference_steps=A ,   timesteps=A ,   )



 def 					a__  (self ,   A ,   A ,   A=None ,   A=None       )    ->     int:


      """simple docstring"""



      _a      =  state.common.alphas_cumprod[t]
      _a      =  jnp.where(t > 0 ,   state.common.alphas_cumprod[t - 1] ,   jnp.array(1.0 ,   dtype=self.dtype       )       )

      # For t > 0, compute predicted variance βt (see formula (6) and (7) from https://arxiv.org/pdf/2006.11239.pdf)
      # and sample from it to get previous sample
      # x_{t-1} ~ N(pred_prev_sample, variance) == add variance to pred_sample
      _a      =  (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * state.common.betas[t]

      if variance_type is None:
           _a      =  self.config.variance_type

      # hacks - were probably added for training stability
      if variance_type == "fixed_small":
           _a      =  jnp.clip(A ,   a_min=1E-20       )
      # for rl-diffuser https://arxiv.org/abs/2205.09991
      elif variance_type == "fixed_small_log":
           _a      =  jnp.log(jnp.clip(A ,   a_min=1E-20       )       )
      elif variance_type == "fixed_large":
           _a      =  state.common.betas[t]
      elif variance_type == "fixed_large_log":
           # Glide max_log
           _a      =  jnp.log(state.common.betas[t]       )
      elif variance_type == "learned":
           return predicted_variance
      elif variance_type == "learned_range":
           _a      =  variance
           _a      =  state.common.betas[t]
           _a      =  (predicted_variance + 1) / 2
           _a      =  frac * max_log + (1 - frac) * min_log

      return variance



 def 					a__  (self ,   A ,   A ,   A ,   A ,   A = None ,   A = True ,   )    ->     Union[FlaxDDPMSchedulerOutput, Tuple]:


      """simple docstring"""



      _a      =  timestep

      if key is None:
           _a      =  jax.random.PRNGKey(0       )

      if model_output.shape[1] == sample.shape[1] * 2 and self.config.variance_type in ["learned", "learned_range"]:
           _a  ,    _a      =  jnp.split(A ,   sample.shape[1] ,   axis=1       )
      else:
           _a      =  None

      # 1. compute alphas, betas
      _a      =  state.common.alphas_cumprod[t]
      _a      =  jnp.where(t > 0 ,   state.common.alphas_cumprod[t - 1] ,   jnp.array(1.0 ,   dtype=self.dtype       )       )
      _a      =  1 - alpha_prod_t
      _a      =  1 - alpha_prod_t_prev

      # 2. compute predicted original sample from predicted noise also called
      # "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
      if self.config.prediction_type == "epsilon":
           _a      =  (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
      elif self.config.prediction_type == "sample":
           _a      =  model_output
      elif self.config.prediction_type == "v_prediction":
           _a      =  (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
      else:
           raise ValueError(
               f'''prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` '''
               ''' for the FlaxDDPMScheduler.'''       )

      # 3. Clip "predicted x_0"
      if self.config.clip_sample:
           _a      =  jnp.clip(A ,   -1 ,   1       )

      # 4. Compute coefficients for pred_original_sample x_0 and current sample x_t
      # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
      _a      =  (alpha_prod_t_prev ** 0.5 * state.common.betas[t]) / beta_prod_t
      _a      =  state.common.alphas[t] ** 0.5 * beta_prod_t_prev / beta_prod_t

      # 5. Compute predicted previous sample µ_t
      # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
      _a      =  pred_original_sample_coeff * pred_original_sample + current_sample_coeff * sample

      # 6. Add noise
      def random_variance():
           _a      =  jax.random.split(A ,   num=1       )
           _a      =  jax.random.normal(A ,   shape=model_output.shape ,   dtype=self.dtype       )
           return (self._get_variance(A ,   A ,   predicted_variance=A       ) ** 0.5) * noise

      _a      =  jnp.where(t > 0 ,   random_variance() ,   jnp.zeros(model_output.shape ,   dtype=self.dtype       )       )

      _a      =  pred_prev_sample + variance

      if not return_dict:
           return (pred_prev_sample, state)

      return FlaxDDPMSchedulerOutput(prev_sample=A ,   state=A       )



 def 					a__  (self ,   A ,   A ,   A ,   A ,   )    ->     jnp.ndarray:


      """simple docstring"""



      return add_noise_common(state.common ,   A ,   A ,   A       )



 def 					a__  (self ,   A ,   A ,   A ,   A ,   )    ->     jnp.ndarray:


      """simple docstring"""



      return get_velocity_common(state.common ,   A ,   A ,   A       )



 def __len__(self       )    ->     Tuple:


      """simple docstring"""



      return self.config.num_train_timesteps



