



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



 '''simple docstring'''







 __lowerCamelCase    :       str          =       field(
     metadata={'help': 'Path to pretrained model or model identifier from huggingface.co/models'}						)
 __lowerCamelCase    :       Optional[str]          =       field(
     default=A   ,   metadata={'help': 'Where do you want to store the pretrained models downloaded from huggingface.co'}   ,   )
 __lowerCamelCase    :       Optional[bool]          =       field(
     default=A   ,   metadata={'help': 'Whether to freeze the feature extractor layers of the model.'}						)
 __lowerCamelCase    :       Optional[bool]          =       field(
     default=A   ,   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.99_99_95   ,   metadata={'help': 'Decay of gumbel temperature during training.'}						)





def lowerCAmelCase  (__A  ,	__A):



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






@dataclass
class __A			:



 '''simple docstring'''







 __lowerCamelCase    :       str          =       field(
     default=A   ,   metadata={'help': 'The name of the dataset to use (via the datasets library).'}						)
 __lowerCamelCase    :       Optional[str]          =       field(
     default=A   ,   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=A   ,   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=A   ,   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 __A			:



 '''simple docstring'''







 __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"""



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

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

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

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



 '''simple docstring'''







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


      """simple docstring"""



      super().__init__(*A ,   **A       )
      _a      =  0
      _a      =  max_gumbel_temp
      _a      =  min_gumbel_temp
      _a      =  gumbel_temp_decay



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


      """simple docstring"""



      model.train()
      _a      =  self._prepare_inputs(A       )

      if self.use_amp:
           with autocast():
                _a      =  self.compute_loss(A ,   A       )
      else:
           _a      =  self.compute_loss(A ,   A       )

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



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

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

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

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

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

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

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

     # normalize and transform to `BatchFeatures`
     _a      =  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
     _a      =  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\'''')

     _a      =  WavaVecaForPreTraining(__A)

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

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





from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


lowercase_		       =  logging.get_logger(__name__)

lowercase_		       =  {
    "microsoft/resnet-50": "https://huggingface.co/microsoft/resnet-50/blob/main/config.json",
}






class __A			(  A   ,   A						):



 '''simple docstring'''







 __lowerCamelCase    :       Tuple          =       'resnet'
 __lowerCamelCase    :       Any          =       ['basic', 'bottleneck']



 def __init__(self ,   A=3 ,   A=64 ,   A=[256, 512, 1_024, 2_048] ,   A=[3, 4, 6, 3] ,   A="bottleneck" ,   A="relu" ,   A=False ,   A=None ,   A=None ,   **A ,   )    ->     Dict:


      """simple docstring"""



      super().__init__(**A       )
      if layer_type not in self.layer_types:
           raise ValueError(f'''layer_type={layer_type} is not one of {','.join(self.layer_types       )}'''       )
      _a      =  num_channels
      _a      =  embedding_size
      _a      =  hidden_sizes
      _a      =  depths
      _a      =  layer_type
      _a      =  hidden_act
      _a      =  downsample_in_first_stage
      _a      =  ['''stem'''] + [f'''stage{idx}''' for idx in range(1 ,   len(A       ) + 1       )]
      _a  ,    _a      =  get_aligned_output_features_output_indices(
          out_features=A ,   out_indices=A ,   stage_names=self.stage_names       )






class __A			(  A						):



 '''simple docstring'''







 __lowerCamelCase    :       Any          =       version.parse('1.11'						)



 @property
 def 					a__  (self       )    ->     Mapping[str, Mapping[int, str]]:


      """simple docstring"""



      return OrderedDict(
          [
              ('''pixel_values''', {0: '''batch''', 1: '''num_channels''', 2: '''height''', 3: '''width'''}),
          ]       )



 @property
 def 					a__  (self       )    ->     float:


      """simple docstring"""



      return 1E-3



