





'''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 					UpperCAmelCase	:
      _lowercase:     str			=  field(
          metadata={'''help''': '''Path to pretrained model or model identifier from huggingface.co/models'''}  )
      _lowercase:     Optional[str]			=  field(
          default=snake_case_				,						metadata={'''help''': '''Where do you want to store the pretrained models downloaded from huggingface.co'''}				,						)
      _lowercase:     Optional[bool]			=  field(
          default=snake_case_				,						metadata={'''help''': '''Whether to freeze the feature extractor layers of the model.'''}  )
      _lowercase:     Optional[bool]			=  field(
          default=snake_case_				,						metadata={'''help''': '''Whether to log verbose messages or not.'''}				,						)
      _lowercase:     Optional[float]			=  field(
          default=2.0				,						metadata={'''help''': '''Maximum temperature for gumbel softmax.'''}  )
      _lowercase:     Optional[float]			=  field(
          default=0.5				,						metadata={'''help''': '''Minimum temperature for gumbel softmax.'''}  )
      _lowercase:     Optional[float]			=  field(
          default=0.999995				,						metadata={'''help''': '''Decay of gumbel temperature during training.'''}  )




def   UpperCamelCase__ (						lowerCAmelCase				,      lowerCAmelCase    ):





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




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




@dataclass
class 					UpperCAmelCase	:
      _lowercase:     WavaVecaForPreTraining
      _lowercase:     WavaVecaFeatureExtractor
      _lowercase:     Union[bool, str]			=  "longest"
      _lowercase:     Optional[int]			=  None
      _lowercase:     Optional[int]			=  None




      def __call__(     self     :					Union[str, Any]     ,     __snake_case     :					List[Dict[str, Union[List[int], torch.Tensor]]]    )							-> Dict[str, torch.Tensor]:
            # reformat list to dict and set to pytorch format
            _lowerCAmelCase					=				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"""     ,     )
            _lowerCAmelCase					=				self.model._get_feat_extract_output_lengths(batch["""input_values"""].shape[-1]    )

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

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

            # sample randomly masked indices
            _lowerCAmelCase					=				_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 					UpperCAmelCase	(     snake_case_  ):




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




      def 							lowercase__       (     self     :					int     ,     __snake_case     :					nn.Module     ,     __snake_case     :					Dict[str, Union[torch.Tensor, Any]]    )							-> torch.Tensor:
            model.train()
            _lowerCAmelCase					=				self._prepare_inputs(__snake_case    )

            if self.use_amp:
                  with autocast():
                        _lowerCAmelCase					=				self.compute_loss(__snake_case     ,     __snake_case    )
            else:
                  _lowerCAmelCase					=				self.compute_loss(__snake_case     ,     __snake_case    )

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





      """simple docstring"""


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

      _lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase					=				parser.parse_args_into_dataclasses()
      configure_logger(lowerCAmelCase				,      lowerCAmelCase    )

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

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

      # load audio files into numpy arrays
      _lowerCAmelCase					=				datasets.map(
          lowerCAmelCase				,      num_proc=data_args.preprocessing_num_workers				,      remove_columns=datasets["""train"""].column_names    )

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

      def normalize(lowerCAmelCase    ):
            return feature_extractor(batch["""speech"""]				,      sampling_rate=feature_extractor.sampling_rate    )

      # normalize and transform to `BatchFeatures`
      _lowerCAmelCase					=				vectorized_datasets.map(
          lowerCAmelCase				,      batched=lowerCAmelCase				,      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
      _lowerCAmelCase					=				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'"""    )

      _lowerCAmelCase					=				WavaVecaForPreTraining(lowerCAmelCase    )

      _lowerCAmelCase					=				DataCollatorForWavaVecaPretraining(model=lowerCAmelCase				,      feature_extractor=lowerCAmelCase    )

      _lowerCAmelCase					=				WavaVecaPreTrainer(
          model=lowerCAmelCase				,      data_collator=lowerCAmelCase				,      args=lowerCAmelCase				,      train_dataset=vectorized_datasets["""train"""]				,      eval_dataset=vectorized_datasets["""validation"""]				,      tokenizer=lowerCAmelCase				,      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'''






# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


A__							:  int    ={
    '''configuration_xmod''': [
        '''XMOD_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''XmodConfig''',
        '''XmodOnnxConfig''',
    ],
}

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      A__							:  List[str]    =[
          '''XMOD_PRETRAINED_MODEL_ARCHIVE_LIST''',
          '''XmodForCausalLM''',
          '''XmodForMaskedLM''',
          '''XmodForMultipleChoice''',
          '''XmodForQuestionAnswering''',
          '''XmodForSequenceClassification''',
          '''XmodForTokenClassification''',
          '''XmodModel''',
          '''XmodPreTrainedModel''',
      ]

if TYPE_CHECKING:
      from .configuration_xmod import XMOD_PRETRAINED_CONFIG_ARCHIVE_MAP, XmodConfig, XmodOnnxConfig

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_xmod import (
                XMOD_PRETRAINED_MODEL_ARCHIVE_LIST,
                XmodForCausalLM,
                XmodForMaskedLM,
                XmodForMultipleChoice,
                XmodForQuestionAnswering,
                XmodForSequenceClassification,
                XmodForTokenClassification,
                XmodModel,
                XmodPreTrainedModel,
            )

else:
      import sys

      A__							:  str    =_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)






