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


A__ :       Optional[Any]				= logging.getLogger(__name__)


@dataclass
class       __magic_name__							:
       UpperCamelCase_ =		field(
           metadata={'''help''': '''Path to pretrained model or model identifier from huggingface.co/models'''}			)
       UpperCamelCase_ =		field(
           default=a__ ,       metadata={'''help''': '''Where do you want to store the pretrained models downloaded from huggingface.co'''} ,       )
       UpperCamelCase_ =		field(
           default=a__ ,       metadata={'''help''': '''Whether to freeze the feature extractor layers of the model.'''}			)
       UpperCamelCase_ =		field(
           default=a__ ,       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.9_9_9_9_9_5 ,       metadata={'''help''': '''Decay of gumbel temperature during training.'''}			)




def 			_lowerCAmelCase		(		_UpperCamelCase			,					_UpperCamelCase    ):






  """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:      int   			=  logging.WARNING
  if model_args.verbose_logging:
    _lowercase:      List[Any]   			=  logging.DEBUG
  elif trainer_utils.is_main_process(training_args.local_rank    ):
    _lowercase:      Union[str, Any]   			=  logging.INFO
  logger.setLevel(A__    )


@dataclass
class       __magic_name__							:
       UpperCamelCase_ =		field(
           default=a__ ,       metadata={'''help''': '''The name of the dataset to use (via the datasets library).'''}			)
       UpperCamelCase_ =		field(
           default=a__ ,       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=a__ ,       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=a__ ,       metadata={'''help''': '''The number of processes to use for the preprocessing.'''} ,       )
       UpperCamelCase_ =		field(
           default=2_0.0 ,       metadata={'''help''': '''Filter audio files that are longer than `max_duration_in_seconds` seconds'''}			)


@dataclass
class       __magic_name__							:
       UpperCamelCase_ =		42
       UpperCamelCase_ =		42
       UpperCamelCase_ =		'''longest'''
       UpperCamelCase_ =		None
       UpperCamelCase_ =		None
       def __call__(  self			,					A_							)  ->						Dict[str, torch.Tensor]:





         """simple docstring"""




         _lowercase:      int   			=  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'''			,					)
         _lowercase:      Optional[Any]   			=  self.model._get_feat_extract_output_lengths(batch['''input_values'''].shape[-1]							)

         _lowercase:      str   			=  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:      Union[str, Any]   			=  self.model._get_feat_extract_output_lengths(batch['''attention_mask'''].sum(-1							)							).to(
               torch.long							)

           _lowercase:      List[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
           _lowercase:      Dict   			=  1
           _lowercase:      List[str]   			=  attention_mask.flip([-1]							).cumsum(-1							).flip([-1]							).bool()

         # sample randomly masked indices
         _lowercase:      Optional[Any]   			=  _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       __magic_name__							(							a__			):
       def __init__(  self			,					*A_			,					A_=1			,					A_=0			,					A_=1.0			,					**A_							)  ->						Any:





         """simple docstring"""




         super().__init__(*__a			,					**__a							)
         _lowercase:      List[str]   			=  0
         _lowercase:      int   			=  max_gumbel_temp
         _lowercase:      List[Any]   			=  min_gumbel_temp
         _lowercase:      Optional[Any]   			=  gumbel_temp_decay



       def 				lowercase_				(  self			,					A_			,					A_							)  ->						torch.Tensor:





         """simple docstring"""




         model.train()
         _lowercase:      Dict   			=  self._prepare_inputs(__a							)

         if self.use_amp:
           with autocast():
             _lowercase:      str   			=  self.compute_loss(__a			,					__a							)
         else:
           _lowercase:      Dict   			=  self.compute_loss(__a			,					__a							)

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






  """simple docstring"""



  _lowercase:      Tuple   			=  HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)    )

  _lowercase:      Union[str, Any]   			=  parser.parse_args_into_dataclasses()
  configure_logger(A__			,					A__    )

  # Downloading and loading a dataset from the hub.
  _lowercase:      List[str]   			=  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:      int   			=  DatasetDict()
    _lowercase:      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			,					)
    _lowercase:      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"
    _lowercase:      str   			=  DatasetDict()
    _lowercase:      Optional[Any]   			=  load_dataset(
        data_args.dataset_name			,					data_args.dataset_config_name			,					split='''validation'''			,					cache_dir=model_args.cache_dir			,					)
    _lowercase:      List[str]   			=  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:      str   			=  WavaVecaFeatureExtractor.from_pretrained(
      model_args.model_name_or_path			,					cache_dir=model_args.cache_dir			,					do_normalize=A__    )

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

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

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

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

  # normalize and transform to `BatchFeatures`
  _lowercase:      Tuple   			=  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
  _lowercase:      Dict   			=  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:      Union[str, Any]   			=  WavaVecaForPreTraining(A__    )

  _lowercase:      Dict   			=  DataCollatorForWavaVecaPretraining(model=A__			,					feature_extractor=A__    )

  _lowercase:      Optional[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()




from dataclasses import dataclass, field
from typing import Optional

from transformers import AutoConfig, AutoImageProcessor, AutoTokenizer, FlaxVisionEncoderDecoderModel, HfArgumentParser





@dataclass
class   SCREAMING_SNAKE_CASE_			:




       '''simple docstring'''



       _a   							=							field(
           metadata={"help": "The output directory where the model will be written."}						,       )
       _a   							=							field(
           metadata={
               "help": (
                   "The encoder model checkpoint for weights initialization."
                   "Don't set if you want to train an encoder model from scratch."
               )
           }						,       )
       _a   							=							field(
           metadata={
               "help": (
                   "The decoder model checkpoint for weights initialization."
                   "Don't set if you want to train a decoder model from scratch."
               )
           }						,       )
       _a   							=							field(
           default=a__						,       metadata={"help": "Pretrained encoder config name or path if not the same as encoder_model_name"}						)
       _a   							=							field(
           default=a__						,       metadata={"help": "Pretrained decoder config name or path if not the same as decoder_model_name"}						)







def    __lowerCamelCase   (      )     ->    Optional[Any]:
   lowerCamelCase_						:   List[Any]		= HfArgumentParser((ModelArguments,)					)
   ((lowerCamelCase_),	)						:   List[Any]		= parser.parse_args_into_dataclasses()

   # Load pretrained model and tokenizer

   # Use explicit specified encoder config
   if model_args.encoder_config_name:
      lowerCamelCase_						:   Any		= AutoConfig.from_pretrained(model_args.encoder_config_name					)
   # Use pretrained encoder model's config
   else:
      lowerCamelCase_						:   Dict		= AutoConfig.from_pretrained(model_args.encoder_model_name_or_path					)

   # Use explicit specified decoder config
   if model_args.decoder_config_name:
      lowerCamelCase_						:   Dict		= AutoConfig.from_pretrained(model_args.decoder_config_name					)
   # Use pretrained decoder model's config
   else:
      lowerCamelCase_						:   Any		= AutoConfig.from_pretrained(model_args.decoder_model_name_or_path					)

   # necessary for `from_encoder_decoder_pretrained` when `decoder_config` is passed
   lowerCamelCase_						:   Optional[Any]		= True
   lowerCamelCase_						:   List[Any]		= True

   lowerCamelCase_						:   List[Any]		= FlaxVisionEncoderDecoderModel.from_encoder_decoder_pretrained(
       encoder_pretrained_model_name_or_path=model_args.encoder_model_name_or_path						,						decoder_pretrained_model_name_or_path=model_args.decoder_model_name_or_path						,						encoder_config=A__						,						decoder_config=A__						,						)

   # GPT2 only has bos/eos tokens but not decoder_start/pad tokens
   lowerCamelCase_						:   Tuple		= decoder_config.decoder_start_token_id
   lowerCamelCase_						:   int		= decoder_config.pad_token_id
   if decoder_start_token_id is None:
      lowerCamelCase_						:   Any		= decoder_config.bos_token_id
   if pad_token_id is None:
      lowerCamelCase_						:   List[str]		= decoder_config.eos_token_id

   # This is necessary to make Flax's generate() work
   lowerCamelCase_						:   List[str]		= decoder_config.eos_token_id
   lowerCamelCase_						:   Tuple		= decoder_start_token_id
   lowerCamelCase_						:   Optional[Any]		= pad_token_id

   lowerCamelCase_						:   str		= AutoImageProcessor.from_pretrained(model_args.encoder_model_name_or_path					)

   lowerCamelCase_						:   Any		= AutoTokenizer.from_pretrained(model_args.decoder_model_name_or_path					)
   lowerCamelCase_						:   Dict		= tokenizer.convert_ids_to_tokens(model.config.pad_token_id					)

   model.save_pretrained(model_args.output_dir					)
   image_processor.save_pretrained(model_args.output_dir					)
   tokenizer.save_pretrained(model_args.output_dir					)


if __name__ == "__main__":
      main()





