Fine-tuning XLS-R for Multi-Lingual ASR with 🤗 Transformers
Patrick von Platen
New (11/2021): This blog post has been updated to feature XLSR'ssuccessor, called XLS-R.Wav2Vec2 is a pretrained model for Automatic Speech Recognition(ASR) and was released in September2020by Alexei Baevski, Michael Auli, and Alex Conneau. Soon after thesuperior performance of Wav2Vec2 was demonstrated on one of the mostpopular English datasets for ASR, calledLibriSpeech,Facebook AI presented a multi-lingual version of Wav2Vec2, calledXLSR. XLSR stands for cross-lingualspeech representations and refers to model's ability to learn speechrepresentations that are useful across multiple languages.XLSR's successor, simply called XLS-R (refering to the''XLM-Rfor Speech''), was released in November 2021 by ArunBabu, Changhan Wang, Andros Tjandra, et al. XLS-R used almost half amillion hours of audio data in 128 languages for self-supervisedpre-training and comes in sizes ranging from 300 milion up to twobillion parameters. You can find the pretrained checkpoints on the 🤗Hub:Wav2Vec2-XLS-R-300MWav2Vec2-XLS-R-1BWav2Vec2-XLS-R-2BSimilar to BERT's masked language modelingobjective, XLS-R learnscontextualized speech representations by randomly masking featurevectors before passing them to a transformer network duringself-supervised pre-training (i.e. diagram on the left below).For fine-tuning, a single linear layer is added on top of thepre-trained network to train the model on labeled data of audiodownstream tasks such as speech recognition, speech translation andaudio classification (i.e. diagram on the right below).XLS-R shows impressive improvements over previous state-of-the-artresults on both speech recognition, speech translation andspeaker/language identification, cf. with Table 3-6, Table 7-10, andTable 11-12 respectively of the official paper.SetupIn this blog, we will give an in-detail explanation of how XLS-R -more specifically the pre-trained checkpointWav2Vec2-XLS-R-300M - can be fine-tuned for ASR.For demonstration purposes, we fine-tune the model on the low resourceASR dataset of CommonVoice that contains onlyca. 4h of validated training data.XLS-R is fine-tuned using Connectionist Temporal Classification (CTC),which is an algorithm that is used to train neural networks forsequence-to-sequence problems, such as ASR and handwriting recognition.I highly recommend reading the well-written blog post SequenceModeling with CTC (2017) by AwniHannun.Before we start, let's install datasets and transformers. Also, weneed the torchaudio to load audio files and jiwer to evaluate ourfine-tuned model using the word error rate(WER) metric 1 {}^1 1.!pip install datasets==1.18.3!pip install transformers==4.11.3!pip install huggingface_hub==0.1!pip install torchaudio!pip install librosa!pip install jiwerWe strongly suggest to upload your training checkpoints directly to theHugging Face Hub while training. The Hugging FaceHub has integrated version control so you canbe sure that no model checkpoint is getting lost during training.To do so you have to store your authentication token from the HuggingFace website (sign up here if youhaven't already!)from huggingface_hub import notebook_loginnotebook_login()Print Output:Login successfulYour token has been saved to /root/.huggingface/tokenThen you need to install Git-LFS to upload your model checkpoints:apt install git-lfs1 {}^1 1 In the paper, the modelwas evaluated using the phoneme error rate (PER), but by far the mostcommon metric in ASR is the word error rate (WER). To keep this notebookas general as possible we decided to evaluate the model using WER.Prepare Data, Tokenizer, Feature ExtractorASR models transcribe speech to text, which means that we both need afeature extractor that processes the speech signal to the model's inputformat, e.g. a feature vector, and a tokenizer that processes themodel's output format to text.In 🤗 Transformers, the XLS-R model is thus accompanied by both atokenizer, calledWav2Vec2CTCTokenizer,and a feature extractor, calledWav2Vec2FeatureExtractor.Let's start by creating the tokenizer to decode the predicted outputclasses to the output transcription.Create Wav2Vec2CTCTokenizerA pre-trained XLS-R model maps the speech signal to a sequence ofcontext representations as illustrated in the figure above. However, forspeech recognition the model has to to map this sequence of contextrepresentations to its corresponding transcription which means that alinear layer has to be added on top of the transformer block (shown inyellow in the diagram above). This linear layer is used to classifyeach context representation to a token class analogous to howa linear layer is added on top of BERT's embeddingsfor further classification after pre-training (cf. with 'BERT' section of the following blogpost).after pretraining a linear layer is added on top of BERT's embeddingsfor further classification - cf. with 'BERT' section of this blogpost.The output size of this layer corresponds to the number of tokens in thevocabulary, which does not depend on XLS-R's pretraining task, butonly on the labeled dataset used for fine-tuning. So in the first step,we will take a look at the chosen dataset of Common Voice and define avocabulary based on the transcriptions.First, let's go to Common Voice officialwebsite and pick alanguage to fine-tune XLS-R on. For this notebook, we will use Turkish.For each language-specific dataset, you can find a language codecorresponding to your chosen language. On CommonVoice, look for the field"Version". The language code then corresponds to the prefix before theunderscore. For Turkish, e.g. the language code is "tr".Great, now we can use 🤗 Datasets' simple API to download the data. Thedataset name is "common_voice", the configuration name corresponds tothe language code, which is "tr" in our case.Common Voice has many different splits including invalidated, whichrefers to data that was not rated as "clean enough" to be considereduseful. In this notebook, we will only make use of the splits "train","validation" and "test".Because the Turkish dataset is so small, we will merge both thevalidation and training data into a training dataset and only use thetest data for validation.from datasets import load_dataset, load_metric, Audiocommon_voice_train = load_dataset("common_voice", "tr", split="train+validation")common_voice_test = load_dataset("common_voice", "tr", split="test")Many ASR datasets only provide the target text, 'sentence' for eachaudio array 'audio' and file 'path'. Common Voice actually providesmuch more information about each audio file, such as the 'accent',etc. Keeping the notebook as general as possible, we only consider thetranscribed text for fine-tuning.common_voice_train = common_voice_train.remove_columns(["accent", "age", "client_id", "down_votes", "gender", "locale", "segment", "up_votes"])common_voice_test = common_voice_test.remove_columns(["accent", "age", "client_id", "down_votes", "gender", "locale", "segment", "up_votes"])Let's write a short function to display some random samples of thedataset and run it a couple of times to get a feeling for thetranscriptions.from datasets import ClassLabelimport randomimport pandas as pdfrom IPython.display import display, HTMLdef show_random_elements(dataset, num_examples=10):assert num_examples <= len(dataset), "Can't pick more elements than there are in 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()))Print Output:IdxSentence1Jonuz, kısa süreli görevi kabul eden tek adaydı.2Biz umudumuzu bu mücadeleden almaktayız.3Sergide beş Hırvat yeniliği sergilendi.4Herşey adıyla bilinmeli.5Kuruluş özelleştirmeye hazır.6Yerleşim yerlerinin manzarası harika.7Olayların failleri bulunamadı.8Fakat bu çabalar boşa çıktı.9Projenin değeri iki virgül yetmiş yedi milyon avro.10Büyük yeniden yapım projesi dört aşamaya bölündü.Alright! The transcriptions look fairly clean. Having translated thetranscribed sentences, it seems that the language corresponds more towritten-out text than noisy dialogue. This makes sense considering thatCommon Voice is acrowd-sourced read speech corpus.We can see that the transcriptions contain some special characters, suchas ,.?!;:. Without a language model, it is much harder to classifyspeech chunks to such special characters because they don't reallycorrespond to a characteristic sound unit. E.g., the letter "s" hasa more or less clear sound, whereas the special character "." doesnot. Also in order to understand the meaning of a speech signal, it isusually not necessary to include special characters in thetranscription.Let's simply remove all characters that don't contribute to themeaning of a word and cannot really be represented by an acoustic soundand normalize the text.import rechars_to_remove_regex = '[\,\?\.\!\-\;\:\"\“\%\‘\”\�\']'def remove_special_characters(batch):batch["sentence"] = re.sub(chars_to_remove_regex, '', batch["sentence"]).lower()return batchcommon_voice_train = common_voice_train.map(remove_special_characters)common_voice_test = common_voice_test.map(remove_special_characters)Let's look at the processed text labels again.show_random_elements(common_voice_train.remove_columns(["path","audio"]))Print Output:IdxTranscription1birisi beyazlar için dediler2maktouf'un cezası haziran ayında sona erdi3orijinalin aksine kıyafetler çıkarılmadı4bunların toplam değeri yüz milyon avroyu buluyor5masada en az iki seçenek bulunuyor6bu hiç de haksız bir heveslilik değil7bu durum bin dokuz yüz doksanlarda ülkenin bölünmesiyle değişti8söz konusu süre altı ay9ancak bedel çok daha yüksek olabilir10başkent fira bir tepenin üzerinde yer alıyorGood! This looks better. We have removed most special characters fromtranscriptions and normalized them to lower-case only.Before finalizing the pre-processing, it is always advantageous toconsult a native speaker of the target language to see whether the textcan be further simplified. For this blog post,Merve was kind enough to take a quicklook and noted that "hatted" characters - like â - aren't reallyused anymore in Turkish and can be replaced by their "un-hatted"equivalent, e.g. a.This means that we should replace a sentence like"yargı sistemi hâlâ sağlıksız" to "yargı sistemi hala sağlıksız".Let's write another short mapping function to further simplify the textlabels. Remember, the simpler the text labels, the easier it is for themodel to learn to predict those labels.def replace_hatted_characters(batch):batch["sentence"] = re.sub('[â]', 'a', batch["sentence"])batch["sentence"] = re.sub('[î]', 'i', batch["sentence"])batch["sentence"] = re.sub('[ô]', 'o', batch["sentence"])batch["sentence"] = re.sub('[û]', 'u', batch["sentence"])return batchcommon_voice_train = common_voice_train.map(replace_hatted_characters)common_voice_test = common_voice_test.map(replace_hatted_characters)In CTC, it is common to classify speech chunks into letters, so we willdo the same here. Let's extract all distinct letters of the trainingand test data and build our vocabulary from this set of letters.We write a mapping function that concatenates all transcriptions intoone long transcription and then transforms the string into a set ofchars. It is important to pass the argument batched=True to themap(...) function so that the mapping function has access to alltranscriptions at once.def extract_all_chars(batch):all_text = " ".join(batch["sentence"])vocab = list(set(all_text))return {"vocab": [vocab], "all_text": [all_text]}vocab_train = common_voice_train.map(extract_all_chars, batched=True, batch_size=-1, keep_in_memory=True, remove_columns=common_voice_train.column_names)vocab_test = common_voice_test.map(extract_all_chars, batched=True, batch_size=-1, keep_in_memory=True, remove_columns=common_voice_test.column_names)Now, we create the union of all distinct letters in the training datasetand test dataset and convert the resulting list into an enumerateddictionary.vocab_list = list(set(vocab_train["vocab"][0]) | set(vocab_test["vocab"][0]))vocab_dict = {v: k for k, v in enumerate(sorted(vocab_list))}vocab_dictPrint Output:{' ': 0,'a': 1,'b': 2,'c': 3,'d': 4,'e': 5,'f': 6,'g': 7,'h': 8,'i': 9,'j': 10,'k': 11,'l': 12,'m': 13,'n': 14,'o': 15,'p': 16,'q': 17,'r': 18,'s': 19,'t': 20,'u': 21,'v': 22,'w': 23,'x': 24,'y': 25,'z': 26,'ç': 27,'ë': 28,'ö': 29,'ü': 30,'ğ': 31,'ı': 32,'ş': 33,'̇': 34}Cool, we see that all letters of the alphabet occur in the dataset(which is not really surprising) and we also extracted the specialcharacters "" and '. Note that we did not exclude those specialcharacters because:The model has to learn to predict when a word is finished or else themodel prediction would always be a sequence of chars which would make itimpossible to separate words from each other.One should always keep in mind that pre-processing is a very importantstep before training your model. E.g., we don't want our model todifferentiate between a and A just because we forgot to normalizethe data. The difference between a and A does not depend on the"sound" of the letter at all, but more on grammatical rules - e.g.use a capitalized letter at the beginning of the sentence. So it issensible to remove the difference between capitalized andnon-capitalized letters so that the model has an easier time learning totranscribe speech.To make it clearer that " " has its own token class, we give it a morevisible character |. In addition, we also add an "unknown" token sothat the model can later deal with characters not encountered in CommonVoice's training set.vocab_dict["|"] = vocab_dict[" "]del vocab_dict[" "]Finally, we also add a padding token that corresponds to CTC's "blanktoken". The "blank token" is a core component of the CTC algorithm.For more information, please take a look at the "Alignment" sectionhere.vocab_dict["[UNK]"] = len(vocab_dict)vocab_dict["[PAD]"] = len(vocab_dict)len(vocab_dict)Cool, now our vocabulary is complete and consists of 39 tokens, whichmeans that the linear layer that we will add on top of the pretrainedXLS-R checkpoint will have an output dimension of 39.Let's now save the vocabulary as a json file.import jsonwith open('vocab.json', 'w') as vocab_file:json.dump(vocab_dict, vocab_file)In a final step, we use the json file to load the vocabulary into aninstance of the Wav2Vec2CTCTokenizer class.from transformers import Wav2Vec2CTCTokenizertokenizer = Wav2Vec2CTCTokenizer.from_pretrained("./", unk_token="[UNK]", pad_token="[PAD]", word_delimiter_token="|")If one wants to re-use the just created tokenizer with the fine-tunedmodel of this notebook, it is strongly advised to upload the tokenizerto the Hugging Face Hub. Let's call the repo to whichwe will upload the files "wav2vec2-large-xlsr-turkish-demo-colab":repo_name = "wav2vec2-large-xls-r-300m-tr-colab"and upload the tokenizer to the 🤗 Hub.tokenizer.push_to_hub(repo_name)Great, you can see the just created repository underhttps://huggingface.co/<your-username>/wav2vec2-large-xls-r-300m-tr-colabCreate Wav2Vec2FeatureExtractorSpeech is a continuous signal, and, to be treated by computers, it firsthas to be discretized, which is usually called sampling. Thesampling rate hereby plays an important role since it defines how manydata points of the speech signal are measured per second. Therefore,sampling with a higher sampling rate results in a better approximationof the real speech signal but also necessitates more values persecond.A pretrained checkpoint expects its input data to have been sampled moreor less from the same distribution as the data it was trained on. Thesame speech signals sampled at two different rates have a very differentdistribution. For example, doubling the sampling rate results in data pointsbeing twice as long. Thus, before fine-tuning a pretrained checkpoint ofan ASR model, it is crucial to verify that the sampling rate of the datathat was used to pretrain the model matches the sampling rate of thedataset used to fine-tune the model.XLS-R was pretrained on audio data ofBabel,Multilingual LibriSpeech(MLS),Common Voice,VoxPopuli, andVoxLingua107 at a sampling rate of16kHz. Common Voice, in its original form, has a sampling rate of 48kHz,thus we will have to downsample the fine-tuning data to 16kHz in thefollowing.A Wav2Vec2FeatureExtractor object requires the following parameters tobe instantiated:feature_size: Speech models take a sequence of feature vectors asan input. While the length of this sequence obviously varies, thefeature size should not. In the case of Wav2Vec2, the feature sizeis 1 because the model was trained on the raw speech signal 2 {}^2 2.sampling_rate: The sampling rate at which the model is trained on.padding_value: For batched inference, shorter inputs need to bepadded with a specific valuedo_normalize: Whether the input should bezero-mean-unit-variance normalized or not. Usually, speech modelsperform better when normalizing the inputreturn_attention_mask: Whether the model should make use of anattention_mask for batched inference. In general, XLS-R modelscheckpoints should always use the attention_mask.from transformers import Wav2Vec2FeatureExtractorfeature_extractor = Wav2Vec2FeatureExtractor(feature_size=1, sampling_rate=16000, padding_value=0.0, do_normalize=True, return_attention_mask=True)Great, XLS-R's feature extraction pipeline is thereby fully defined!For improved user-friendliness, the feature extractor and tokenizer arewrapped into a single Wav2Vec2Processor class so that one only needsa model and processor object.from transformers import Wav2Vec2Processorprocessor = Wav2Vec2Processor(feature_extractor=feature_extractor, tokenizer=tokenizer)Next, we can prepare the dataset.Preprocess DataSo far, we have not looked at the actual values of the speech signal butjust the transcription. In addition to sentence, our datasets includetwo more column names path and audio. path states the absolutepath of the audio file. Let's take a look.common_voice_train[0]["path"]XLS-R expects the input in the format of a 1-dimensional array of 16kHz. This means that the audio file has to be loaded and resampled.Thankfully, datasets does this automatically by calling the othercolumn audio. Let try it out.common_voice_train[0]["audio"]{'array': array([ 0.0000000e+00, 0.0000000e+00, 0.0000000e+00, ...,-8.8930130e-05, -3.8027763e-05, -2.9146671e-05], dtype=float32),'path': '/root/.cache/huggingface/datasets/downloads/extracted/05be0c29807a73c9b099873d2f5975dae6d05e9f7d577458a2466ecb9a2b0c6b/cv-corpus-6.1-2020-12-11/tr/clips/common_voice_tr_21921195.mp3','sampling_rate': 48000}Great, we can see that the audio file has automatically been loaded.This is thanks to the new "Audio"featureintroduced in datasets == 1.18.3, which loads and resamples audiofiles on-the-fly upon calling.In the example above we can see that the audio data is loaded with asampling rate of 48kHz whereas 16kHz are expected by the model. We canset the audio feature to the correct sampling rate by making use ofcast_column:common_voice_train = common_voice_train.cast_column("audio", Audio(sampling_rate=16_000))common_voice_test = common_voice_test.cast_column("audio", Audio(sampling_rate=16_000))Let's take a look at "audio" again.common_voice_train[0]["audio"]{'array': array([ 0.0000000e+00, 0.0000000e+00, 0.0000000e+00, ...,-7.4556941e-05, -1.4621433e-05, -5.7861507e-05], dtype=float32),'path': '/root/.cache/huggingface/datasets/downloads/extracted/05be0c29807a73c9b099873d2f5975dae6d05e9f7d577458a2466ecb9a2b0c6b/cv-corpus-6.1-2020-12-11/tr/clips/common_voice_tr_21921195.mp3','sampling_rate': 16000}This seemed to have worked! Let's listen to a couple of audio files tobetter understand the dataset and verify that the audio was correctlyloaded.import IPython.display as ipdimport numpy as npimport randomrand_int = random.randint(0, len(common_voice_train)-1)print(common_voice_train[rand_int]["sentence"])ipd.Audio(data=common_voice_train[rand_int]["audio"]["array"], autoplay=True, rate=16000)Print Output:sunulan bütün teklifler i̇ngilizce idiIt seems like the data is now correctly loaded and resampled.It can be heard, that the speakers change along with their speakingrate, accent, and background environment, etc. Overall, the recordingssound acceptably clear though, which is to be expected from acrowd-sourced read speech corpus.Let's do a final check that the data is correctly prepared, by printingthe shape of the speech input, its transcription, and the correspondingsampling rate.rand_int = random.randint(0, len(common_voice_train)-1)print("Target text:", common_voice_train[rand_int]["sentence"])print("Input array shape:", common_voice_train[rand_int]["audio"]["array"].shape)print("Sampling rate:", common_voice_train[rand_int]["audio"]["sampling_rate"])Print Output:Target text: makedonya bu yıl otuz adet tyetmiş iki tankı aldıInput array shape: (71040,)Sampling rate: 16000Good! Everything looks fine - the data is a 1-dimensional array, thesampling rate always corresponds to 16kHz, and the target text isnormalized.Finally, we can leverage Wav2Vec2Processor to process the data to theformat expected by Wav2Vec2ForCTC for training. To do so let's makeuse of Dataset'smap(...)function.First, we load and resample the audio data, simply by callingbatch["audio"]. Second, we extract the input_values from the loadedaudio file. In our case, the Wav2Vec2Processor only normalizes thedata. For other speech models, however, this step can include morecomplex feature extraction, such as Log-Mel featureextraction.Third, we encode the transcriptions to label ids.Note: This mapping function is a good example of how theWav2Vec2Processor class should be used. In "normal" context, callingprocessor(...) is redirected to Wav2Vec2FeatureExtractor's callmethod. When wrapping the processor into the as_target_processorcontext, however, the same method is redirected toWav2Vec2CTCTokenizer's call method. For more information please checkthedocs.def prepare_dataset(batch):audio = batch["audio"]# batched output is "un-batched"batch["input_values"] = processor(audio["array"], sampling_rate=audio["sampling_rate"]).input_values[0]batch["input_length"] = len(batch["input_values"])with processor.as_target_processor():batch["labels"] = processor(batch["sentence"]).input_idsreturn batchLet's apply the data preparation function to all examples.common_voice_train = common_voice_train.map(prepare_dataset, remove_columns=common_voice_train.column_names)common_voice_test = common_voice_test.map(prepare_dataset, remove_columns=common_voice_test.column_names)Note: Currently datasets make use oftorchaudio andlibrosa for audio loadingand resampling. If you wish to implement your own costumized dataloading/sampling, feel free to just make use of the "path" columninstead and disregard the "audio" column.Long input sequences require a lot of memory. XLS-R is based onself-attention. The memory requirement scales quadratically with theinput length for long input sequences (cf. withthisreddit post). In case this demo crashes with an "Out-of-memory" errorfor you, you might want to uncomment the following lines to filter allsequences that are longer than 5 seconds for training.#max_input_length_in_sec = 5.0#common_voice_train = common_voice_train.filter(lambda x: x < max_input_length_in_sec * processor.feature_extractor.sampling_rate, input_columns=["input_length"])Awesome, now we are ready to start training!TrainingThe data is processed so that we are ready to start setting up thetraining pipeline. We will make use of 🤗'sTrainerfor which we essentially need to do the following:Define a data collator. In contrast to most NLP models, XLS-R has amuch larger input length than output length. E.g., a sample ofinput length 50000 has an output length of no more than 100. Giventhe large input sizes, it is much more efficient to pad the trainingbatches dynamically meaning that all training samples should only bepadded to the longest sample in their batch and not the overalllongest sample. Therefore, fine-tuning XLS-R requires a specialpadding data collator, which we will define belowEvaluation metric. During training, the model should be evaluated onthe word error rate. We should define a compute_metrics functionaccordinglyLoad a pretrained checkpoint. We need to load a pretrainedcheckpoint and configure it correctly for training.Define the training configuration.After having fine-tuned the model, we will correctly evaluate it on thetest data and verify that it has indeed learned to correctly transcribespeech.Set-up TrainerLet's start by defining the data collator. The code for the datacollator was copied from thisexample.Without going into too many details, in contrast to the common datacollators, this data collator treats the input_values and labelsdifferently and thus applies to separate padding functions on them(again making use of XLS-R processor's context manager). This isnecessary because in speech input and output are of different modalitiesmeaning that they should not be treated by the same padding function.Analogous to the common data collators, the padding tokens in the labelswith -100 so that those tokens are not taken into account whencomputing the loss.import torchfrom dataclasses import dataclass, fieldfrom typing import Any, Dict, List, Optional, Union@dataclassclass DataCollatorCTCWithPadding:"""Data collator that will dynamically pad the inputs received.Args:processor (:class:`~transformers.Wav2Vec2Processor`)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 singlesequence if provided).* :obj:`'max_length'`: Pad to a maximum length specified with the argument :obj:`max_length` or to themaximum 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 ofdifferent lengths)."""processor: Wav2Vec2Processorpadding: Union[bool, str] = Truedef __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:# split inputs and labels since they have to be of different lengths and need# different padding methodsinput_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,return_tensors="pt",)with self.processor.as_target_processor():labels_batch = self.processor.pad(label_features,padding=self.padding,return_tensors="pt",)# replace padding with -100 to ignore loss correctlylabels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)batch["labels"] = labelsreturn batchdata_collator = DataCollatorCTCWithPadding(processor=processor, padding=True)Next, the evaluation metric is defined. As mentioned earlier, thepredominant metric in ASR is the word error rate (WER), hence we willuse it in this notebook as well.wer_metric = load_metric("wer")The model will return a sequence of logit vectors: y1,…,ym \mathbf{y}_1, \ldots, \mathbf{y}_m y1​,…,ym​ with y1=fθ(x1,…,xn)[0] \mathbf{y}_1 = f_{\theta}(x_1, \ldots, x_n)[0] y1​=fθ​(x1​,…,xn​)[0] and n>>m n >> m n>>m.A logit vector y1 \mathbf{y}_1 y1​ contains the log-odds for each word in thevocabulary we defined earlier, thus len(yi)= \text{len}(\mathbf{y}_i) = len(yi​)=config.vocab_size. We are interested in the most likely prediction ofthe model and thus take the argmax(...) of the logits. Also, wetransform the encoded labels back to the original string by replacing-100 with the pad_token_id and decoding the ids while making surethat consecutive tokens are not grouped to the same token in CTCstyle 1 {}^1 1.def compute_metrics(pred):pred_logits = pred.predictionspred_ids = np.argmax(pred_logits, axis=-1)pred.label_ids[pred.label_ids == -100] = processor.tokenizer.pad_token_idpred_str = processor.batch_decode(pred_ids)# we do not want to group tokens when computing the metricslabel_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 can load the pretrained checkpoint ofWav2Vec2-XLS-R-300M.The tokenizer's pad_token_id must be to define the model'spad_token_id or in the case of Wav2Vec2ForCTC also CTC's blanktoken 2 {}^2 2. To save GPU memory, we enable PyTorch's gradientcheckpointing and alsoset the loss reduction to "mean".Because the dataset is quite small (~6h of training data) and becauseCommon Voice is quite noisy, fine-tuning Facebook'swav2vec2-xls-r-300m checkpoint seems to require somehyper-parameter tuning. Therefore, I had to play around a bit withdifferent values for dropout,SpecAugment's masking dropout rate,layer dropout, and the learning rate until training seemed to be stableenough.Note: When using this notebook to train XLS-R on another language ofCommon Voice those hyper-parameter settings might not work very well.Feel free to adapt those depending on your use case.from transformers import Wav2Vec2ForCTCmodel = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-xls-r-300m", attention_dropout=0.0,hidden_dropout=0.0,feat_proj_dropout=0.0,mask_time_prob=0.05,layerdrop=0.0,ctc_loss_reduction="mean", pad_token_id=processor.tokenizer.pad_token_id,vocab_size=len(processor.tokenizer),)The first component of XLS-R consists of a stack of CNN layers that areused to extract acoustically meaningful - but contextually independent -features from the raw speech signal. This part of the model has alreadybeen sufficiently trained during pretraining and as stated in thepaper does not need to befine-tuned anymore. Thus, we can set the requires_grad to False forall parameters of the feature extraction part.model.freeze_feature_extractor()In a final step, we define all parameters related to training. To givemore explanation on some of the parameters:group_by_length makes training more efficient by grouping trainingsamples of similar input length into one batch. This cansignificantly speed up training time by heavily reducing the overallnumber of useless padding tokens that are passed through the modellearning_rate and weight_decay were heuristically tuned untilfine-tuning has become stable. Note that those parameters stronglydepend on the Common Voice dataset and might be suboptimal for otherspeech datasets.For more explanations on other parameters, one can take a look at thedocs.During training, a checkpoint will be uploaded asynchronously to the Hubevery 400 training steps. It allows you to also play around with thedemo widget even while your model is still training.Note: If one does not want to upload the model checkpoints to theHub, simply set push_to_hub=False.from transformers import TrainingArgumentstraining_args = TrainingArguments(output_dir=repo_name,group_by_length=True,per_device_train_batch_size=16,gradient_accumulation_steps=2,evaluation_strategy="steps",num_train_epochs=30,gradient_checkpointing=True,fp16=True,save_steps=400,eval_steps=400,logging_steps=400,learning_rate=3e-4,warmup_steps=500,save_total_limit=2,push_to_hub=True,)Now, all instances can be passed to Trainer and we are ready to starttraining!from transformers import Trainertrainer = Trainer(model=model,data_collator=data_collator,args=training_args,compute_metrics=compute_metrics,train_dataset=common_voice_train,eval_dataset=common_voice_test,tokenizer=processor.feature_extractor,)1 {}^1 1 To allow models to become independent of the speaker rate, inCTC, consecutive tokens that are identical are simply grouped as asingle token. However, the encoded labels should not be grouped whendecoding since they don't correspond to the predicted tokens of themodel, which is why the group_tokens=False parameter has to be passed.If we wouldn't pass this parameter a word like "hello" wouldincorrectly be encoded, and decoded as "helo".2 {}^2 2 The blank token allows the model to predict a word, such as"hello" by forcing it to insert the blank token between the two l's.A CTC-conform prediction of "hello" of our model would be[PAD] [PAD] "h" "e" "e" "l" "l" [PAD] "l" "o" "o" [PAD].TrainingTraining will take multiple hours depending on the GPU allocated to thisnotebook. While the trained model yields somewhat satisfying results onCommon Voice's test data of Turkish, it is by no means an optimallyfine-tuned model. The purpose of this notebook is just to demonstratehow to fine-tune XLS-R XLSR-Wav2Vec2's on an ASR dataset.Depending on what GPU was allocated to your google colab it might bepossible that you are seeing an "out-of-memory" error here. In thiscase, it's probably best to reduce per_device_train_batch_size to 8or even less and increasegradient_accumulation.trainer.train()Print Output:Training LossEpochStepValidation LossWer3.88423.674000.67940.70000.41157.348000.43040.45480.194611.0112000.44660.42160.130814.6816000.45260.39610.099718.3520000.45670.36960.078422.0224000.41930.34420.063325.6928000.41530.33470.049829.3632000.40770.3195The training loss and validation WER go down nicely.You can now upload the result of the training to the Hub, just executethis instruction:trainer.push_to_hub()You can now share this model with all your friends, family, favoritepets: they can all load it with the identifier"your-username/the-name-you-picked" so for instance:from transformers import AutoModelForCTC, Wav2Vec2Processormodel = AutoModelForCTC.from_pretrained("patrickvonplaten/wav2vec2-large-xls-r-300m-tr-colab")processor = Wav2Vec2Processor.from_pretrained("patrickvonplaten/wav2vec2-large-xls-r-300m-tr-colab")For more examples of how XLS-R can be fine-tuned, please take a look at the official 🤗 Transformers examples.EvaluationAs a final check, let's load the model and verify that it indeed haslearned to transcribe Turkish speech.Let's first load the pretrained checkpoint.model = Wav2Vec2ForCTC.from_pretrained(repo_name).to("cuda")processor = Wav2Vec2Processor.from_pretrained(repo_name)Now, we will just take the first example of the test set, run it throughthe model and take the argmax(...) of the logits to retrieve thepredicted token ids.input_dict = processor(common_voice_test[0]["input_values"], return_tensors="pt", padding=True)logits = model(input_dict.input_values.to("cuda")).logitspred_ids = torch.argmax(logits, dim=-1)[0]It is strongly recommended to pass the sampling_rate argument to this function.Failing to do so can result in silent errors that might be hard to debug.We adapted common_voice_test quite a bit so that the dataset instancedoes not contain the original sentence label anymore. Thus, we re-usethe original dataset to get the label of the first example.common_voice_test_transcription = load_dataset("common_voice", "tr", data_dir="./cv-corpus-6.1-2020-12-11", split="test")Finally, we can decode the example.print("Prediction:")print(processor.decode(pred_ids))print("Reference:")print(common_voice_test_transcription[0]["sentence"].lower())Print Output:pred_strtarget_texthatta küçük şeyleri için bir büyt bir şeyleri kolluyor veyınıki çuk şeyler için bir bir mizi inciltiyoruzhayatta küçük şeyleri kovalıyor ve yine küçük şeyler için birbirimizi incitiyoruz.Alright! The transcription can definitely be recognized from ourprediction, but it is not perfect yet. Training the model a bit longer,spending more time on the data preprocessing, and especially using alanguage model for decoding would certainly improve the model's overallperformance.For a demonstration model on a low-resource language, the results arequite acceptable however 🤗.