Views
No views yet
1from datasets import load_dataset
2from transformers import Speech2TextTransformerForConditionalGeneration, Speech2TextTransformerTokenizer
3import soundfile as sf
4from jiwer import wer
5
6librispeech_eval = load_dataset("librispeech_asr", "clean", split="test") # change to "other" for other test dataset
7
8model = Speech2TextTransformerForConditionalGeneration.from_pretrained("valhalla/s2t_librispeech_medium").to("cuda")
9tokenizer = Speech2TextTransformerTokenizer.from_pretrained("valhalla/s2t_librispeech_medium", do_upper_case=True)
10
11def map_to_array(batch):
12 speech, _ = sf.read(batch["file"])
13 batch["speech"] = speech
14 return batch
15
16librispeech_eval = librispeech_eval.map(map_to_array)
17
18def map_to_pred(batch):
19 features = tokenizer(batch["speech"], sample_rate=16000, padding=True, return_tensors="pt")
20 input_features = features.input_features.to("cuda")
21 attention_mask = features.attention_mask.to("cuda")
22
23 gen_tokens = model.generate(input_ids=input_features, attention_mask=attention_mask)
24 batch["transcription"] = tokenizer.batch_decode(gen_tokens, skip_special_tokens=True)
25 return batch
26
27result = librispeech_eval.map(map_to_pred, batched=True, batch_size=8, remove_columns=["speech"])
28
29print("WER:", wer(result["text"], result["transcription"]))| "clean" | "other" |
|---|---|
| 3.5 | 7.8 |