



'''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'):
     lowerCAmelCase:   Optional[int]					=		True
     from torch.cuda.amp import autocast


lowerCAmelCase:   Optional[int]					=		logging.getLogger(__name__)

@dataclass
class 				a__:
    lowercase__					=       field(
        metadata={"""help""": """Path to pretrained model or model identifier from huggingface.co/models"""}      )
    lowercase__					=       field(
        default=lowerCamelCase__						,					metadata={"""help""": """Where do you want to store the pretrained models downloaded from huggingface.co"""}						,					)
    lowercase__					=       field(
        default=lowerCamelCase__						,					metadata={"""help""": """Whether to freeze the feature extractor layers of the model."""}      )
    lowercase__					=       field(
        default=lowerCamelCase__						,					metadata={"""help""": """Whether to log verbose messages or not."""}						,					)
    lowercase__					=       field(
        default=2.0						,					metadata={"""help""": """Maximum temperature for gumbel softmax."""}      )
    lowercase__					=       field(
        default=0.5						,					metadata={"""help""": """Minimum temperature for gumbel softmax."""}      )
    lowercase__					=       field(
        default=0.99_99_95						,					metadata={"""help""": """Decay of gumbel temperature during training."""}      )
def  lowerCamelCase__			(		_A					,     _A							):
  logging.basicConfig(
      format='%(asctime)s - %(levelname)s - %(name)s - %(message)s'					,     datefmt='%m/%d/%Y %H:%M:%S'					,     handlers=[logging.StreamHandler(sys.stdout							)]					,     )
  a		:   Optional[int]	  =	logging.WARNING
  if model_args.verbose_logging:
    a		:   Any	  =	logging.DEBUG
  elif trainer_utils.is_main_process(training_args.local_rank							):
    a		:   List[str]	  =	logging.INFO
  logger.setLevel(_A							)

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

@dataclass
class 				a__:
    lowercase__					=       42
    lowercase__					=       42
    lowercase__					=       "longest"
    lowercase__					=       None
    lowercase__					=       None







    def __call__( self					:    Optional[Any]		,     __snake_case					:    List[Dict[str, Union[List[int], torch.Tensor]]]      ):
      # reformat list to dict and set to pytorch format
      a		:   str	  =	self.feature_extractor.pad(
          __snake_case		,     max_length=self.max_length		,     padding=self.padding		,     pad_to_multiple_of=self.pad_to_multiple_of		,     return_tensors='pt'		,     )
      a		:   List[Any]	  =	self.model._get_feat_extract_output_lengths(batch['input_values'].shape[-1]      )

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

        a		:   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
        a		:   Union[str, Any]	  =	1
        a		:   Dict	  =	attention_mask.flip([-1]      ).cumsum(-1      ).flip([-1]      ).bool()

      # sample randomly masked indices
      a		:   Union[str, Any]	  =	_compute_mask_indices(
          (batch_size, mask_indices_seq_length)		,     self.model.config.mask_time_prob		,     self.model.config.mask_time_length		,     attention_mask=__snake_case		,     min_masks=2		,     )

      return batch

class 				a__(   lowerCamelCase__      ):







    def __init__( self					:    Any		,     *__snake_case					:    int		,     __snake_case					:    List[Any]=1		,     __snake_case					:    Union[str, Any]=0		,     __snake_case					:    Optional[Any]=1.0		,     **__snake_case					:    Union[str, Any]      ):
      super().__init__(*__snake_case		,     **__snake_case      )
      a		:   int	  =	0
      a		:   Tuple	  =	max_gumbel_temp
      a		:   Tuple	  =	min_gumbel_temp
      a		:   Any	  =	gumbel_temp_decay







    def      lowercase_			( self					:    Union[str, Any]		,     __snake_case					:    nn.Module		,     __snake_case					:    Dict[str, Union[torch.Tensor, Any]]      ):
      model.train()
      a		:   Dict	  =	self._prepare_inputs(__snake_case      )

      if self.use_amp:
        with autocast():
          a		:   int	  =	self.compute_loss(__snake_case		,     __snake_case      )
      else:
        a		:   Optional[int]	  =	self.compute_loss(__snake_case		,     __snake_case      )

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

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

  a		:   List[Any]	  =	HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)							)

  a  ,      a  ,      a		:   Dict	  =	parser.parse_args_into_dataclasses()
  configure_logger(_A					,     _A							)

  # Downloading and loading a dataset from the hub.
  a		:   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"
    a		:   str	  =	DatasetDict()
    a		:   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					,     )
    a		:   List[str]	  =	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		:   Optional[int]	  =	DatasetDict()
    a		:   int	  =	load_dataset(
        data_args.dataset_name					,     data_args.dataset_config_name					,     split='validation'					,     cache_dir=model_args.cache_dir					,     )
    a		:   Dict	  =	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		:   Tuple	  =	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		:   Union[str, Any]	  =	librosa.load(batch[data_args.speech_file_column]					,     sr=feature_extractor.sampling_rate							)
    return batch

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

  # filter audio files that are too long
  a		:   Union[str, Any]	  =	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		:   Dict	  =	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		:   str	  =	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		:   Optional[int]	  =	WavaVecaForPreTraining(_A							)

  a		:   str	  =	DataCollatorForWavaVecaPretraining(model=_A					,     feature_extractor=_A							)

  a		:   str	  =	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 ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCAmelCase:   Optional[int]					=		logging.get_logger(__name__)

lowerCAmelCase:   int					=		{'ctrl': 'https://huggingface.co/ctrl/resolve/main/config.json'}

class 				a__(   lowerCamelCase__      ):
    lowercase__					=       """ctrl"""
    lowercase__					=       ["""past_key_values"""]
    lowercase__					=       {
        """max_position_embeddings""": """n_positions""",
        """hidden_size""": """n_embd""",
        """num_attention_heads""": """n_head""",
        """num_hidden_layers""": """n_layer""",
    }







    def __init__( self					:    Optional[int]		,     __snake_case					:    Union[str, Any]=24_65_34		,     __snake_case					:    Dict=2_56		,     __snake_case					:    Optional[int]=12_80		,     __snake_case					:    Optional[int]=81_92		,     __snake_case					:    Union[str, Any]=48		,     __snake_case					:    str=16		,     __snake_case					:    Any=0.1		,     __snake_case					:    List[Any]=0.1		,     __snake_case					:    str=1e-6		,     __snake_case					:    List[str]=0.02		,     __snake_case					:    int=True		,     **__snake_case					:    Any		,     ):
      a		:   int	  =	vocab_size
      a		:   Union[str, Any]	  =	n_positions
      a		:   Optional[int]	  =	n_embd
      a		:   Dict	  =	n_layer
      a		:   Union[str, Any]	  =	n_head
      a		:   Tuple	  =	dff
      a		:   Union[str, Any]	  =	resid_pdrop
      a		:   Tuple	  =	embd_pdrop
      a		:   int	  =	layer_norm_epsilon
      a		:   str	  =	initializer_range

      a		:   Any	  =	use_cache

      super().__init__(**__snake_case      )