Views
No views yet
1 from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
2 from datasets import load_dataset
3 import torch
4
5 # load model and processor
6 processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-large-960h-lv60-self")
7 model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-large-960h-lv60-self")
8
9 # load dummy dataset and read soundfiles
10 ds = load_dataset("patrickvonplaten/librispeech_asr_dummy", "clean", split="validation")
11
12 # tokenize
13 input_values = processor(ds[0]["audio"]["array"], return_tensors="pt", padding="longest").input_values
14
15 # retrieve logits
16 logits = model(input_values).logits
17
18 # take argmax and decode
19 predicted_ids = torch.argmax(logits, dim=-1)
20 transcription = processor.batch_decode(predicted_ids)1from datasets import load_dataset
2from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
3import torch
4from jiwer import wer
5
6
7librispeech_eval = load_dataset("librispeech_asr", "clean", split="test")
8
9model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-large-960h-lv60-self").to("cuda")
10processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-large-960h-lv60-self")
11
12def map_to_pred(batch):
13 inputs = processor(batch["audio"]["array"], return_tensors="pt", padding="longest")
14 input_values = inputs.input_values.to("cuda")
15 attention_mask = inputs.attention_mask.to("cuda")
16
17 with torch.no_grad():
18 logits = model(input_values, attention_mask=attention_mask).logits
19
20 predicted_ids = torch.argmax(logits, dim=-1)
21 transcription = processor.batch_decode(predicted_ids)
22 batch["transcription"] = transcription
23 return batch
24
25result = librispeech_eval.map(map_to_pred, remove_columns=["audio"])
26
27print("WER:", wer(result["text"], result["transcription"]))| "clean" | "other" |
|---|---|
| 1.9 | 3.9 |