






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


SCREAMING_SNAKE_CASE						=	logging.getLogger(__name__)
@dataclass
class UpperCAmelCase_					:




      """simple docstring"""




      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.999995				,       metadata={'''help''': '''Decay of gumbel temperature during training.'''}				)





def 							lowercase_				(    __A :       ModelArguments					,      __A :       TrainingArguments			)       ->      int:



 """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					:   Optional[int]				  =logging.WARNING
 if model_args.verbose_logging:
  lowercase					:   str				  =logging.DEBUG
 elif trainer_utils.is_main_process(training_args.local_rank			):
  lowercase					:   Union[str, Any]				  =logging.INFO
 logger.setLevel(__A			)
@dataclass
class UpperCAmelCase_					:




      """simple docstring"""




      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=20.0				,       metadata={'''help''': '''Filter audio files that are longer than `max_duration_in_seconds` seconds'''}				)





@dataclass
class UpperCAmelCase_					:




      """simple docstring"""




      UpperCamelCase_								=			42
      UpperCamelCase_								=			42
      UpperCamelCase_								=			"longest"
      UpperCamelCase_								=			None
      UpperCamelCase_								=			None






      def __call__( self			:		Tuple     ,      UpperCAmelCase			:		List[Dict[str, Union[List[int], torch.Tensor]]]		)  ->  Dict[str, torch.Tensor]:


       '''simple docstring'''

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

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

        lowercase					:   Optional[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					:   Tuple				  =1
        lowercase					:   Optional[Any]				  =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=UpperCAmelCase     ,      min_masks=2     ,      )

       return batch





class UpperCAmelCase_					(    __A				):




      """simple docstring"""




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


       '''simple docstring'''

       super().__init__(*UpperCAmelCase     ,      **UpperCAmelCase		)
       lowercase					:   Any				  =0
       lowercase					:   List[str]				  =max_gumbel_temp
       lowercase					:   List[str]				  =min_gumbel_temp
       lowercase					:   str				  =gumbel_temp_decay






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


       '''simple docstring'''

       model.train()
       lowercase					:   Any				  =self._prepare_inputs(UpperCAmelCase		)

       if self.use_amp:
        with autocast():
         lowercase					:   Union[str, Any]				  =self.compute_loss(UpperCAmelCase     ,      UpperCAmelCase		)
       else:
        lowercase					:   List[str]				  =self.compute_loss(UpperCAmelCase     ,      UpperCAmelCase		)

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



 """simple docstring"""




 lowercase					:   List[str]				  =HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)			)

 lowercase  ,		lowercase  ,		lowercase					:   Union[str, Any]				  =parser.parse_args_into_dataclasses()
 configure_logger(__A					,      __A			)

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

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

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

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

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

 # normalize and transform to `BatchFeatures`
 lowercase					:   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
 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					:   Dict				  =WavaVecaForPreTraining(__A			)

 lowercase					:   Any				  =DataCollatorForWavaVecaPretraining(model=__A					,      feature_extractor=__A			)

 lowercase					:   Optional[int]				  =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 typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


SCREAMING_SNAKE_CASE						=	{
    'configuration_rembert': ['REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'RemBertConfig', 'RemBertOnnxConfig']
}

try:
      if not is_sentencepiece_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      SCREAMING_SNAKE_CASE						=	['RemBertTokenizer']

try:
      if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      SCREAMING_SNAKE_CASE						=	['RemBertTokenizerFast']

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      SCREAMING_SNAKE_CASE						=	[
          'REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
          'RemBertForCausalLM',
          'RemBertForMaskedLM',
          'RemBertForMultipleChoice',
          'RemBertForQuestionAnswering',
          'RemBertForSequenceClassification',
          'RemBertForTokenClassification',
          'RemBertLayer',
          'RemBertModel',
          'RemBertPreTrainedModel',
          'load_tf_weights_in_rembert',
      ]


try:
      if not is_tf_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      SCREAMING_SNAKE_CASE						=	[
          'TF_REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST',
          'TFRemBertForCausalLM',
          'TFRemBertForMaskedLM',
          'TFRemBertForMultipleChoice',
          'TFRemBertForQuestionAnswering',
          'TFRemBertForSequenceClassification',
          'TFRemBertForTokenClassification',
          'TFRemBertLayer',
          'TFRemBertModel',
          'TFRemBertPreTrainedModel',
      ]


if TYPE_CHECKING:
      from .configuration_rembert import REMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, RemBertConfig, RemBertOnnxConfig

      try:
            if not is_sentencepiece_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_rembert import RemBertTokenizer

      try:
            if not is_tokenizers_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_rembert_fast import RemBertTokenizerFast

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_rembert import (
                REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
                RemBertForCausalLM,
                RemBertForMaskedLM,
                RemBertForMultipleChoice,
                RemBertForQuestionAnswering,
                RemBertForSequenceClassification,
                RemBertForTokenClassification,
                RemBertLayer,
                RemBertModel,
                RemBertPreTrainedModel,
                load_tf_weights_in_rembert,
            )

      try:
            if not is_tf_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_tf_rembert import (
                TF_REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
                TFRemBertForCausalLM,
                TFRemBertForMaskedLM,
                TFRemBertForMultipleChoice,
                TFRemBertForQuestionAnswering,
                TFRemBertForSequenceClassification,
                TFRemBertForTokenClassification,
                TFRemBertLayer,
                TFRemBertModel,
                TFRemBertPreTrainedModel,
            )


else:
      import sys

      SCREAMING_SNAKE_CASE						=	_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)



