File size: 1,590 Bytes
44f4f13 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
#!/usr/bin/env python3
from datasets import load_dataset, load_metric
import datasets
import torch
from transformers import AutoModelForCTC, AutoProcessor
device = "cuda:0" if torch.cuda.is_available() else "cpu"
ds = load_dataset("mozilla-foundation/common_voice_3_0", "it", split="train+validation+test+other")
wer = load_metric("wer")
model = AutoModelForCTC.from_pretrained("microsoft/unispeech-1350-en-90-it-ft-1h")
model = model.to(device)
processor = AutoProcessor.from_pretrained("microsoft/unispeech-1350-en-90-it-ft-1h")
# taken from
# https://github.com/microsoft/UniSpeech/blob/main/UniSpeech/examples/unispeech/data/it/phonesMatches_reduced.json
with open("./testSeqs_uniform_new_version.text", "r") as f:
lines = f.readlines()
# retrieve ids model is evaluated on
ids = [x.split("\t")[0] for x in lines]
ds = ds.filter(lambda p: p.split("/")[-1].split(".")[0] in ids, input_columns=["path"])
ds = ds.cast_column("audio", datasets.Audio(sampling_rate=16_000))
def decode(batch):
input_values = processor(batch["audio"]["array"], return_tensors="pt", sampling_rate=16_000).input_values
input_values = input_values.to(device)
logits = model(input_values).logits
pred_ids = torch.argmax(logits, axis=-1)
batch["id"] = batch["path"]
batch["prediction"] = processor.batch_decode(pred_ids)[0]
batch["target"] = processor.tokenizer.phonemize(batch["sentence"])
return batch
out = ds.map(decode, remove_columns=ds.column_names)
wer_out = wer.compute(predictions=out["prediction"], references=out["target"])
print("wer", wer_out)
|