File size: 1,541 Bytes
6db1d2d 21443b2 6db1d2d 21443b2 6db1d2d 21443b2 6db1d2d 21443b2 6db1d2d 21443b2 6db1d2d |
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 |
#!/usr/bin/env python3
from datasets import load_dataset, load_metric, Audio
from transformers import AutoModelForCTC, AutoProcessor, Wav2Vec2Processor
import torch
import re
lang = "sv-SE"
model_id = "./xls-r-300m-sv"
device = "cuda" if torch.cuda.is_available() else "cpu"
dataset = load_dataset("mozilla-foundation/common_voice_7_0", lang, split="test", use_auth_token=True)
wer = load_metric("wer")
dataset = dataset.cast_column("audio", Audio(sampling_rate=16_000))
model = AutoModelForCTC.from_pretrained(model_id).to(device)
processor = AutoProcessor.from_pretrained(model_id)
chars_to_ignore_regex = '[,?.!\-\;\:\"“%‘”�—’…–]' # change to the ignored characters of your fine-tuned model
def map_to_pred(batch):
input_values = processor(batch["audio"]["array"], return_tensors="pt", padding="longest", sampling_rate=16_000).input_values
with torch.no_grad():
logits = model(input_values.to(device)).logits
if processor.__class__.__name__ == "Wav2Vec2Processor":
predicted_ids = torch.argmax(logits, dim=-1)
transcription = processor.batch_decode(predicted_ids)[0]
else:
transcription = processor.batch_decode(logits.cpu().numpy()).text[0]
batch["transcription"] = transcription
batch["text"] = re.sub(chars_to_ignore_regex, "", batch["sentence"].lower())
return batch
result = dataset.map(map_to_pred, remove_columns=["audio"])
wer_result = wer.compute(references=result["text"], predictions=result["transcription"])
print("WER", wer_result)
|