






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__       		=  True
     from torch.cuda.amp import autocast


lowerCamelCase__       		=  logging.getLogger(__name__)




@dataclass
class       snake_case__ :

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






def 	UpperCamelCase	(							snake_case__  : ModelArguments						,snake_case__  : TrainingArguments	):




     '''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	)]						,)
     __snake_case      :Tuple 							=    logging.WARNING
     if model_args.verbose_logging:
          __snake_case      :Optional[int] 							=    logging.DEBUG
     elif trainer_utils.is_main_process(training_args.local_rank	):
          __snake_case      :Dict 							=    logging.INFO
     logger.setLevel(snake_case__	)




@dataclass
class       snake_case__ :

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




@dataclass
class       snake_case__ :

       '''simple docstring'''
       lowerCamelCase				:					WavaVecaForPreTraining
       lowerCamelCase				:					WavaVecaFeatureExtractor
       lowerCamelCase				:					Union[bool, str]								=			"longest"
       lowerCamelCase				:					Optional[int]								=			None
       lowerCamelCase				:					Optional[int]								=			None



       def __call__(				self      ,    a__     )			->   Dict[str, torch.Tensor]:






            '''simple docstring'''





            __snake_case      :Tuple 							=    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"""      ,    )
            __snake_case      :str 							=    self.model._get_feat_extract_output_lengths(batch["""input_values"""].shape[-1]     )

            __snake_case      :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
                 __snake_case      :Union[str, Any] 							=    self.model._get_feat_extract_output_lengths(batch["""attention_mask"""].sum(-1     )     ).to(
                     torch.long     )

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

            # sample randomly masked indices
            __snake_case      :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       snake_case__ (				lowercase_):

       '''simple docstring'''



       def __init__(				self      ,    *a__      ,    a__=1      ,    a__=0      ,    a__=1.0      ,    **a__     )			->   int:






            '''simple docstring'''





            super().__init__(*a__      ,    **a__     )
            __snake_case      :List[str] 							=    0
            __snake_case      :Any 							=    max_gumbel_temp
            __snake_case      :Optional[int] 							=    min_gumbel_temp
            __snake_case      :Union[str, Any] 							=    gumbel_temp_decay



       def        __lowercase    (				self      ,    a__      ,    a__     )			->   torch.Tensor:






            '''simple docstring'''





            model.train()
            __snake_case      :str 							=    self._prepare_inputs(a__     )

            if self.use_amp:
                 with autocast():
                      __snake_case      :str 							=    self.compute_loss(a__      ,    a__     )
            else:
                 __snake_case      :Optional[Any] 							=    self.compute_loss(a__      ,    a__     )

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




     '''simple docstring'''






     __snake_case      :Any 							=    HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)	)

     __snake_case							,			__snake_case							,			__snake_case      :str 							=    parser.parse_args_into_dataclasses()
     configure_logger(snake_case__						,snake_case__	)

     # Downloading and loading a dataset from the hub.
     __snake_case      :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"
          __snake_case      :Tuple 							=    DatasetDict()
          __snake_case      :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						,)
          __snake_case      :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						,)
     else:
          # make sure only "validation" and "train" keys remain"
          __snake_case      :Tuple 							=    DatasetDict()
          __snake_case      :Tuple 							=    load_dataset(
              data_args.dataset_name						,data_args.dataset_config_name						,split="""validation"""						,cache_dir=model_args.cache_dir						,)
          __snake_case      :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
     __snake_case      :Optional[int] 							=    WavaVecaFeatureExtractor.from_pretrained(
         model_args.model_name_or_path						,cache_dir=model_args.cache_dir						,do_normalize=snake_case__	)

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

     # load audio files into numpy arrays
     __snake_case      :List[Any] 							=    datasets.map(
         snake_case__						,num_proc=data_args.preprocessing_num_workers						,remove_columns=datasets["""train"""].column_names	)

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

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

     # normalize and transform to `BatchFeatures`
     __snake_case      :List[str] 							=    vectorized_datasets.map(
         snake_case__						,batched=snake_case__						,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
     __snake_case      :Optional[int] 							=    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'"""	)

     __snake_case      :List[str] 							=    WavaVecaForPreTraining(snake_case__	)

     __snake_case      :Optional[Any] 							=    DataCollatorForWavaVecaPretraining(model=snake_case__						,feature_extractor=snake_case__	)

     __snake_case      :List[Any] 							=    WavaVecaPreTrainer(
         model=snake_case__						,data_collator=snake_case__						,args=snake_case__						,train_dataset=vectorized_datasets["""train"""]						,eval_dataset=vectorized_datasets["""validation"""]						,tokenizer=snake_case__						,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()










import json
import os
import tempfile

import transformers

import datasets
from utils import generate_example_dataset, get_duration


lowerCamelCase__       		=  50_0000

lowerCamelCase__		,      lowerCamelCase__       		=  os.path.split(__file__)
lowerCamelCase__       		=  os.path.join(RESULTS_BASEPATH, """results""", RESULTS_FILENAME.replace(""".py""", """.json"""))






@get_duration
def 	UpperCamelCase	(							snake_case__  : datasets.Dataset						,**snake_case__  : Any	):




     '''simple docstring'''






     __snake_case      :List[str] 							=    dataset.map(**snake_case__	)






@get_duration
def 	UpperCamelCase	(							snake_case__  : datasets.Dataset						,**snake_case__  : Dict	):




     '''simple docstring'''






     __snake_case      :int 							=    dataset.filter(**snake_case__	)






def 	UpperCamelCase	(							):




     '''simple docstring'''






     __snake_case      :str 							=    {"""num examples""": SPEED_TEST_N_EXAMPLES}
     with tempfile.TemporaryDirectory() as tmp_dir:
          __snake_case      :Dict 							=    datasets.Features({"""text""": datasets.Value("""string"""	), """numbers""": datasets.Value("""float32"""	)}	)
          __snake_case      :Optional[int] 							=    generate_example_dataset(
              os.path.join(snake_case__						,"""dataset.arrow"""	)						,snake_case__						,num_examples=snake_case__	)

          __snake_case      :str 							=    transformers.AutoTokenizer.from_pretrained("""bert-base-cased"""						,use_fast=snake_case__	)

          def tokenize(snake_case__  : Optional[int]	):
               return tokenizer(examples["""text"""]	)

          __snake_case      :Optional[Any] 							=    map(snake_case__	)

          __snake_case      :Optional[Any] 							=    map(snake_case__						,batched=snake_case__	)

          __snake_case      :Optional[Any] 							=    map(snake_case__						,function=lambda snake_case__	: None						,batched=snake_case__	)

          with dataset.formatted_as(type="""numpy"""	):
               __snake_case      :int 							=    map(snake_case__						,function=lambda snake_case__	: None						,batched=snake_case__	)

          with dataset.formatted_as(type="""pandas"""	):
               __snake_case      :List[Any] 							=    map(snake_case__						,function=lambda snake_case__	: None						,batched=snake_case__	)

          with dataset.formatted_as(type="""torch"""						,columns="""numbers"""	):
               __snake_case      :str 							=    map(snake_case__						,function=lambda snake_case__	: None						,batched=snake_case__	)

          with dataset.formatted_as(type="""tensorflow"""						,columns="""numbers"""	):
               __snake_case      :Optional[int] 							=    map(snake_case__						,function=lambda snake_case__	: None						,batched=snake_case__	)

          __snake_case      :Dict 							=    map(snake_case__						,function=snake_case__						,batched=snake_case__	)

          __snake_case      :Tuple 							=    filter(snake_case__	)

          # Activate later when tokenizer support batched inputs
          # with dataset.formatted_as(type='numpy'):
          #     times[func.__name__ + " fast-tokenizer batched numpy"] = func(dataset, function=tokenize, batched=True)

     with open(snake_case__						,"""wb"""	) as f:
          f.write(json.dumps(snake_case__	).encode("""utf-8"""	)	)


if __name__ == "__main__":  # useful to run the profiler
     benchmark_map_filter()


