Tremendous-Tune Wav2Vec2 for English ASR in Hugging Face with 🤗 Transformers

Wav2Vec2 is a pretrained model for Automatic Speech Recognition (ASR)
and was released in September
2020
by Alexei Baevski, Michael Auli, and Alex Conneau.

Using a novel contrastive pretraining objective, Wav2Vec2 learns
powerful speech representations from greater than 50.000 hours of unlabeled
speech. Similar, to BERT’s masked language
modeling, the model learns
contextualized speech representations by randomly masking feature
vectors before passing them to a transformer network.

For the primary time, it has been shown that pretraining, followed by
fine-tuning on little or no labeled speech data achieves competitive
results to state-of-the-art ASR systems. Using as little as 10 minutes
of labeled data, Wav2Vec2 yields a word error rate (WER) of lower than 5%
on the clean test set of
LibriSpeech – cf.
with Table 9 of the paper.

On this notebook, we’ll give an in-detail explanation of how
Wav2Vec2’s pretrained checkpoints may be fine-tuned on any English ASR
dataset. Note that on this notebook, we’ll fine-tune Wav2Vec2 without
making use of a language model. It is far simpler to make use of Wav2Vec2
and not using a language model as an end-to-end ASR system and it has been
shown that a standalone Wav2Vec2 acoustic model achieves impressive
results. For demonstration purposes, we fine-tune the “base”-sized
pretrained checkpoint
on the quite small Timit
dataset that comprises just 5h of coaching data.

Wav2Vec2 is fine-tuned using Connectionist Temporal Classification
(CTC), which is an algorithm that’s used to coach neural networks for
sequence-to-sequence problems and mainly in Automatic Speech Recognition
and handwriting recognition.

I highly recommend reading the blog post Sequence Modeling with CTC
(2017) very well-written blog post by
Awni Hannun.

Before we start, let’s install each datasets and transformers from
master. Also, we’d like the soundfile package to load audio files and
the jiwer to guage our fine-tuned model using the word error rate
(WER) metric ${}^1\{\}^1$

!pip install datasets>=1.18.3
!pip install transformers==4.11.3
!pip install librosa
!pip install jiwer

Next we strongly suggest to upload your training checkpoints on to the Hugging Face Hub while training. The Hub has integrated version control so you’ll be able to ensure that no model checkpoint is getting lost during training.

To accomplish that you’ve to store your authentication token from the Hugging Face website (join here if you happen to have not already!)

from huggingface_hub import notebook_login

notebook_login()

Print Output:

Login successful
Your token has been saved to /root/.huggingface/token
Authenticated through git-crendential store but this isn't the helper defined in your machine.
You'll have to re-authenticate when pushing to the Hugging Face Hub. Run the next command in your terminal to set it because the default

git config --global credential.helper store

Then it is advisable install Git-LFS to upload your model checkpoints:

!apt install git-lfs

${}^1\{\}^1$

Prepare Data, Tokenizer, Feature Extractor

ASR models transcribe speech to text, which implies that we each need a
feature extractor that processes the speech signal to the model’s input
format, e.g. a feature vector, and a tokenizer that processes the
model’s output format to text.

In 🤗 Transformers, the Wav2Vec2 model is thus accompanied by each a
tokenizer, called
Wav2Vec2CTCTokenizer,
and a feature extractor, called
Wav2Vec2FeatureExtractor.

Let’s start by creating the tokenizer chargeable for decoding the
model’s predictions.

Create Wav2Vec2CTCTokenizer

The pretrained Wav2Vec2 checkpoint maps the speech signal to a
sequence of context representations as illustrated within the figure above.
A fine-tuned Wav2Vec2 checkpoint must map this sequence of context
representations to its corresponding transcription in order that a linear
layer needs to be added on top of the transformer block (shown in yellow).
This linear layer is used to classifies each context representation to a
token class analogous how, e.g., after pretraining a linear layer is
added on top of BERT’s embeddings for further classification – cf.
with “BERT” section of this blog post.

The output size of this layer corresponds to the variety of tokens within the
vocabulary, which does not rely on Wav2Vec2’s pretraining task,
but only on the labeled dataset used for fine-tuning. So in the primary
step, we’ll take a have a look at Timit and define a vocabulary based on the
dataset’s transcriptions.

Let’s start by loading the dataset and taking a have a look at its structure.

from datasets import load_dataset, load_metric

timit = load_dataset("timit_asr")

print(timit)

Print Output:

    DatasetDict({
        train: Dataset({
            features: ['file', 'audio', 'text', 'phonetic_detail', 'word_detail', 'dialect_region', 'sentence_type', 'speaker_id', 'id'],
            num_rows: 4620
        })
        test: Dataset({
            features: ['file', 'audio', 'text', 'phonetic_detail', 'word_detail', 'dialect_region', 'sentence_type', 'speaker_id', 'id'],
            num_rows: 1680
        })
    })

Many ASR datasets only provide the goal text, 'text' for every audio
file 'file'. Timit actually provides far more details about each
audio file, equivalent to the 'phonetic_detail', etc., which is why many
researchers select to guage their models on phoneme classification
as an alternative of speech recognition when working with Timit. Nonetheless, we wish
to maintain the notebook as general as possible, so that we are going to only
consider the transcribed text for fine-tuning.

timit = timit.remove_columns(["phonetic_detail", "word_detail", "dialect_region", "id", "sentence_type", "speaker_id"])

Let’s write a brief function to display some random samples of the
dataset and run it a few times to get a sense for the
transcriptions.

from datasets import ClassLabel
import random
import pandas as pd
from IPython.display import display, HTML

def show_random_elements(dataset, num_examples=10):
    assert num_examples <= len(dataset), "Cannot pick more elements than there are within the dataset."
    picks = []
    for _ in range(num_examples):
        pick = random.randint(0, len(dataset)-1)
        while pick in picks:
            pick = random.randint(0, len(dataset)-1)
        picks.append(pick)
    
    df = pd.DataFrame(dataset[picks])
    display(HTML(df.to_html()))

show_random_elements(timit["train"].remove_columns(["file", "audio"]))

Print Output:

Alright! The transcriptions look very clean and the language seems to
correspond more to written text than dialogue. This is sensible taking
under consideration that Timit is
a read speech corpus.

We will see that the transcriptions contain some special characters, such
as ,.?!;:. With out a language model, it is far harder to categorise
speech chunks to such special characters because they do not really
correspond to a characteristic sound unit. E.g., the letter "s" has
a kind of clear sound, whereas the special character "." does
not. Also so as to understand the meaning of a speech signal, it’s
often not essential to incorporate special characters within the
transcription.

As well as, we normalize the text to only have lower case letters.

import re
chars_to_ignore_regex = '[,?.!-;:"]'

def remove_special_characters(batch):
    batch["text"] = re.sub(chars_to_ignore_regex, '', batch["text"]).lower()
    return batch

timit = timit.map(remove_special_characters)

Let’s take a have a look at the preprocessed transcriptions.

show_random_elements(timit["train"].remove_columns(["file", "audio"]))

Print Output:

Good! This looks higher. We have now removed most special characters from
transcriptions and normalized them to lower-case only.

In CTC, it is not uncommon to categorise speech chunks into letters, so we’ll
do the identical here. Let’s extract all distinct letters of the training
and test data and construct our vocabulary from this set of letters.

We write a mapping function that concatenates all transcriptions into
one long transcription after which transforms the string right into a set of
chars. It’s important to pass the argument batched=True to the
map(...) function in order that the mapping function has access to all
transcriptions directly.

def extract_all_chars(batch):
  all_text = " ".join(batch["text"])
  vocab = list(set(all_text))
  return {"vocab": [vocab], "all_text": [all_text]}

vocabs = timit.map(extract_all_chars, batched=True, batch_size=-1, keep_in_memory=True, remove_columns=timit.column_names["train"])

Now, we create the union of all distinct letters within the training dataset
and test dataset and convert the resulting list into an enumerated
dictionary.

vocab_list = list(set(vocabs["train"]["vocab"][0]) | set(vocabs["test"]["vocab"][0]))

vocab_dict = {v: k for k, v in enumerate(vocab_list)}
vocab_dict

Print Output:

{    
     ' ': 21,
     "'": 13,
     'a': 24,
     'b': 17,
     'c': 25,
     'd': 2,
     'e': 9,
     'f': 14,
     'g': 22,
     'h': 8,
     'i': 4,
     'j': 18,
     'k': 5,
     'l': 16,
     'm': 6,
     'n': 7,
     'o': 10,
     'p': 19,
     'q': 3,
     'r': 20,
     's': 11,
     't': 0,
     'u': 26,
     'v': 27,
     'w': 1,
     'x': 23,
     'y': 15,
     'z': 12
}

Cool, we see that each one letters of the alphabet occur within the dataset
(which is just not really surprising) and we also extracted the special
characters " " and '. Note that we didn’t exclude those special
characters because:

• The model has to learn to predict when a word finished or else the
  model prediction would at all times be a sequence of chars which might
  make it unattainable to separate words from one another.
• In English, we’d like to maintain the ' character to distinguish
  between words, e.g., "it's" and "its" which have very
  different meanings.

To make it clearer that " " has its own token class, we give it a more
visible character |. As well as, we also add an “unknown” token so
that the model can later take care of characters not encountered in
Timit’s training set.

vocab_dict["|"] = vocab_dict[" "]
del vocab_dict[" "]

Finally, we also add a padding token that corresponds to CTC’s “blank
token“. The “blank token” is a core component of the CTC algorithm.
For more information, please take a have a look at the “Alignment” section
here.

vocab_dict["[UNK]"] = len(vocab_dict)
vocab_dict["[PAD]"] = len(vocab_dict)
print(len(vocab_dict))

Print Output:

    30

Cool, now our vocabulary is complete and consists of 30 tokens, which
signifies that the linear layer that we are going to add on top of the pretrained
Wav2Vec2 checkpoint can have an output dimension of 30.

Let’s now save the vocabulary as a json file.

import json
with open('vocab.json', 'w') as vocab_file:
    json.dump(vocab_dict, vocab_file)

In a final step, we use the json file to instantiate an object of the
Wav2Vec2CTCTokenizer class.

from transformers import Wav2Vec2CTCTokenizer

tokenizer = Wav2Vec2CTCTokenizer("./vocab.json", unk_token="[UNK]", pad_token="[PAD]", word_delimiter_token="|")

If one desires to re-use the just created tokenizer with the fine-tuned model of this notebook, it’s strongly advised to upload the tokenizer to the 🤗 Hub. Let’s call the repo to which we’ll upload the files
"wav2vec2-large-xlsr-turkish-demo-colab":

repo_name = "wav2vec2-base-timit-demo-colab"

and upload the tokenizer to the 🤗 Hub.

tokenizer.push_to_hub(repo_name)

Great, you’ll be able to see the just created repository under https://huggingface.co//wav2vec2-base-timit-demo-colab

Create Wav2Vec2 Feature Extractor

Speech is a continuous signal and to be treated by computers, it first
needs to be discretized, which is frequently called sampling. The
sampling rate hereby plays a crucial role in that it defines what number of
data points of the speech signal are measured per second. Subsequently,
sampling with a better sampling rate leads to a greater approximation
of the real speech signal but in addition necessitates more values per
second.

A pretrained checkpoint expects its input data to have been sampled more
or less from the identical distribution as the information it was trained on. The
same speech signals sampled at two different rates have a really different
distribution, e.g., doubling the sampling rate leads to data points
being twice as long. Thus, before fine-tuning a pretrained checkpoint of
an ASR model, it’s crucial to confirm that the sampling rate of the information
that was used to pretrain the model matches the sampling rate of the
dataset used to fine-tune the model.

Wav2Vec2 was pretrained on the audio data of
LibriSpeech and
LibriVox which each were sampling with 16kHz. Our fine-tuning dataset,
Timit, was luckily also
sampled with 16kHz. If the fine-tuning dataset would have been sampled
with a rate lower or higher than 16kHz, we first would have needed to up or
downsample the speech signal to match the sampling rate of the information used
for pretraining.

A Wav2Vec2 feature extractor object requires the next parameters to
be instantiated:

• feature_size: Speech models take a sequence of feature vectors as
  an input. While the length of this sequence obviously varies, the
  feature size shouldn’t. Within the case of Wav2Vec2, the feature size
  is 1 since the model was trained on the raw speech signal ${}^2\{\}^2$
• sampling_rate: The sampling rate at which the model is trained on.
• padding_value: For batched inference, shorter inputs must be
  padded with a selected value
• do_normalize: Whether the input ought to be
  zero-mean-unit-variance normalized or not. Normally, speech models
  perform higher when normalizing the input
• return_attention_mask: Whether the model should make use of an
  attention_mask for batched inference. Basically, models should
  at all times make use of the attention_mask to mask padded tokens.
  Nonetheless, resulting from a really specific design alternative of Wav2Vec2‘s
  “base” checkpoint, higher results are achieved when using no
  attention_mask. That is not really helpful for other speech
  models. For more information, one can take a have a look at
  this issue.
  Vital If you should use this notebook to fine-tune
  large-lv60,
  this parameter ought to be set to True.

from transformers import Wav2Vec2FeatureExtractor

feature_extractor = Wav2Vec2FeatureExtractor(feature_size=1, sampling_rate=16000, padding_value=0.0, do_normalize=True, return_attention_mask=False)

Great, Wav2Vec2’s feature extraction pipeline is thereby fully defined!

To make the usage of Wav2Vec2 as user-friendly as possible, the feature
extractor and tokenizer are wrapped right into a single Wav2Vec2Processor
class in order that one only needs a model and processor object.

from transformers import Wav2Vec2Processor

processor = Wav2Vec2Processor(feature_extractor=feature_extractor, tokenizer=tokenizer)

Preprocess Data

Up to now, we have now not checked out the actual values of the speech signal but just the transcription. Along with sentence, our datasets include two more column names path and audio. path states absolutely the path of the audio file. Let’s have a look.

print(timit[0]["path"])

Print Output:

'/root/.cache/huggingface/datasets/downloads/extracted/404950a46da14eac65eb4e2a8317b1372fb3971d980d91d5d5b221275b1fd7e0/data/TRAIN/DR4/MMDM0/SI681.WAV'

Wav2Vec2 expects the input within the format of a 1-dimensional array of 16 kHz. Because of this the audio file needs to be loaded and resampled.

Thankfully, datasets does this routinely by calling the opposite column audio. Let try it out.

common_voice_train[0]["audio"]

Print Output:

{'array': array([-2.1362305e-04,  6.1035156e-05,  3.0517578e-05, ...,
        -3.0517578e-05, -9.1552734e-05, -6.1035156e-05], dtype=float32),
 'path': '/root/.cache/huggingface/datasets/downloads/extracted/404950a46da14eac65eb4e2a8317b1372fb3971d980d91d5d5b221275b1fd7e0/data/TRAIN/DR4/MMDM0/SI681.WAV',
 'sampling_rate': 16000}

We will see that the audio file has routinely been loaded. That is because of the brand new "Audio" feature introduced in datasets == 4.13.3, which loads and resamples audio files on-the-fly upon calling.

The sampling rate is ready to 16kHz which is what Wav2Vec2 expects as an input.

Great, let’s take heed to a few audio files to higher understand the dataset and confirm that the audio was accurately loaded.

import IPython.display as ipd
import numpy as np
import random

rand_int = random.randint(0, len(timit["train"]))

print(timit["train"][rand_int]["text"])
ipd.Audio(data=np.asarray(timit["train"][rand_int]["audio"]["array"]), autoplay=True, rate=16000)

It might probably be heard, that the speakers change together with their speaking rate, accent, etc. Overall, the recordings sound relatively clear though, which is to be expected from a read speech corpus.

Let’s do a final check that the information is accurately prepared, by printing the form of the speech input, its transcription, and the corresponding sampling rate.

rand_int = random.randint(0, len(timit["train"]))

print("Goal text:", timit["train"][rand_int]["text"])
print("Input array shape:", np.asarray(timit["train"][rand_int]["audio"]["array"]).shape)
print("Sampling rate:", timit["train"][rand_int]["audio"]["sampling_rate"])

Print Output:

    Goal text: she had your dark suit in greasy wash water all yr
    Input array shape: (52941,)
    Sampling rate: 16000

Good! The whole lot looks effective – the information is a 1-dimensional array, the
sampling rate at all times corresponds to 16kHz, and the goal text is
normalized.

Finally, we will process the dataset to the format expected by the model for training. We’ll make use of the map(...) function.

First, we load and resample the audio data, just by calling batch["audio"].
Second, we extract the input_values from the loaded audio file. In our case, the Wav2Vec2Processor only normalizes the information. For other speech models, nevertheless, this step can include more complex feature extraction, equivalent to Log-Mel feature extraction.
Third, we encode the transcriptions to label ids.

Note: This mapping function is example of how the Wav2Vec2Processor class ought to be used. In “normal” context, calling processor(...) is redirected to Wav2Vec2FeatureExtractor‘s call method. When wrapping the processor into the as_target_processor context, nevertheless, the identical method is redirected to Wav2Vec2CTCTokenizer‘s call method.
For more information please check the docs.

def prepare_dataset(batch):
    audio = batch["audio"]

    
    batch["input_values"] = processor(audio["array"], sampling_rate=audio["sampling_rate"]).input_values[0]
    
    with processor.as_target_processor():
        batch["labels"] = processor(batch["text"]).input_ids
    return batch

Let’s apply the information preparation function to all examples.

timit = timit.map(prepare_dataset, remove_columns=timit.column_names["train"], num_proc=4)

Note: Currently datasets make use of torchaudio and librosa for audio loading and resampling. When you want to implement your individual costumized data loading/sampling, be happy to simply make use of the "path" column as an alternative and disrespect the "audio" column.

Training & Evaluation

The information is processed in order that we’re ready to start out organising the
training pipeline. We’ll make use of 🤗’s
Trainer
for which we essentially have to do the next:

• Define a knowledge collator. In contrast to most NLP models, Wav2Vec2 has
  a much larger input length than output length. E.g., a sample of
  input length 50000 has an output length of not more than 100. Given
  the big input sizes, it’s far more efficient to pad the training
  batches dynamically meaning that each one training samples should only be
  padded to the longest sample of their batch and never the general
  longest sample. Subsequently, fine-tuning Wav2Vec2 requires a special
  padding data collator, which we’ll define below

• Evaluation metric. During training, the model ought to be evaluated on
  the word error rate. We must always define a compute_metrics function
  accordingly

• Load a pretrained checkpoint. We want to load a pretrained
  checkpoint and configure it accurately for training.

• Define the training configuration.

After having fine-tuned the model, we’ll accurately evaluate it on the
test data and confirm that it has indeed learned to accurately transcribe
speech.

Set-up Trainer

Let’s start by defining the information collator. The code for the information
collator was copied from this
example.

Without going into too many details, in contrast to the common data
collators, this data collator treats the input_values and labels
otherwise and thus applies to separate padding functions on them
(again making use of Wav2Vec2’s context manager). That is essential
because in speech input and output are of various modalities meaning
that they shouldn’t be treated by the identical padding function. Analogous
to the common data collators, the padding tokens within the labels with
-100 in order that those tokens are not taken under consideration when
computing the loss.

import torch

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union

@dataclass
class DataCollatorCTCWithPadding:
    """
    Data collator that can dynamically pad the inputs received.
    Args:
        processor (:class:`~transformers.Wav2Vec2Processor`)
            The processor used for processing the information.
        padding (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.PaddingStrategy`, `optional`, defaults to :obj:`True`):
            Select a technique to pad the returned sequences (in accordance with the model's padding side and padding index)
            amongst:
            * :obj:`True` or :obj:`'longest'`: Pad to the longest sequence within 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 just 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).
        max_length_labels (:obj:`int`, `optional`):
            Maximum length of the ``labels`` 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 particularly useful to enable the usage of Tensor Cores on NVIDIA hardware with compute capability >=
            7.5 (Volta).
    """

    processor: Wav2Vec2Processor
    padding: Union[bool, str] = True
    max_length: Optional[int] = None
    max_length_labels: Optional[int] = None
    pad_to_multiple_of: Optional[int] = None
    pad_to_multiple_of_labels: Optional[int] = None

    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
        
        
        input_features = [{"input_values": feature["input_values"]} for feature in features]
        label_features = [{"input_ids": feature["labels"]} for feature in features]

        batch = self.processor.pad(
            input_features,
            padding=self.padding,
            max_length=self.max_length,
            pad_to_multiple_of=self.pad_to_multiple_of,
            return_tensors="pt",
        )
        with self.processor.as_target_processor():
            labels_batch = self.processor.pad(
                label_features,
                padding=self.padding,
                max_length=self.max_length_labels,
                pad_to_multiple_of=self.pad_to_multiple_of_labels,
                return_tensors="pt",
            )

        
        labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)

        batch["labels"] = labels

        return batch

Let’s initialize the information collator.

data_collator = DataCollatorCTCWithPadding(processor=processor, padding=True)

Next, the evaluation metric is defined. As mentioned earlier, the
predominant metric in ASR is the word error rate (WER), hence we’ll
use it on this notebook as well.

wer_metric = load_metric("wer")

The model will return a sequence of logit vectors:

$${\mathbf{y}}_1,\dots ,{\mathbf{y}}_{m}mathbf\{y\}\_1,\; ldots,\; mathbf\{y\}\_m$$

with ${\mathbf{y}}_1=f_{\theta }(x_1,\dots ,x_n)[0]mathbf\{y\}\_1\; =\; f\_\{theta\}(x\_1,\; ldots,\; x\_n)[0]$

A logit vector ${\mathbf{y}}_{1}mathbf\{y\}\_1$

def compute_metrics(pred):
    pred_logits = pred.predictions
    pred_ids = np.argmax(pred_logits, axis=-1)

    pred.label_ids[pred.label_ids == -100] = processor.tokenizer.pad_token_id

    pred_str = processor.batch_decode(pred_ids)
    
    label_str = processor.batch_decode(pred.label_ids, group_tokens=False)

    wer = wer_metric.compute(predictions=pred_str, references=label_str)

    return {"wer": wer}

Now, we will load the pretrained Wav2Vec2 checkpoint. The tokenizer’s
pad_token_id have to be to define the model’s pad_token_id or within the
case of Wav2Vec2ForCTC also CTC’s blank token ${}^2\{\}^2$

from transformers import Wav2Vec2ForCTC

model = Wav2Vec2ForCTC.from_pretrained(
    "facebook/wav2vec2-base", 
    ctc_loss_reduction="mean", 
    pad_token_id=processor.tokenizer.pad_token_id,
)

Print Output:

    Some weights of Wav2Vec2ForCTC weren't initialized from the model checkpoint at facebook/wav2vec2-base and are newly initialized: ['lm_head.weight', 'lm_head.bias']
    You need to probably TRAIN this model on a down-stream task to find a way to make use of it for predictions and inference.

The primary component of Wav2Vec2 consists of a stack of CNN layers that
are used to extract acoustically meaningful – but contextually
independent – features from the raw speech signal. This a part of the
model has already been sufficiently trained during pretrainind and as
stated within the paper doesn’t have to
be fine-tuned anymore. Thus, we will set the requires_grad to False
for all parameters of the feature extraction part.

model.freeze_feature_extractor()

In a final step, we define all parameters related to training. To offer
more explanation on among the parameters:

• group_by_length makes training more efficient by grouping training
  samples of comparable input length into one batch. This may
  significantly speed up training time by heavily reducing the general
  variety of useless padding tokens which are passed through the model
• learning_rate and weight_decay were heuristically tuned until
  fine-tuning has change into stable. Note that those parameters strongly
  rely on the Timit dataset and is likely to be suboptimal for other speech
  datasets.

For more explanations on other parameters, one can take a have a look at the
docs.

During training, a checkpoint shall be uploaded asynchronously to the hub every 400 training steps. It means that you can also mess around with the demo widget even while your model continues to be training.

Note: If one doesn’t wish to upload the model checkpoints to the hub, simply set push_to_hub=False.

from transformers import TrainingArguments

training_args = TrainingArguments(
  output_dir=repo_name,
  group_by_length=True,
  per_device_train_batch_size=32,
  evaluation_strategy="steps",
  num_train_epochs=30,
  fp16=True,
  gradient_checkpointing=True, 
  save_steps=500,
  eval_steps=500,
  logging_steps=500,
  learning_rate=1e-4,
  weight_decay=0.005,
  warmup_steps=1000,
  save_total_limit=2,
)

Now, all instances may be passed to Trainer and we’re ready to start out
training!

from transformers import Trainer

trainer = Trainer(
    model=model,
    data_collator=data_collator,
    args=training_args,
    compute_metrics=compute_metrics,
    train_dataset=timit_prepared["train"],
    eval_dataset=timit_prepared["test"],
    tokenizer=processor.feature_extractor,
)

${}^1\{\}^1$

Training

Training will take between 90 and 180 minutes depending on the GPU
allocated to the google colab attached to this notebook. While the trained model yields satisfying
results on Timit‘s test data, it’s under no circumstances an optimally
fine-tuned model. The aim of this notebook is to display how
Wav2Vec2’s base,
large, and
large-lv60
checkpoints may be fine-tuned on any English dataset.

In case you should use this google colab to fine-tune your model, you
should be certain that that your training doesn’t stop resulting from inactivity. A
easy hack to forestall that is to stick the next code into the
console of this tab (right mouse click -> inspect -> Console tab and
insert code).

function ConnectButton(){
    console.log("Connect pushed"); 
    document.querySelector("#top-toolbar > colab-connect-button").shadowRoot.querySelector("#connect").click() 
}
setInterval(ConnectButton,60000);

trainer.train()

Depending in your GPU, it is likely to be possible that you just are seeing an "out-of-memory" error here. On this case, it’s probably best to scale back per_device_train_batch_size to 16 and even less and eventually make use of gradient_accumulation.

Print Output:

The ultimate WER ought to be below 0.3 which is cheap provided that
state-of-the-art phoneme error rates (PER) are slightly below 0.1 (see
leaderboard)
and that WER is frequently worse than PER.

You possibly can now upload the results of the training to the Hub, just execute this instruction:

trainer.push_to_hub()

You possibly can now share this model with all your mates, family, favorite pets: they’ll all load it with the identifier “your-username/the-name-you-picked” so for example:

from transformers import AutoModelForCTC, Wav2Vec2Processor

model = AutoModelForCTC.from_pretrained("patrickvonplaten/wav2vec2-base-timit-demo-colab")
processor = Wav2Vec2Processor.from_pretrained("patrickvonplaten/wav2vec2-base-timit-demo-colab")

Evaluation

In the ultimate part, we evaluate our fine-tuned model on the test set and
mess around with it a bit.

Let’s load the processor and model.

processor = Wav2Vec2Processor.from_pretrained(repo_name)
model = Wav2Vec2ForCTC.from_pretrained(repo_name)

Now, we’ll make use of the map(...) function to predict the
transcription of each test sample and to save lots of the prediction within the
dataset itself. We’ll call the resulting dictionary "results".

Note: we evaluate the test data set with batch_size=1 on purpose
resulting from this issue.
Since padded inputs don’t yield the very same output as non-padded
inputs, a greater WER may be achieved by not padding the input in any respect.

def map_to_result(batch):
  with torch.no_grad():
    input_values = torch.tensor(batch["input_values"], device="cuda").unsqueeze(0)
    logits = model(input_values).logits

  pred_ids = torch.argmax(logits, dim=-1)
  batch["pred_str"] = processor.batch_decode(pred_ids)[0]
  batch["text"] = processor.decode(batch["labels"], group_tokens=False)
  
  return batch

results = timit["test"].map(map_to_result, remove_columns=timit["test"].column_names)

Let’s compute the general WER now.

print("Test WER: {:.3f}".format(wer_metric.compute(predictions=results["pred_str"], references=results["text"])))

Print Output:

    Test WER: 0.221

22.1% WER – not bad! Our demo model would have probably made it on the official leaderboard.

Let’s take a have a look at some predictions to see what errors are made by the model.

Print Output:

show_random_elements(results.remove_columns(["speech", "sampling_rate"]))

It becomes clear that the anticipated transcriptions are acoustically very
much like the goal transcriptions, but often contain spelling or
grammatical errors. This should not be very surprising though provided that
we purely depend on Wav2Vec2 without making use of a language model.

Finally, to higher understand how CTC works, it’s value taking a deeper
have a look at the precise output of the model. Let’s run the primary test sample
through the model, take the anticipated ids and convert them to their
corresponding tokens.

model.to("cuda")

with torch.no_grad():
  logits = model(torch.tensor(timit["test"][:1]["input_values"], device="cuda")).logits

pred_ids = torch.argmax(logits, dim=-1)


" ".join(processor.tokenizer.convert_ids_to_tokens(pred_ids[0].tolist()))

Print Output:

[PAD] [PAD] [PAD] [PAD] [PAD] [PAD] t t h e e | | b b [PAD] u u n n n g g [PAD] a [PAD] [PAD] l l [PAD] o o o [PAD] | w w a a [PAD] s s | | [PAD] [PAD] p l l e e [PAD] [PAD] s s e n n t t t [PAD] l l y y | | | s s [PAD] i i [PAD] t t t [PAD] u u u u [PAD] [PAD] [PAD] a a [PAD] t t e e e d d d | n n e e a a a r | | t h h e | | s s h h h [PAD] o o o [PAD] o o r r [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD]

The output should make it a bit clearer how CTC works in practice. The
model is to some extent invariant to speaking rate because it has learned
to either just repeat the identical token in case the speech chunk to be
classified still corresponds to the identical token. This makes CTC a really
powerful algorithm for speech recognition because the speech file’s
transcription is usually very much independent of its length.

I again advise the reader to try
this very nice blog post to higher
understand CTC.