





'''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[int]							=				logging.getLogger(__name__)






@dataclass
class UpperCAmelCase							:
    _A :			str						   =       field(
        metadata={"""help""": """Path to pretrained model or model identifier from huggingface.co/models"""}							)
    _A :			Optional[str]						   =       field(
        default=UpperCAmelCase_					,			metadata={"""help""": """Where do you want to store the pretrained models downloaded from huggingface.co"""}					,			)
    _A :			Optional[bool]						   =       field(
        default=UpperCAmelCase_					,			metadata={"""help""": """Whether to freeze the feature extractor layers of the model."""}							)
    _A :			Optional[bool]						   =       field(
        default=UpperCAmelCase_					,			metadata={"""help""": """Whether to log verbose messages or not."""}					,			)
    _A :			Optional[float]						   =       field(
        default=2.0					,			metadata={"""help""": """Maximum temperature for gumbel softmax."""}							)
    _A :			Optional[float]						   =       field(
        default=0.5					,			metadata={"""help""": """Minimum temperature for gumbel softmax."""}							)
    _A :			Optional[float]						   =       field(
        default=0.99_9995					,			metadata={"""help""": """Decay of gumbel temperature during training."""}							)







def _lowerCAmelCase     (      _lowerCAmelCase       ,    _lowerCAmelCase    )->     Optional[Any]:
 logging.basicConfig(
     format='%(asctime)s - %(levelname)s - %(name)s - %(message)s'       ,    datefmt='%m/%d/%Y %H:%M:%S'       ,    handlers=[logging.StreamHandler(sys.stdout    )]       ,    )
 __UpperCAmelCase					=   logging.WARNING
 if model_args.verbose_logging:
  __UpperCAmelCase					=   logging.DEBUG
 elif trainer_utils.is_main_process(training_args.local_rank    ):
  __UpperCAmelCase					=   logging.INFO
 logger.setLevel(_lowerCAmelCase    )






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






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



    def __call__(  self     ,			__A					):
     # reformat list to dict and set to pytorch format
     __UpperCAmelCase					=   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					=   self.model._get_feat_extract_output_lengths(batch['input_values'].shape[-1]					)

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

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

     # sample randomly masked indices
     __UpperCAmelCase					=   _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							(		UpperCAmelCase_							):



    def __init__(  self     ,			*__A     ,			__A=1     ,			__A=0     ,			__A=1.0     ,			**__A					):
     super().__init__(*__A     ,			**__A					)
     __UpperCAmelCase					=   0
     __UpperCAmelCase					=   max_gumbel_temp
     __UpperCAmelCase					=   min_gumbel_temp
     __UpperCAmelCase					=   gumbel_temp_decay




    def 					__lowerCamelCase						(  self     ,			__A     ,			__A					):
     model.train()
     __UpperCAmelCase					=   self._prepare_inputs(__A					)

     if self.use_amp:
      with autocast():
       __UpperCAmelCase					=   self.compute_loss(__A     ,			__A					)
     else:
      __UpperCAmelCase					=   self.compute_loss(__A     ,			__A					)

     if self.args.n_gpu > 1 or self.deepspeed:
      if model.module.config.ctc_loss_reduction == "mean":
       __UpperCAmelCase					=   loss.mean()
      elif model.module.config.ctc_loss_reduction == "sum":
       __UpperCAmelCase					=   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					=   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     (      )->     Dict:
 # 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					=   HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)    )

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

 # Downloading and loading a dataset from the hub.
 __UpperCAmelCase					=   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					=   DatasetDict()
  __UpperCAmelCase					=   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					=   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					=   DatasetDict()
  __UpperCAmelCase					=   load_dataset(
      data_args.dataset_name       ,    data_args.dataset_config_name       ,    split='validation'       ,    cache_dir=model_args.cache_dir       ,    )
  __UpperCAmelCase					=   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					=   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
  __UpperCAmelCase					,    __UpperCAmelCase					=   librosa.load(batch[data_args.speech_file_column]       ,    sr=feature_extractor.sampling_rate    )
  return batch

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

 # filter audio files that are too long
 __UpperCAmelCase					=   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`
 __UpperCAmelCase					=   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
 __UpperCAmelCase					=   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					=   WavaVecaForPreTraining(_lowerCAmelCase    )

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

 __UpperCAmelCase					=   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 2021 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 pathlib import Path

import torch

from ...utils import is_npu_available, is_xpu_available
from .config_args import ClusterConfig, default_json_config_file
from .config_utils import SubcommandHelpFormatter


_A:					str							=				"""Create a default config file for Accelerate with only a few flags set."""







def _lowerCAmelCase     (      _lowerCAmelCase="no"       ,    _lowerCAmelCase = default_json_config_file       ,    _lowerCAmelCase = False    )->     List[Any]:
 __UpperCAmelCase					=   Path(_lowerCAmelCase    )
 path.parent.mkdir(parents=_lowerCAmelCase       ,    exist_ok=_lowerCAmelCase    )
 if path.exists():
  print(
      F'Configuration already exists at {save_location}, will not override. Run `accelerate config` manually or pass a different `save_location`.'    )
  return False
 __UpperCAmelCase					=   mixed_precision.lower()
 if mixed_precision not in ["no", "fp16", "bf16", "fp8"]:
  raise ValueError(
      F'`mixed_precision` should be one of \'no\', \'fp16\', \'bf16\', or \'fp8\'. Received {mixed_precision}'    )
 __UpperCAmelCase					=   {
     'compute_environment': 'LOCAL_MACHINE',
     'mixed_precision': mixed_precision,
 }
 if torch.cuda.is_available():
  __UpperCAmelCase					=   torch.cuda.device_count()
  __UpperCAmelCase					=   num_gpus
  __UpperCAmelCase					=   False
  if num_gpus > 1:
   __UpperCAmelCase					=   'MULTI_GPU'
  else:
   __UpperCAmelCase					=   'NO'
 elif is_xpu_available() and use_xpu:
  __UpperCAmelCase					=   torch.xpu.device_count()
  __UpperCAmelCase					=   num_xpus
  __UpperCAmelCase					=   False
  if num_xpus > 1:
   __UpperCAmelCase					=   'MULTI_XPU'
  else:
   __UpperCAmelCase					=   'NO'
 elif is_npu_available():
  __UpperCAmelCase					=   torch.npu.device_count()
  __UpperCAmelCase					=   num_npus
  __UpperCAmelCase					=   False
  if num_npus > 1:
   __UpperCAmelCase					=   'MULTI_NPU'
  else:
   __UpperCAmelCase					=   'NO'
 else:
  __UpperCAmelCase					=   0
  __UpperCAmelCase					=   True
  __UpperCAmelCase					=   1
  __UpperCAmelCase					=   'NO'
 __UpperCAmelCase					=   ClusterConfig(**_lowerCAmelCase    )
 config.to_json_file(_lowerCAmelCase    )
 return path







def _lowerCAmelCase     (      _lowerCAmelCase       ,    _lowerCAmelCase    )->     List[str]:
 __UpperCAmelCase					=   parser.add_parser('default'       ,    parents=_lowerCAmelCase       ,    help=_lowerCAmelCase       ,    formatter_class=_lowerCAmelCase    )
 parser.add_argument(
     '--config_file'       ,    default=_lowerCAmelCase       ,    help=(
         'The path to use to store the config file. Will default to a file named default_config.yaml in the cache '
         'location, which is the content of the environment `HF_HOME` suffixed with \'accelerate\', or if you don\'t have '
         'such an environment variable, your cache directory (\'~/.cache\' or the content of `XDG_CACHE_HOME`) suffixed '
         'with \'huggingface\'.'
     )       ,    dest='save_location'       ,    )

 parser.add_argument(
     '--mixed_precision'       ,    choices=['no', 'fp16', 'bf16']       ,    type=_lowerCAmelCase       ,    help='Whether or not to use mixed precision training. '
     'Choose between FP16 and BF16 (bfloat16) training. '
     'BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.'       ,    default='no'       ,    )
 parser.set_defaults(func=_lowerCAmelCase    )
 return parser







def _lowerCAmelCase     (      _lowerCAmelCase    )->     Union[str, Any]:
 __UpperCAmelCase					=   write_basic_config(args.mixed_precision       ,    args.save_location    )
 if config_file:
  print(F'accelerate configuration saved at {config_file}'    )





