When using this model, make sure that your speech input is sampled at 16kHz.
1import json
2import random
3import torch
4import torchaudio
5from datasets import load_dataset
6from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
7
8
9#Load test_dataset from saved files in folder
10from datasets import load_dataset, load_metric
11
12#for test
13for root, dirs, files in os.walk(test/):
14 test_dataset= load_dataset("json", data_files=[os.path.join(root,i) for i in files],split="train")
15
16#Remove unnecessary chars
17chars_to_ignore_regex = '[\\,\\?\\.\\!\\-\\;\\:\\"\\“\\%\\‘\\”]'
18def remove_special_characters(batch):
19 batch["sentence"] = re.sub(chars_to_ignore_regex, '', batch["sentence"]).lower() + " "
20 return batch
21
22test_dataset = test_dataset.map(remove_special_characters)
23
24processor = Wav2Vec2Processor.from_pretrained("chrisjay/wav2vec2-large-xlsr-53-fon")
25model = Wav2Vec2ForCTC.from_pretrained("chrisjay/wav2vec2-large-xlsr-53-fon")
26
27#No need for resampling because audio dataset already at 16kHz
28#resampler = torchaudio.transforms.Resample(48_000, 16_000)
29
30# Preprocessing the datasets.
31# We need to read the audio files as arrays
32def speech_file_to_array_fn(batch):
33 speech_array, sampling_rate = torchaudio.load(batch["path"])
34 batch["speech"]=speech_array.squeeze().numpy()
35 return batch
36
37test_dataset = test_dataset.map(speech_file_to_array_fn)
38inputs = processor(test_dataset["speech"][:2], sampling_rate=16_000, return_tensors="pt", padding=True)
39
40with torch.no_grad():
41 tlogits = model(inputs.input_values, attention_mask=inputs.attention_mask).logits
42
43predicted_ids = torch.argmax(logits, dim=-1)
44
45print("Prediction:", processor.batch_decode(predicted_ids))
46print("Reference:", test_dataset["sentence"][:2])
The model can be evaluated as follows on our unique Fon test data.
1import torch
2import torchaudio
3from datasets import load_dataset, load_metric
4from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
5import re
6
7for root, dirs, files in os.walk(test/):
8 test_dataset = load_dataset("json", data_files=[os.path.join(root,i) for i in files],split="train")
9
10chars_to_ignore_regex = '[\\,\\?\\.\\!\\-\\;\\:\\"\\“\\%\\‘\\”]'
11 batch["sentence"] = re.sub(chars_to_ignore_regex, '', batch["sentence"]).lower() + " "
12 return batch
13
14test_dataset = test_dataset.map(remove_special_characters)
15wer = load_metric("wer")
16
17processor = Wav2Vec2Processor.from_pretrained("chrisjay/wav2vec2-large-xlsr-53-fon")
18model = Wav2Vec2ForCTC.from_pretrained("chrisjay/wav2vec2-large-xlsr-53-fon")
19model.to("cuda")
20
21# Preprocessing the datasets.
22# We need to read the aduio files as arrays
23def speech_file_to_array_fn(batch):
24 speech_array, sampling_rate = torchaudio.load(batch["path"])
25 batch["speech"] = speech_array[0].numpy()
26 batch["sampling_rate"] = sampling_rate
27 batch["target_text"] = batch["sentence"]
28 return batch
29
30test_dataset = test_dataset.map(speech_file_to_array_fn)
31
32#Evaluation on test dataset
33def evaluate(batch):
34 inputs = processor(batch["speech"], sampling_rate=16_000, return_tensors="pt", padding=True)
35
36 with torch.no_grad():
37 logits = model(inputs.input_values.to("cuda"), attention_mask=inputs.attention_mask.to("cuda")).logits
38
39 pred_ids = torch.argmax(logits, dim=-1)
40 batch["pred_strings"] = processor.batch_decode(pred_ids)
41 return batch
42
43result = test_dataset.map(evaluate, batched=True, batch_size=8)
44
45print("WER: {:2f}".format(100 * wer.compute(predictions=result["pred_strings"], references=result["sentence"])))
46
The script used for training can be found
here