Views
No views yet
1 from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
2 from datasets import load_dataset
3 import soundfile as sf
4 import torch
5
6 # load model and tokenizer
7 processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
8 model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")
9
10 # load dummy dataset and read soundfiles
11 ds = load_dataset("patrickvonplaten/librispeech_asr_dummy", "clean", split="validation")
12
13 # tokenize
14 input_values = processor(ds[0]["audio"]["array"], return_tensors="pt", padding="longest").input_values # Batch size 1
15
16 # retrieve logits
17 logits = model(input_values).logits
18
19 # take argmax and decode
20 predicted_ids = torch.argmax(logits, dim=-1)
21 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-base-960h").to("cuda")
10processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
11
12def map_to_pred(batch):
13 input_values = processor(batch["audio"]["array"], return_tensors="pt", padding="longest").input_values
14 with torch.no_grad():
15 logits = model(input_values.to("cuda")).logits
16
17 predicted_ids = torch.argmax(logits, dim=-1)
18 transcription = processor.batch_decode(predicted_ids)
19 batch["transcription"] = transcription
20 return batch
21
22result = librispeech_eval.map(map_to_pred, batched=True, batch_size=1, remove_columns=["audio"])
23
24print("WER:", wer(result["text"], result["transcription"]))| "clean" | "other" |
|---|---|
| 3.4 | 8.6 |