



"""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    :Any     =					logging.getLogger(__name__)
@dataclass
class 							_UpperCAmelCase							:





     '''simple docstring'''



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





def  lowerCamelCase_						(UpperCamelCase__		:		ModelArguments     ,       UpperCamelCase__		:		TrainingArguments							):
    logging.basicConfig(
        format='''%(asctime)s - %(levelname)s - %(name)s - %(message)s'''     ,       datefmt='''%m/%d/%Y %H:%M:%S'''     ,       handlers=[logging.StreamHandler(sys.stdout							)]     ,       )
    _UpperCAmelCase							:  str											=    logging.WARNING
    if model_args.verbose_logging:
        _UpperCAmelCase							:  Union[str, Any]											=    logging.DEBUG
    elif trainer_utils.is_main_process(training_args.local_rank							):
        _UpperCAmelCase							:  int											=    logging.INFO
    logger.setLevel(UpperCamelCase__							)
@dataclass
class 							_UpperCAmelCase							:





     '''simple docstring'''



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



@dataclass
class 							_UpperCAmelCase							:





     '''simple docstring'''



     a__									=42
     a__									=42
     a__									="longest"
     a__									=None
     a__									=None





     def __call__(					self       ,	A		)							->					Dict[str, torch.Tensor]:
         # reformat list to dict and set to pytorch format
         _UpperCAmelCase							:  Optional[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'''       ,	)
         _UpperCAmelCase							:  Tuple											=    self.model._get_feat_extract_output_lengths(batch['''input_values'''].shape[-1]		)

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

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

         # sample randomly masked indices
         _UpperCAmelCase							:  Dict											=    _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 							_UpperCAmelCase							( a      ):





     '''simple docstring'''



     def __init__(					self       ,	*A       ,	A=1       ,	A=0       ,	A=1.0       ,	**A		)							->					List[str]:
         super().__init__(*A       ,	**A		)
         _UpperCAmelCase							:  str											=    0
         _UpperCAmelCase							:  List[str]											=    max_gumbel_temp
         _UpperCAmelCase							:  List[Any]											=    min_gumbel_temp
         _UpperCAmelCase							:  str											=    gumbel_temp_decay





     def  __lowerCAmelCase	(					self       ,	A       ,	A		)							->					torch.Tensor:
         model.train()
         _UpperCAmelCase							:  Union[str, Any]											=    self._prepare_inputs(A		)

         if self.use_amp:
             with autocast():
                 _UpperCAmelCase							:  Tuple											=    self.compute_loss(A       ,	A		)
         else:
             _UpperCAmelCase							:  List[str]											=    self.compute_loss(A       ,	A		)

         if self.args.n_gpu > 1 or self.deepspeed:
             if model.module.config.ctc_loss_reduction == "mean":
                 _UpperCAmelCase							:  str											=    loss.mean()
             elif model.module.config.ctc_loss_reduction == "sum":
                 _UpperCAmelCase							:  Tuple											=    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:
             _UpperCAmelCase							:  Union[str, Any]											=    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_						():
    # 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.

    _UpperCAmelCase							:  Optional[int]											=    HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)							)

    _UpperCAmelCase							,				_UpperCAmelCase							,				_UpperCAmelCase							:  Any											=    parser.parse_args_into_dataclasses()
    configure_logger(UpperCamelCase__     ,       UpperCamelCase__							)

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

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

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

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

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

    # normalize and transform to `BatchFeatures`
    _UpperCAmelCase							:  int											=    vectorized_datasets.map(
        UpperCamelCase__     ,       batched=UpperCamelCase__     ,       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
    _UpperCAmelCase							:  Optional[Any]											=    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\''''							)

    _UpperCAmelCase							:  Optional[Any]											=    WavaVecaForPreTraining(UpperCamelCase__							)

    _UpperCAmelCase							:  Any											=    DataCollatorForWavaVecaPretraining(model=UpperCamelCase__     ,       feature_extractor=UpperCamelCase__							)

    _UpperCAmelCase							:  Tuple											=    WavaVecaPreTrainer(
        model=UpperCamelCase__     ,       data_collator=UpperCamelCase__     ,       args=UpperCamelCase__     ,       train_dataset=vectorized_datasets['''train''']     ,       eval_dataset=vectorized_datasets['''validation''']     ,       tokenizer=UpperCamelCase__     ,       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 typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


_lowerCAmelCase    :str     =					{
    'configuration_squeezebert': [
        'SQUEEZEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'SqueezeBertConfig',
        'SqueezeBertOnnxConfig',
    ],
    'tokenization_squeezebert': ['SqueezeBertTokenizer'],
}

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase    :Optional[int]     =					['SqueezeBertTokenizerFast']

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _lowerCAmelCase    :str     =					[
      'SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
      'SqueezeBertForMaskedLM',
      'SqueezeBertForMultipleChoice',
      'SqueezeBertForQuestionAnswering',
      'SqueezeBertForSequenceClassification',
      'SqueezeBertForTokenClassification',
      'SqueezeBertModel',
      'SqueezeBertModule',
      'SqueezeBertPreTrainedModel',
  ]


if TYPE_CHECKING:
  from .configuration_squeezebert import (
      SQUEEZEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
      SqueezeBertConfig,
      SqueezeBertOnnxConfig,
  )
  from .tokenization_squeezebert import SqueezeBertTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_squeezebert_fast import SqueezeBertTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_squeezebert import (
        SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        SqueezeBertForMaskedLM,
        SqueezeBertForMultipleChoice,
        SqueezeBertForQuestionAnswering,
        SqueezeBertForSequenceClassification,
        SqueezeBertForTokenClassification,
        SqueezeBertModel,
        SqueezeBertModule,
        SqueezeBertPreTrainedModel,
    )

else:
  import sys

  _lowerCAmelCase    :Optional[Any]     =					_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)


