|
|
|
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") |
|
|
|
|
|
|
|
|
|
with open("./testSeqs_uniform_new_version.text", "r") as f: |
|
lines = f.readlines() |
|
|
|
|
|
|
|
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) |
|
|
