




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


lowerCamelCase    :		str			   =       logging.getLogger(__name__)







@dataclass
class   A:



 '''simple docstring'''


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

def       _SCREAMING_SNAKE_CASE      (							lowercase     :	ModelArguments     ,							lowercase     :	TrainingArguments					):


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







@dataclass
class   A:



 '''simple docstring'''


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







@dataclass
class   A:



 '''simple docstring'''


 UpperCamelCase						       =						42
 UpperCamelCase						       =						42
 UpperCamelCase						       =						"longest"
 UpperCamelCase						       =						None
 UpperCamelCase						       =						None

 def __call__(    self :  Optional[Any]						,							A_ :  List[Dict[str, Union[List[int], torch.Tensor]]]  )       ->     Dict:

       """simple docstring"""






       lowerCamelCase_							=   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'						,							)
       lowerCamelCase_							=   self.model._get_feat_extract_output_lengths(batch['input_values'].shape[-1]  )

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

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

       # sample randomly masked indices
       lowerCamelCase_							=   _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   A(						__lowerCamelCase       ):



 '''simple docstring'''


 def __init__(    self :  List[Any]						,							*A_ :  int						,							A_ :  Any=1						,							A_ :  List[Any]=0						,							A_ :  List[str]=1.0						,							**A_ :  Dict  )       ->     int:

       """simple docstring"""






       super().__init__(*UpperCamelCase_						,							**UpperCamelCase_  )
       lowerCamelCase_							=   0
       lowerCamelCase_							=   max_gumbel_temp
       lowerCamelCase_							=   min_gumbel_temp
       lowerCamelCase_							=   gumbel_temp_decay





 def     a__			(    self :  Optional[int]						,							A_ :  nn.Module						,							A_ :  Dict[str, Union[torch.Tensor, Any]]  )       ->     Any:

       """simple docstring"""






       model.train()
       lowerCamelCase_							=   self._prepare_inputs(UpperCamelCase_  )

       if self.use_amp:
             with autocast():
                   lowerCamelCase_							=   self.compute_loss(UpperCamelCase_						,							UpperCamelCase_  )
       else:
             lowerCamelCase_							=   self.compute_loss(UpperCamelCase_						,							UpperCamelCase_  )

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


      '''simple docstring'''



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

      lowerCamelCase_					,    lowerCamelCase_					,    lowerCamelCase_							=   parser.parse_args_into_dataclasses()
      configure_logger(A__     ,							A__					)

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

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

      # load audio files into numpy arrays
      lowerCamelCase_							=   datasets.map(
          A__     ,							num_proc=data_args.preprocessing_num_workers     ,							remove_columns=datasets['train'].column_names					)

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

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

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

      lowerCamelCase_							=   DataCollatorForWavaVecaPretraining(model=A__     ,							feature_extractor=A__					)

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









from ..utils import is_flax_available, is_torch_available


if is_torch_available():
 from .autoencoder_kl import AutoencoderKL
 from .controlnet import ControlNetModel
 from .dual_transformer_ad import DualTransformeraDModel
 from .modeling_utils import ModelMixin
 from .prior_transformer import PriorTransformer
 from .ta_film_transformer import TaFilmDecoder
 from .transformer_ad import TransformeraDModel
 from .unet_ad import UNetaDModel
 from .unet_ad import UNetaDModel
 from .unet_ad_condition import UNetaDConditionModel
 from .unet_ad_condition import UNetaDConditionModel
 from .vq_model import VQModel

if is_flax_available():
 from .controlnet_flax import FlaxControlNetModel
 from .unet_ad_condition_flax import FlaxUNetaDConditionModel
 from .vae_flax import FlaxAutoencoderKL



