This is a the demonstration of a fine-tuned Wav2vec model for Brazilian Portuguese using the following datasets:
These datasets were combined to build a larger Brazilian Portuguese dataset. All data was used for training except Common Voice dev/test sets, that were used for validation/test respectively. We also made test sets for all the gathered datasets.
The original model was fine-tuned using
fairseq. This notebook uses a converted version of the original one. The link to the original fairseq model is available
here.
1%%capture
2!pip install torch==1.8.2+cu111 torchvision==0.9.2+cu111 torchaudio===0.8.2 -f https://download.pytorch.org/whl/lts/1.8/torch_lts.html
3!pip install datasets
4!pip install jiwer
5!pip install transformers
6!pip install soundfile
7!pip install pyctcdecode
8!pip install https://github.com/kpu/kenlm/archive/master.zip
1import jiwer
2import torchaudio
3from datasets import load_dataset, load_metric
4from transformers import (
5 Wav2Vec2ForCTC,
6 Wav2Vec2Processor,
7)
8from pyctcdecode import build_ctcdecoder
9import torch
10import re
11import sys
1chars_to_ignore_regex = '[\,\?\.\!\;\:\"]' # noqa: W605
2
3def map_to_array(batch):
4 speech, _ = torchaudio.load(batch["path"])
5 batch["speech"] = speech.squeeze(0).numpy()
6 batch["sampling_rate"] = 16_000
7 batch["sentence"] = re.sub(chars_to_ignore_regex, '', batch["sentence"]).lower().replace("’", "'")
8 batch["target"] = batch["sentence"]
9 return batch
1def calc_metrics(truths, hypos):
2 wers = []
3 mers = []
4 wils = []
5 for t, h in zip(truths, hypos):
6 try:
7 wers.append(jiwer.wer(t, h))
8 mers.append(jiwer.mer(t, h))
9 wils.append(jiwer.wil(t, h))
10 except: # Empty string?
11 pass
12 wer = sum(wers)/len(wers)
13 mer = sum(mers)/len(mers)
14 wil = sum(wils)/len(wils)
15 return wer, mer, wil
1def load_data(dataset):
2 data_files = {'test': f'{dataset}/test.csv'}
3 dataset = load_dataset('csv', data_files=data_files)["test"]
4 return dataset.map(map_to_array)
1class STT:
2
3 def __init__(self,
4 model_name,
5 device='cuda' if torch.cuda.is_available() else 'cpu',
6 lm=None):
7 self.model_name = model_name
8 self.model = Wav2Vec2ForCTC.from_pretrained(model_name).to(device)
9 self.processor = Wav2Vec2Processor.from_pretrained(model_name)
10 self.vocab_dict = self.processor.tokenizer.get_vocab()
11 self.sorted_dict = {
12 k.lower(): v for k, v in sorted(self.vocab_dict.items(),
13 key=lambda item: item[1])
14 }
15 self.device = device
16 self.lm = lm
17 if self.lm:
18 self.lm_decoder = build_ctcdecoder(
19 list(self.sorted_dict.keys()),
20 self.lm
21 )
22
23 def batch_predict(self, batch):
24 features = self.processor(batch["speech"],
25 sampling_rate=batch["sampling_rate"][0],
26 padding=True,
27 return_tensors="pt")
28 input_values = features.input_values.to(self.device)
29 with torch.no_grad():
30 logits = self.model(input_values).logits
31 if self.lm:
32 logits = logits.cpu().numpy()
33 batch["predicted"] = []
34 for sample_logits in logits:
35 batch["predicted"].append(self.lm_decoder.decode(sample_logits))
36 else:
37 pred_ids = torch.argmax(logits, dim=-1)
38 batch["predicted"] = self.processor.batch_decode(pred_ids)
39 return batch
1%%capture
2!gdown --id 1HFECzIizf-bmkQRLiQD0QVqcGtOG5upI
3!mkdir bp_dataset
4!unzip bp_dataset -d bp_dataset/
1ds = load_data('cetuc_dataset')
2result = ds.map(stt.batch_predict, batched=True, batch_size=8)
3wer, mer, wil = calc_metrics(result["sentence"], result["predicted"])
4print("CETUC WER:", wer)
1ds = load_data('commonvoice_dataset')
2result = ds.map(stt.batch_predict, batched=True, batch_size=8)
3wer, mer, wil = calc_metrics(result["sentence"], result["predicted"])
4print("CV WER:", wer)
1ds = load_data('lapsbm_dataset')
2result = ds.map(stt.batch_predict, batched=True, batch_size=8)
3wer, mer, wil = calc_metrics(result["sentence"], result["predicted"])
4print("Laps WER:", wer)
1ds = load_data('mls_dataset')
2result = ds.map(stt.batch_predict, batched=True, batch_size=8)
3wer, mer, wil = calc_metrics(result["sentence"], result["predicted"])
4print("MLS WER:", wer)
1ds = load_data('sid_dataset')
2result = ds.map(stt.batch_predict, batched=True, batch_size=8)
3wer, mer, wil = calc_metrics(result["sentence"], result["predicted"])
4print("Sid WER:", wer)
1ds = load_data('tedx_dataset')
2result = ds.map(stt.batch_predict, batched=True, batch_size=8)
3wer, mer, wil = calc_metrics(result["sentence"], result["predicted"])
4print("TEDx WER:", wer)
1ds = load_data('voxforge_dataset')
2result = ds.map(stt.batch_predict, batched=True, batch_size=8)
3wer, mer, wil = calc_metrics(result["sentence"], result["predicted"])
4print("VoxForge WER:", wer)
1!rm -rf ~/.cache
2!gdown --id 1GJIKseP5ZkTbllQVgOL98R4yYAcIySFP # trained with wikipedia
3stt = STT(MODEL_NAME, lm='pt-BR-wiki.word.4-gram.arpa')
4# !gdown --id 1dLFldy7eguPtyJj5OAlI4Emnx0BpFywg # trained with bp
5# stt = STT(MODEL_NAME, lm='pt-BR.word.4-gram.arpa')
1ds = load_data('cetuc_dataset')
2result = ds.map(stt.batch_predict, batched=True, batch_size=8)
3wer, mer, wil = calc_metrics(result["sentence"], result["predicted"])
4print("CETUC WER:", wer)
1ds = load_data('commonvoice_dataset')
2result = ds.map(stt.batch_predict, batched=True, batch_size=8)
3wer, mer, wil = calc_metrics(result["sentence"], result["predicted"])
4print("CV WER:", wer)
1ds = load_data('lapsbm_dataset')
2result = ds.map(stt.batch_predict, batched=True, batch_size=8)
3wer, mer, wil = calc_metrics(result["sentence"], result["predicted"])
4print("Laps WER:", wer)
1ds = load_data('mls_dataset')
2result = ds.map(stt.batch_predict, batched=True, batch_size=8)
3wer, mer, wil = calc_metrics(result["sentence"], result["predicted"])
4print("MLS WER:", wer)
1ds = load_data('sid_dataset')
2result = ds.map(stt.batch_predict, batched=True, batch_size=8)
3wer, mer, wil = calc_metrics(result["sentence"], result["predicted"])
4print("Sid WER:", wer)
1ds = load_data('tedx_dataset')
2result = ds.map(stt.batch_predict, batched=True, batch_size=8)
3wer, mer, wil = calc_metrics(result["sentence"], result["predicted"])
4print("TEDx WER:", wer)
1ds = load_data('voxforge_dataset')
2result = ds.map(stt.batch_predict, batched=True, batch_size=8)
3wer, mer, wil = calc_metrics(result["sentence"], result["predicted"])
4print("VoxForge WER:", wer)