hf-test commited on
Commit
6db1d2d
1 Parent(s): 1830e0f
Files changed (1) hide show
  1. eval.py +39 -0
eval.py ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ from datasets import load_dataset, load_metric, Audio
3
+ from transformers import AutoModelForCTC, AutoProcessor, Wav2Vec2Processor
4
+ import torch
5
+
6
+ lang = "sv-SE"
7
+ model_id = "./xls-r-300m-sv"
8
+
9
+ device = "cuda" if torch.cuda.is_available() else "cpu"
10
+
11
+ dataset = load_dataset("mozilla-foundation/common_voice_7_0", lang, split="test", use_auth_token=True)
12
+ wer = load_metric("wer")
13
+
14
+ dataset = dataset.select(range(100))
15
+
16
+ dataset = dataset.cast_column("audio", Audio(sampling_rate=16_000))
17
+
18
+ model = AutoModelForCTC.from_pretrained(model_id).to(device)
19
+ processor = Wav2Vec2Processor.from_pretrained(model_id)
20
+
21
+
22
+ def map_to_pred(batch):
23
+ input_values = processor(batch["audio"]["array"], return_tensors="pt", padding="longest", sampling_rate=16_000).input_values
24
+
25
+ with torch.no_grad():
26
+ logits = model(input_values.to(device)).logits
27
+
28
+ predicted_ids = torch.argmax(logits, dim=-1)
29
+ transcription = processor.batch_decode(predicted_ids)[0]
30
+ batch["transcription"] = transcription
31
+ return batch
32
+
33
+
34
+ result = dataset.map(map_to_pred, remove_columns=["audio"])
35
+
36
+ import ipdb; ipdb.set_trace()
37
+ wer_result = wer.compute(references=result["sentence"], predictions=result["transcription"])
38
+
39
+ print("WER", wer_result)