usrbinenv python codingutf8 2021 the huggingface inc team 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 orglicenseslicense2 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 pretraining a wav2vec2 model on unlabeled audio data import argparse import math import os from dataclasses import dataclass from pathlib import path from typing import dict list optional union import datasets import torch from accelerate import accelerator from accelerate logging import getlogger from datasets import datasetdict concatenatedatasets loaddataset from huggingfacehub import repository createrepo from torch utils data dataloader import dataloader from tqdm auto import tqdm import transformers from transformers import adamw schedulertype wav2vec2config wav2vec2featureextractor wav2vec2forpretraining getscheduler iswandbavailable setseed from transformers models wav2vec2 modelingwav2vec2 import computemaskindices samplenegativeindices from transformers utils import sendexampletelemetry logger getloggername def parseargs parser argparse argumentparserdescriptionfinetune a transformers model on a text classification task parser addargument datasetname typestr defaultnone helpthe name of the dataset to use via the datasets library parser addargument datasetconfignames nargs typestr requiredtrue helpthe configuration names of the dataset to use via the datasets library parser addargument datasetsplitnames nargs typestr requiredtrue helpthe names of the training data set splits to use via the datasets library parser addargument preprocessingnumworkers typeint defaultnone helpthe number of processes to use for the preprocessing parser addargument overwritecache actionstoretrue helpoverwrite the cached training and evaluation sets parser addargument preprocessingonly actionstoretrue helponly run the preprocessing script to be cached for future use parser addargument cachedir typestr defaultnone helpwhere do you want to store the pretrained models downloaded from huggingface co parser addargument validationsplitpercentage typeint default1 helppercentage of training data that should be used for validation if no validation is present in dataset parser addargument loggingsteps typeint default500 helpnumber of steps between each logging parser addargument savingsteps typeint default500 helpnumber of steps between each logging parser addargument audiocolumnname typestr defaultaudio helpcolumn in the dataset that contains speech file path defaults to audio parser addargument modelnameorpath typestr helppath to pretrained model or model identifier from huggingface comodels requiredtrue parser addargument configname typestr defaultnone helppretrained config name or path if not the same as modelname parser addargument traincachefilename typestr defaultnone helppath to the train cached file name parser addargument validationcachefilename typestr defaultnone helppath to the validation cached file name parser addargument perdevicetrainbatchsize typeint default8 helpbatch size per device for the training dataloader parser addargument perdeviceevalbatchsize typeint default8 helpbatch size per device for the evaluation dataloader parser addargument learningrate typefloat default5e5 helpinitial learning rate after the potential warmup period to use parser addargumentweightdecay typefloat default0 0 helpweight decay to use parser addargumentnumtrainepochs typeint default3 helptotal number of training epochs to perform parser addargument maxtrainsteps typeint defaultnone helptotal number of training steps to perform if provided overrides numtrainepochs parser addargument gradientaccumulationsteps typeint default1 helpnumber of updates steps to accumulate before performing a backwardupdate pass parser addargument gradientcheckpointing actionstoretrue helpif true use gradient checkpointing to save memory at the expense of slower backward pass parser addargument lrschedulertype typeschedulertype defaultlinear helpthe scheduler type to use choiceslinear cosine cosinewithrestarts polynomial constant constantwithwarmup parser addargument numwarmupsteps typeint default0 helpnumber of steps for the warmup in the lr scheduler parser addargumentoutputdir typestr defaultnone helpwhere to store the final model parser addargumentseed typeint default0 helpa seed for reproducible training parser addargument maxgumbeltemperature typefloat default2 0 helpmaximum temperature for gumbel softmax parser addargument mingumbeltemperature typefloat default0 5 helpminimum temperature for gumbel softmax parser addargument gumbeltemperaturedecay typefloat default0 999995 helpdecay of gumbel temperature during training parser addargument maxdurationinseconds typefloat default5 0 helpfilter out audio files that are longer than maxdurationinseconds seconds parser addargument mindurationinseconds typefloat default3 0 helpfilter out audio files that are shorter than mindurationinseconds seconds parser addargument padtomultipleof typeint defaultnone help if set will pad the sequence to a multiple of the provided value this is especially useful to enable the use of tensor cores on nvidia hardware with compute capability 7 5 volta parser addargument adambeta1 typefloat default0 9 helpbeta1 for adamw optimizer parser addargument adambeta2 typefloat default0 999 helpbeta2 for adamw optimizer parser addargument adamepsilon typefloat default1e8 helpepsilon for adamw optimizer parser addargumentpushtohub actionstoretrue helpwhether or not to push the model to the hub parser addargument hubmodelid typestr helpthe name of the repository to keep in sync with the local outputdir parser addargumenthubtoken typestr helpthe token to use to push to the model hub parser addargument masktimeprob typefloat defaultnone help percentage between 0 and 1 of all feature vectors along the time axis which will be masked in the contrastive task if omitted will pull value from model config parser addargument masktimelength typeint defaultnone help length of each vector mask span to mask along the time axis in the contrastive task if omitted will pull value from model config args parser parseargs if args pushtohub assert args outputdir is not none need an outputdir to create a repo when pushtohub is passed if args outputdir is not none os makedirsargs outputdir existoktrue return args dataclass class datacollatorforwav2vec2pretraining model wav2vec2forpretraining featureextractor wav2vec2featureextractor padding unionbool str longest padtomultipleof optionalint none masktimeprob optionalfloat 0 65 masktimelength optionalint 10 def callself features listdictstr unionlistint torch tensor dictstr torch tensor reformat list to dict and set to pytorch format batch self featureextractor pad features paddingself padding padtomultipleofself padtomultipleof returntensorspt device batchinputvalues device batchsize batchinputvalues shape0 maskindicesseqlength self model getfeatextractoutputlengthsbatchinputvalues shape1 make sure masked sequence length is a python scalar maskindicesseqlength intmaskindicesseqlength make sure that no loss is computed on padded inputs if batch getattentionmask is not none compute real output lengths according to convolution formula batchsubattentionmask self model getfeaturevectorattentionmask maskindicesseqlength batchattentionmask featuresshape batchsize maskindicesseqlength sample randomly masked indices masktimeindices computemaskindices featuresshape self masktimeprob self masktimelength attentionmaskbatch getsubattentionmask sample negative indices samplednegativeindices samplenegativeindices featuresshape self model config numnegatives masktimeindicesmasktimeindices batchmasktimeindices torch tensormasktimeindices dtypetorch long devicedevice batchsamplednegativeindices torch tensorsamplednegativeindices dtypetorch long devicedevice return batch def multiplygradsparams c compute grad norm given a gradient scale totalnorm 0 0 for p in params if p grad is not none paramnorm p grad detach data scale norm2 totalnorm paramnorm item 2 totalnorm totalnorm0 5 return totalnorm def main see all possible arguments in srctransformersargs py or by passing the help flag to this script we now keep distinct sets of args for a cleaner separation of concerns args parseargs sending telemetry tracking the example usage helps us better allocate resources to maintain them the information sent is the one passed as arguments along with your pythonpytorch versions sendexampletelemetryrunwav2vec2pretrainingnotrainer args initialize the accelerator we will let the accelerator handle device placement for us in this example accelerator accelerator logger infoaccelerator state mainprocessonlyfalse if accelerator islocalmainprocess datasets utils logging setverbositywarning transformers utils logging setverbosityinfo set up weights and biases if available if iswandbavailable import wandb wandb initprojectargs outputdir split1 else datasets utils logging setverbosityerror transformers utils logging setverbosityerror if passed along set the training seed now if args seed is not none setseedargs seed handle the repository creation if accelerator ismainprocess if args pushtohub and not args preprocessingonly retrieve of infer reponame reponame args hubmodelid if reponame is none reponame pathargs outputdir absolute name create repo and retrieve repoid repoid createreporeponame existoktrue tokenargs hubtoken repoid clone repo locally repo repositoryargs outputdir clonefromrepoid tokenargs hubtoken elif args outputdir is not none os makedirsargs outputdir existoktrue accelerator waitforeveryone 1 download and create train validation dataset we load all dataset configuration and datset split pairs passed in args datasetconfignames and args datasetsplitnames datasetssplits for datasetconfigname trainsplitname in zipargs datasetconfignames args datasetsplitnames load dataset datasetsplit loaddataset args datasetname datasetconfigname splittrainsplitname cachedirargs cachedir datasetssplits appenddatasetsplit next we concatenate all configurations and splits into a single training dataset rawdatasets datasetdict if lendatasetssplits 1 rawdatasetstrain concatenatedatasetsdatasetssplits shuffleseedargs seed else rawdatasetstrain datasetssplits0 take args validationsplitpercentage from the training dataset for the validationsplitpercentage numvalidationsamples rawdatasetstrain numrows args validationsplitpercentage 100 if numvalidationsamples 0 raise valueerror args validationsplitpercentage is less than a single sample ffor lenrawdatasets train training samples increase args numvalidationsplitpercentage rawdatasetsvalidation rawdatasetstrain selectrangenumvalidationsamples rawdatasetstrain rawdatasetstrain selectrangenumvalidationsamples rawdatasetstrain numrows 2 now we preprocess the datasets including loading the audio resampling and normalization thankfully datasets takes care of automatically loading and resampling the audio so that we just need to set the correct target sampling rate and normalize the input via the featureextractor featureextractor wav2vec2featureextractor frompretrainedargs modelnameorpath make sure that dataset decodes audio with correct sampling rate rawdatasets rawdatasets castcolumn args audiocolumnname datasets features audiosamplingratefeatureextractor samplingrate only normalizedinputstraining is supported if not featureextractor donormalize raise valueerror training is only supported for normalized inputs make sure featureextractor donormalize true set max min audio length in number of samples maxlength intargs maxdurationinseconds featureextractor samplingrate minlength intargs mindurationinseconds featureextractor samplingrate def preparedatasetbatch sample batchargs audiocolumnname inputs featureextractor samplearray samplingratesamplesamplingrate maxlengthmaxlength truncationtrue batchinputvalues inputs inputvalues0 batchinputlength leninputs inputvalues0 return batch load via mapped files via path cachefilenames none if args traincachefilename is not none cachefilenames train args traincachefilename validation args validationcachefilename load audio files into numpy arrays with accelerator mainprocessfirst vectorizeddatasets rawdatasets map preparedataset numprocargs preprocessingnumworkers removecolumnsrawdatasetstrain columnnames cachefilenamescachefilenames if minlength 0 0 vectorizeddatasets vectorizeddatasets filter lambda x x minlength numprocargs preprocessingnumworkers inputcolumnsinputlength vectorizeddatasets vectorizeddatasets removecolumnsinputlength for large datasets it is advised to run the preprocessing on a single machine first with args preprocessingonly since there will mostly likely be a timeout when running the script in distributed mode in a second step args preprocessingonly can then be set to false to load the cached dataset if args preprocessingonly return 3 load model config wav2vec2config frompretrainedargs modelnameorpath pretraining is only supported for newer stable layer norm architecture applyspecaugment has to be true maskfeatureprob has to be 0 0 if not config dostablelayernorm or config featextractnorm layer raise valueerror pretraining is only supported for config dostablelayernormtrue and config featextractnorm layer initialize random model model wav2vec2forpretrainingconfig activate gradient checkpointing if needed if args gradientcheckpointing model gradientcheckpointingenable 4 define data collator optimizer and scheduler masktimeprob config masktimeprob if args masktimeprob is none else args masktimeprob masktimelength config masktimelength if args masktimelength is none else args masktimelength datacollator datacollatorforwav2vec2pretraining modelmodel featureextractorfeatureextractor padtomultipleofargs padtomultipleof masktimeprobmasktimeprob masktimelengthmasktimelength traindataloader dataloader vectorizeddatasetstrain shuffletrue collatefndatacollator batchsizeargs perdevicetrainbatchsize evaldataloader dataloader vectorizeddatasetsvalidation collatefndatacollator batchsizeargs perdeviceevalbatchsize optimizer optimizer adamw listmodel parameters lrargs learningrate betasargs adambeta1 args adambeta2 epsargs adamepsilon prepare everything with our accelerator model optimizer traindataloader evaldataloader accelerator prepare model optimizer traindataloader evaldataloader scheduler and math around the number of training steps numupdatestepsperepoch math ceillentraindataloader args gradientaccumulationsteps if args maxtrainsteps is none args maxtrainsteps args numtrainepochs numupdatestepsperepoch lrscheduler getscheduler nameargs lrschedulertype optimizeroptimizer numwarmupstepsargs numwarmupsteps numtrainingstepsargs maxtrainsteps afterwards we recalculate our number of training epochs args numtrainepochs math ceilargs maxtrainsteps numupdatestepsperepoch 5 train totalbatchsize args perdevicetrainbatchsize accelerator numprocesses args gradientaccumulationsteps logger info running training logger infof num examples lenvectorizeddatasets train logger infof num epochs args numtrainepochs logger infof instantaneous batch size per device args perdevicetrainbatchsize logger infof total train batch size w parallel distributed accumulation totalbatchsize logger infof gradient accumulation steps args gradientaccumulationsteps logger infof total optimization steps args maxtrainsteps completedsteps 0 startingepoch 0 only show the progress bar once on each machine progressbar tqdmrangeargs maxtrainsteps disablenot accelerator islocalmainprocess completedsteps 0 startingepoch 0 for epoch in rangestartingepoch args numtrainepochs model train for step batch in enumeratetraindataloader compute num of losses numlosses batchmasktimeindices sum subattentionmask batch popsubattentionmask none subattentionmask subattentionmask if subattentionmask is not none else torch oneslikebatchmasktimeindices percentmasked numlosses subattentionmask sum forward outputs modelbatch divide loss by gradient accumulation steps since gradients are accumulated for multiple backward passes in pytorch loss outputs loss args gradientaccumulationsteps accelerator backwardloss make sure that numlosses is summed for distributed training and average gradients over losses of all devices if accelerator state numprocesses 1 numlosses accelerator gatherformetricsnumlosses sum gradientmultiplier accelerator state numprocesses numlosses multiplygradsmodel module parameters gradientmultiplier else multiplygradsmodel parameters 1 numlosses update step if step 1 args gradientaccumulationsteps 0 or step lentraindataloader 1 compute grad norm for monitoring scale accelerator scaler scale item if hasattraccelerator scaler and accelerator scaler is not none else 1 if accelerator state numprocesses 1 gradnorm getgradnormmodel module parameters scale else gradnorm getgradnormmodel parameters scale update parameters optimizer step optimizer zerograd if not accelerator optimizerstepwasskipped lrscheduler step elif accelerator islocalmainprocess progressbar write fgradients have overflown skipping update step updating gradient scale to scale update gumbel temperature gumbeltemperature max args maxgumbeltemperature args gumbeltemperaturedecaycompletedsteps args mingumbeltemperature if hasattrmodel module model module setgumbeltemperaturegumbeltemperature else model setgumbeltemperaturegumbeltemperature progressbar update1 completedsteps 1 6 log all results if step 1 args gradientaccumulationsteps args loggingsteps 0 loss detach outputs contrastiveloss detach outputs diversityloss detach if accelerator state numprocesses 1 loss accelerator gatherformetricsloss sum outputs contrastiveloss accelerator gatherformetricsoutputs contrastiveloss sum outputs diversityloss accelerator gatherformetricsoutputs diversityloss sum percentmasked accelerator gatherformetricspercentmasked sum trainlogs loss loss args gradientaccumulationsteps numlosses constrastloss outputs contrastiveloss numlosses divloss outputs diversityloss numlosses maskidx percentmasked accelerator numprocesses ppl outputs codevectorperplexity lr torch tensoroptimizer paramgroups0lr temp torch tensorgumbeltemperature gradnorm torch tensorgradnorm logstr for k v in trainlogs items logstr 3e formatk v item if accelerator islocalmainprocess progressbar writelogstr if iswandbavailable wandb logtrainlogs save model every args savingsteps steps if step 1 args gradientaccumulationsteps args savingsteps 0 if args pushtohub and epoch args numtrainepochs 1 or args outputdir is not none accelerator waitforeveryone unwrappedmodel accelerator unwrapmodelmodel unwrappedmodel savepretrained args outputdir ismainprocessaccelerator ismainprocess savefunctionaccelerator save if args pushtohub and epoch args numtrainepochs 1 and accelerator ismainprocess repo pushtohub commitmessageftraining in progress step completedsteps blockingfalse autolfsprunetrue if completed steps args maxtrainsteps stop if completedsteps args maxtrainsteps break 7 validate model eval init logs vallogs valloss 0 valcontrastiveloss 0 valdiversityloss 0 valnumlosses 0 for step batch in enumerateevaldataloader with torch nograd batch popsubattentionmask none outputs modelbatch vallogsvalloss outputs loss vallogsvalcontrastiveloss outputs contrastiveloss vallogsvaldiversityloss outputs diversityloss vallogsvalnumlosses batchmasktimeindices sum sum over devices in multiprocessing if accelerator numprocesses 1 vallogs k accelerator gatherformetricsv sum for k v in vallogs items vallogs k v vallogsvalnumlosses for k v in vallogs items logstr for k v in vallogs items logstr 3e formatk v item if accelerator islocalmainprocess progressbar writelogstr if iswandbavailable wandb logvallogs if args outputdir is not none accelerator waitforeveryone unwrappedmodel accelerator unwrapmodelmodel unwrappedmodel savepretrained args outputdir ismainprocessaccelerator ismainprocess savefunctionaccelerator save if accelerator ismainprocess if args pushtohub repo pushtohubcommitmessageend of training autolfsprunetrue if name main main usr bin env python coding utf 8 2021 the huggingface inc team 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 pre training a wav2vec2 model on unlabeled audio data data collator that will dynamically pad the inputs received and prepare masked indices for self supervised pretraining args model class transformers wav2vec2forpretraining the wav2vec2 model used for pretraining the data collator needs to have access to config and _get_feat_extract_output_lengths function for correct padding feature_extractor class transformers wav2vec2featureextractor the processor used for proccessing the data padding obj bool obj str or class transformers tokenization_utils_base paddingstrategy optional defaults to obj true select a strategy to pad the returned sequences according to the model s padding side and padding index among obj true or obj longest pad to the longest sequence in the batch or no padding if only a single sequence if provided obj max_length pad to a maximum length specified with the argument obj max_length or to the maximum acceptable input length for the model if that argument is not provided obj false or obj do_not_pad default no padding i e can output a batch with sequences of different lengths max_length obj int optional maximum length of the input_values of the returned list and optionally padding length see above pad_to_multiple_of obj int optional if set will pad the sequence to a multiple of the provided value this is especially useful to enable the use of tensor cores on nvidia hardware with compute capability 7 5 volta mask_time_prob obj float optional defaults to obj 0 65 percentage between 0 and 1 of all feature vectors along the time axis which will be masked for the contrastive task note that overlap between masked sequences may decrease the actual percentage of masked vectors the default value is taken from the original wav2vec 2 0 article https arxiv org abs 2006 11477 and results in about 49 percent of each sequence being masked on average mask_time_length obj int optional defaults to obj 10 length of each vector mask span to mask along the time axis in the contrastive task the default value originates from the original wav2vec 2 0 article and corresponds to the m variable mentioned there reformat list to dict and set to pytorch format make sure masked sequence length is a python scalar make sure that no loss is computed on padded inputs compute real output lengths according to convolution formula sample randomly masked indices sample negative indices multiplies grads by a constant c compute grad norm given a gradient scale see all possible arguments in src transformers args py or by passing the help flag to this script we now keep distinct sets of args for a cleaner separation of concerns sending telemetry tracking the example usage helps us better allocate resources to maintain them the information sent is the one passed as arguments along with your python pytorch versions initialize the accelerator we will let the accelerator handle device placement for us in this example set up weights and biases if available if passed along set the training seed now handle the repository creation retrieve of infer repo_name create repo and retrieve repo_id clone repo locally 1 download and create train validation dataset we load all dataset configuration and datset split pairs passed in args dataset_config_names and args dataset_split_names load dataset next we concatenate all configurations and splits into a single training dataset take args validation_split_percentage from the training dataset for the validation_split_percentage 2 now we preprocess the datasets including loading the audio resampling and normalization thankfully datasets takes care of automatically loading and resampling the audio so that we just need to set the correct target sampling rate and normalize the input via the feature_extractor make sure that dataset decodes audio with correct sampling rate only normalized inputs training is supported set max min audio length in number of samples load via mapped files via path load audio files into numpy arrays for large datasets it is advised to run the preprocessing on a single machine first with args preprocessing_only since there will mostly likely be a timeout when running the script in distributed mode in a second step args preprocessing_only can then be set to false to load the cached dataset 3 load model 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 initialize random model activate gradient checkpointing if needed 4 define data collator optimizer and scheduler optimizer prepare everything with our accelerator scheduler and math around the number of training steps afterwards we recalculate our number of training epochs 5 train only show the progress bar once on each machine compute num of losses forward divide loss by gradient accumulation steps since gradients are accumulated for multiple backward passes in pytorch make sure that num_losses is summed for distributed training and average gradients over losses of all devices update step compute grad norm for monitoring update parameters update gumbel temperature 6 log all results save model every args saving_steps steps if completed steps args max_train_steps stop 7 validate init logs sum over devices in multi processing
import argparse
import math
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional, Union

import datasets
import torch
from accelerate import Accelerator
from accelerate.logging import get_logger
from datasets import DatasetDict, concatenate_datasets, load_dataset
from huggingface_hub import Repository, create_repo
from torch.utils.data.dataloader import DataLoader
from tqdm.auto import tqdm

import transformers
from transformers import (
    AdamW,
    SchedulerType,
    Wav2Vec2Config,
    Wav2Vec2FeatureExtractor,
    Wav2Vec2ForPreTraining,
    get_scheduler,
    is_wandb_available,
    set_seed,
)
from transformers.models.wav2vec2.modeling_wav2vec2 import _compute_mask_indices, _sample_negative_indices
from transformers.utils import send_example_telemetry


logger = get_logger(__name__)


def parse_args():
    parser = argparse.ArgumentParser(description="Finetune a transformers model on a text classification task")
    parser.add_argument(
        "--dataset_name",
        type=str,
        default=None,
        help="The name of the dataset to use (via the datasets library).",
    )
    parser.add_argument(
        "--dataset_config_names",
        nargs="+",
        type=str,
        required=True,
        help="The configuration names of the dataset to use (via the datasets library).",
    )
    parser.add_argument(
        "--dataset_split_names",
        nargs="+",
        type=str,
        required=True,
        help="The names of the training data set splits to use (via the datasets library).",
    )
    parser.add_argument(
        "--preprocessing_num_workers",
        type=int,
        default=None,
        help="The number of processes to use for the preprocessing.",
    )
    parser.add_argument(
        "--overwrite_cache", action="store_true", help="Overwrite the cached training and evaluation sets"
    )
    parser.add_argument(
        "--preprocessing_only",
        action="store_true",
        help="Only run the preprocessing script to be cached for future use",
    )
    parser.add_argument(
        "--cache_dir",
        type=str,
        default=None,
        help="Where do you want to store the pretrained models downloaded from huggingface.co",
    )
    parser.add_argument(
        "--validation_split_percentage",
        type=int,
        default=1,
        help="Percentage of training data that should be used for validation if no validation is present in dataset.",
    )
    parser.add_argument(
        "--logging_steps",
        type=int,
        default=500,
        help="Number of steps between each logging",
    )
    parser.add_argument(
        "--saving_steps",
        type=int,
        default=500,
        help="Number of steps between each logging",
    )
    parser.add_argument(
        "--audio_column_name",
        type=str,
        default="audio",
        help="Column in the dataset that contains speech file path. Defaults to 'audio'",
    )
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        help="Path to pretrained model or model identifier from huggingface.co/models.",
        required=True,
    )
    parser.add_argument(
        "--config_name",
        type=str,
        default=None,
        help="Pretrained config name or path if not the same as model_name",
    )
    parser.add_argument(
        "--train_cache_file_name",
        type=str,
        default=None,
        help="Path to the train cached file name",
    )
    parser.add_argument(
        "--validation_cache_file_name",
        type=str,
        default=None,
        help="Path to the validation cached file name",
    )
    parser.add_argument(
        "--per_device_train_batch_size",
        type=int,
        default=8,
        help="Batch size (per device) for the training dataloader.",
    )
    parser.add_argument(
        "--per_device_eval_batch_size",
        type=int,
        default=8,
        help="Batch size (per device) for the evaluation dataloader.",
    )
    parser.add_argument(
        "--learning_rate",
        type=float,
        default=5e-5,
        help="Initial learning rate (after the potential warmup period) to use.",
    )
    parser.add_argument("--weight_decay", type=float, default=0.0, help="Weight decay to use.")
    parser.add_argument("--num_train_epochs", type=int, default=3, help="Total number of training epochs to perform.")
    parser.add_argument(
        "--max_train_steps",
        type=int,
        default=None,
        help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
    )
    parser.add_argument(
        "--gradient_accumulation_steps",
        type=int,
        default=1,
        help="Number of updates steps to accumulate before performing a backward/update pass.",
    )
    parser.add_argument(
        "--gradient_checkpointing",
        action="store_true",
        help="If True, use gradient checkpointing to save memory at the expense of slower backward pass.",
    )
    parser.add_argument(
        "--lr_scheduler_type",
        type=SchedulerType,
        default="linear",
        help="The scheduler type to use.",
        choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"],
    )
    parser.add_argument(
        "--num_warmup_steps", type=int, default=0, help="Number of steps for the warmup in the lr scheduler."
    )
    parser.add_argument("--output_dir", type=str, default=None, help="Where to store the final model.")
    parser.add_argument("--seed", type=int, default=0, help="A seed for reproducible training.")
    parser.add_argument(
        "--max_gumbel_temperature",
        type=float,
        default=2.0,
        help="Maximum temperature for gumbel softmax.",
    )
    parser.add_argument(
        "--min_gumbel_temperature",
        type=float,
        default=0.5,
        help="Minimum temperature for gumbel softmax.",
    )
    parser.add_argument(
        "--gumbel_temperature_decay", type=float, default=0.999995, help="Decay of gumbel temperature during training."
    )
    parser.add_argument(
        "--max_duration_in_seconds",
        type=float,
        default=5.0,
        help="Filter out audio files that are longer than `max_duration_in_seconds` seconds",
    )
    parser.add_argument(
        "--min_duration_in_seconds",
        type=float,
        default=3.0,
        help="Filter out audio files that are shorter than `min_duration_in_seconds` seconds",
    )
    parser.add_argument(
        "--pad_to_multiple_of",
        type=int,
        default=None,
        help=(
            "If set will pad the sequence to a multiple of the provided value. This is especially useful to enable the"
            " use of Tensor Cores on NVIDIA hardware with compute capability >= 7.5 (Volta)."
        ),
    )
    parser.add_argument(
        "--adam_beta1",
        type=float,
        default=0.9,
        help="Beta1 for AdamW optimizer",
    )
    parser.add_argument(
        "--adam_beta2",
        type=float,
        default=0.999,
        help="Beta2 for AdamW optimizer",
    )
    parser.add_argument(
        "--adam_epsilon",
        type=float,
        default=1e-8,
        help="Epsilon for AdamW optimizer",
    )
    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
    parser.add_argument(
        "--hub_model_id", type=str, help="The name of the repository to keep in sync with the local `output_dir`."
    )
    parser.add_argument("--hub_token", type=str, help="The token to use to push to the Model Hub.")
    parser.add_argument(
        "--mask_time_prob",
        type=float,
        default=None,
        help=(
            "Percentage (between 0 and 1) of all feature vectors along the time axis which will be masked in the"
            " contrastive task. If omitted, will pull value from model config."
        ),
    )
    parser.add_argument(
        "--mask_time_length",
        type=int,
        default=None,
        help=(
            "Length of each vector mask span to mask along the time axis in the contrastive task."
            " If omitted, will pull value from model config."
        ),
    )
    args = parser.parse_args()

    if args.push_to_hub:
        assert args.output_dir is not None, "Need an `output_dir` to create a repo when `--push_to_hub` is passed."

    if args.output_dir is not None:
        os.makedirs(args.output_dir, exist_ok=True)

    return args


@dataclass
class DataCollatorForWav2Vec2Pretraining:
    

    model: Wav2Vec2ForPreTraining
    feature_extractor: Wav2Vec2FeatureExtractor
    padding: Union[bool, str] = "longest"
    pad_to_multiple_of: Optional[int] = None
    mask_time_prob: Optional[float] = 0.65
    mask_time_length: Optional[int] = 10

    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
        
        batch = self.feature_extractor.pad(
            features,
            padding=self.padding,
            pad_to_multiple_of=self.pad_to_multiple_of,
            return_tensors="pt",
        )

        device = batch["input_values"].device
        batch_size = batch["input_values"].shape[0]

        mask_indices_seq_length = self.model._get_feat_extract_output_lengths(batch["input_values"].shape[-1])
        
        mask_indices_seq_length = int(mask_indices_seq_length)

        
        if batch.get("attention_mask") is not None:
            
            batch["sub_attention_mask"] = self.model._get_feature_vector_attention_mask(
                mask_indices_seq_length, batch["attention_mask"]
            )

        features_shape = (batch_size, mask_indices_seq_length)

        
        mask_time_indices = _compute_mask_indices(
            features_shape,
            self.mask_time_prob,
            self.mask_time_length,
            attention_mask=batch.get("sub_attention_mask"),
        )

        
        sampled_negative_indices = _sample_negative_indices(
            features_shape,
            self.model.config.num_negatives,
            mask_time_indices=mask_time_indices,
        )
        batch["mask_time_indices"] = torch.tensor(mask_time_indices, dtype=torch.long, device=device)
        batch["sampled_negative_indices"] = torch.tensor(sampled_negative_indices, dtype=torch.long, device=device)

        return batch


def multiply_grads(params, c):
    
    for p in params:
        if p.grad is not None:
            if torch.is_tensor(c):
                c = c.to(p.grad.device)
            p.grad.data.mul_(c)


def get_grad_norm(params, scale=1):
    
    total_norm = 0.0
    for p in params:
        if p.grad is not None:
            param_norm = (p.grad.detach().data / scale).norm(2)
            total_norm += param_norm.item() ** 2
    total_norm = total_norm**0.5
    return total_norm


def main():
    
    
    
    args = parse_args()

    
    
    send_example_telemetry("run_wav2vec2_pretraining_no_trainer", args)

    
    accelerator = Accelerator()
    logger.info(accelerator.state, main_process_only=False)
    if accelerator.is_local_main_process:
        datasets.utils.logging.set_verbosity_warning()
        transformers.utils.logging.set_verbosity_info()

        
        if is_wandb_available():
            import wandb

            wandb.init(project=args.output_dir.split("/")[-1])
    else:
        datasets.utils.logging.set_verbosity_error()
        transformers.utils.logging.set_verbosity_error()

    
    if args.seed is not None:
        set_seed(args.seed)

    
    if accelerator.is_main_process:
        if args.push_to_hub and not args.preprocessing_only:
            
            repo_name = args.hub_model_id
            if repo_name is None:
                repo_name = Path(args.output_dir).absolute().name
            
            repo_id = create_repo(repo_name, exist_ok=True, token=args.hub_token).repo_id
            
            repo = Repository(args.output_dir, clone_from=repo_id, token=args.hub_token)
        elif args.output_dir is not None:
            os.makedirs(args.output_dir, exist_ok=True)
    accelerator.wait_for_everyone()

    
    
    
    datasets_splits = []
    for dataset_config_name, train_split_name in zip(args.dataset_config_names, args.dataset_split_names):
        
        dataset_split = load_dataset(
            args.dataset_name,
            dataset_config_name,
            split=train_split_name,
            cache_dir=args.cache_dir,
        )
        datasets_splits.append(dataset_split)

    
    raw_datasets = DatasetDict()
    if len(datasets_splits) > 1:
        raw_datasets["train"] = concatenate_datasets(datasets_splits).shuffle(seed=args.seed)
    else:
        raw_datasets["train"] = datasets_splits[0]

    
    num_validation_samples = raw_datasets["train"].num_rows * args.validation_split_percentage // 100

    if num_validation_samples == 0:
        raise ValueError(
            "`args.validation_split_percentage` is less than a single sample "
            f"for {len(raw_datasets['train'])} training samples. Increase "
            "`args.num_validation_split_percentage`. "
        )

    raw_datasets["validation"] = raw_datasets["train"].select(range(num_validation_samples))
    raw_datasets["train"] = raw_datasets["train"].select(range(num_validation_samples, raw_datasets["train"].num_rows))

    
    
    
    
    feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(args.model_name_or_path)

    
    raw_datasets = raw_datasets.cast_column(
        args.audio_column_name, datasets.features.Audio(sampling_rate=feature_extractor.sampling_rate)
    )

    
    if not feature_extractor.do_normalize:
        raise ValueError(
            "Training is only supported for normalized inputs. Make sure ``feature_extractor.do_normalize == True``"
        )

    
    max_length = int(args.max_duration_in_seconds * feature_extractor.sampling_rate)
    min_length = int(args.min_duration_in_seconds * feature_extractor.sampling_rate)

    def prepare_dataset(batch):
        sample = batch[args.audio_column_name]

        inputs = feature_extractor(
            sample["array"], sampling_rate=sample["sampling_rate"], max_length=max_length, truncation=True
        )
        batch["input_values"] = inputs.input_values[0]
        batch["input_length"] = len(inputs.input_values[0])

        return batch

    
    cache_file_names = None
    if args.train_cache_file_name is not None:
        cache_file_names = {"train": args.train_cache_file_name, "validation": args.validation_cache_file_name}

    
    with accelerator.main_process_first():
        vectorized_datasets = raw_datasets.map(
            prepare_dataset,
            num_proc=args.preprocessing_num_workers,
            remove_columns=raw_datasets["train"].column_names,
            cache_file_names=cache_file_names,
        )

        if min_length > 0.0:
            vectorized_datasets = vectorized_datasets.filter(
                lambda x: x > min_length,
                num_proc=args.preprocessing_num_workers,
                input_columns=["input_length"],
            )

        vectorized_datasets = vectorized_datasets.remove_columns("input_length")

    
    
    
    
    
    if args.preprocessing_only:
        return

    
    config = Wav2Vec2Config.from_pretrained(args.model_name_or_path)

    
    
    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'"
        )

    
    model = Wav2Vec2ForPreTraining(config)

    
    if args.gradient_checkpointing:
        model.gradient_checkpointing_enable()

    

    mask_time_prob = config.mask_time_prob if args.mask_time_prob is None else args.mask_time_prob
    mask_time_length = config.mask_time_length if args.mask_time_length is None else args.mask_time_length

    data_collator = DataCollatorForWav2Vec2Pretraining(
        model=model,
        feature_extractor=feature_extractor,
        pad_to_multiple_of=args.pad_to_multiple_of,
        mask_time_prob=mask_time_prob,
        mask_time_length=mask_time_length,
    )
    train_dataloader = DataLoader(
        vectorized_datasets["train"],
        shuffle=True,
        collate_fn=data_collator,
        batch_size=args.per_device_train_batch_size,
    )
    eval_dataloader = DataLoader(
        vectorized_datasets["validation"], collate_fn=data_collator, batch_size=args.per_device_eval_batch_size
    )

    
    optimizer = AdamW(
        list(model.parameters()),
        lr=args.learning_rate,
        betas=[args.adam_beta1, args.adam_beta2],
        eps=args.adam_epsilon,
    )

    
    model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
        model, optimizer, train_dataloader, eval_dataloader
    )

    
    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)

    if args.max_train_steps is None:
        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch

    lr_scheduler = get_scheduler(
        name=args.lr_scheduler_type,
        optimizer=optimizer,
        num_warmup_steps=args.num_warmup_steps,
        num_training_steps=args.max_train_steps,
    )

    
    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)

    
    total_batch_size = args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps

    logger.info("***** Running training *****")
    logger.info(f"  Num examples = {len(vectorized_datasets['train'])}")
    logger.info(f"  Num Epochs = {args.num_train_epochs}")
    logger.info(f"  Instantaneous batch size per device = {args.per_device_train_batch_size}")
    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {args.max_train_steps}")
    completed_steps = 0
    starting_epoch = 0

    
    progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)
    completed_steps = 0
    starting_epoch = 0
    for epoch in range(starting_epoch, args.num_train_epochs):
        model.train()
        for step, batch in enumerate(train_dataloader):
            
            num_losses = batch["mask_time_indices"].sum()
            sub_attention_mask = batch.pop("sub_attention_mask", None)
            sub_attention_mask = (
                sub_attention_mask if sub_attention_mask is not None else torch.ones_like(batch["mask_time_indices"])
            )
            percent_masked = num_losses / sub_attention_mask.sum()

            
            outputs = model(**batch)

            
            
            loss = outputs.loss / args.gradient_accumulation_steps
            accelerator.backward(loss)

            
            
            if accelerator.state.num_processes > 1:
                num_losses = accelerator.gather_for_metrics(num_losses).sum()
                gradient_multiplier = accelerator.state.num_processes / num_losses
                multiply_grads(model.module.parameters(), gradient_multiplier)
            else:
                multiply_grads(model.parameters(), 1 / num_losses)

            
            if (step + 1) % args.gradient_accumulation_steps == 0 or step == len(train_dataloader) - 1:
                
                scale = (
                    accelerator.scaler._scale.item()
                    if hasattr(accelerator, "scaler") and accelerator.scaler is not None
                    else 1
                )
                if accelerator.state.num_processes > 1:
                    grad_norm = get_grad_norm(model.module.parameters(), scale)
                else:
                    grad_norm = get_grad_norm(model.parameters(), scale)

                
                optimizer.step()
                optimizer.zero_grad()

                if not accelerator.optimizer_step_was_skipped:
                    lr_scheduler.step()
                elif accelerator.is_local_main_process:
                    progress_bar.write(
                        f"Gradients have overflown - skipping update step... Updating gradient scale to {scale}..."
                    )

                
                gumbel_temperature = max(
                    args.max_gumbel_temperature * args.gumbel_temperature_decay**completed_steps,
                    args.min_gumbel_temperature,
                )
                if hasattr(model, "module"):
                    model.module.set_gumbel_temperature(gumbel_temperature)
                else:
                    model.set_gumbel_temperature(gumbel_temperature)

                progress_bar.update(1)
                completed_steps += 1

            
            if (step + 1) % (args.gradient_accumulation_steps * args.logging_steps) == 0:
                loss.detach()
                outputs.contrastive_loss.detach()
                outputs.diversity_loss.detach()

                if accelerator.state.num_processes > 1:
                    loss = accelerator.gather_for_metrics(loss).sum()
                    outputs.contrastive_loss = accelerator.gather_for_metrics(outputs.contrastive_loss).sum()
                    outputs.diversity_loss = accelerator.gather_for_metrics(outputs.diversity_loss).sum()
                    percent_masked = accelerator.gather_for_metrics(percent_masked).sum()

                train_logs = {
                    "loss": (loss * args.gradient_accumulation_steps) / num_losses,
                    "constrast_loss": outputs.contrastive_loss / num_losses,
                    "div_loss": outputs.diversity_loss / num_losses,
                    "%_mask_idx": percent_masked / accelerator.num_processes,
                    "ppl": outputs.codevector_perplexity,
                    "lr": torch.tensor(optimizer.param_groups[0]["lr"]),
                    "temp": torch.tensor(gumbel_temperature),
                    "grad_norm": torch.tensor(grad_norm),
                }
                log_str = ""
                for k, v in train_logs.items():
                    log_str += "| {}: {:.3e}".format(k, v.item())

                if accelerator.is_local_main_process:
                    progress_bar.write(log_str)
                    if is_wandb_available():
                        wandb.log(train_logs)

            
            if (step + 1) % (args.gradient_accumulation_steps * args.saving_steps) == 0:
                if (args.push_to_hub and epoch < args.num_train_epochs - 1) or args.output_dir is not None:
                    accelerator.wait_for_everyone()
                    unwrapped_model = accelerator.unwrap_model(model)
                    unwrapped_model.save_pretrained(
                        args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save
                    )

                if (args.push_to_hub and epoch < args.num_train_epochs - 1) and accelerator.is_main_process:
                    repo.push_to_hub(
                        commit_message=f"Training in progress step {completed_steps}",
                        blocking=False,
                        auto_lfs_prune=True,
                    )

            
            if completed_steps >= args.max_train_steps:
                break

        
        model.eval()

        
        val_logs = {
            "val_loss": 0,
            "val_contrastive_loss": 0,
            "val_diversity_loss": 0,
            "val_num_losses": 0,
        }
        for step, batch in enumerate(eval_dataloader):
            with torch.no_grad():
                batch.pop("sub_attention_mask", None)
                outputs = model(**batch)

            val_logs["val_loss"] += outputs.loss
            val_logs["val_contrastive_loss"] += outputs.contrastive_loss
            val_logs["val_diversity_loss"] += outputs.diversity_loss
            val_logs["val_num_losses"] += batch["mask_time_indices"].sum()

        
        if accelerator.num_processes > 1:
            val_logs = {k: accelerator.gather_for_metrics(v).sum() for k, v in val_logs.items()}

        val_logs = {k: v / val_logs["val_num_losses"] for k, v in val_logs.items()}

        log_str = ""
        for k, v in val_logs.items():
            log_str += "| {}: {:.3e}".format(k, v.item())

        if accelerator.is_local_main_process:
            progress_bar.write(log_str)
            if is_wandb_available():
                wandb.log(val_logs)

        if args.output_dir is not None:
            accelerator.wait_for_everyone()
            unwrapped_model = accelerator.unwrap_model(model)
            unwrapped_model.save_pretrained(
                args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save
            )
            if accelerator.is_main_process:
                if args.push_to_hub:
                    repo.push_to_hub(commit_message="End of training", auto_lfs_prune=True)


if __name__ == "__main__":
    main()