Views
No views yet
1 from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
2 from datasets import load_dataset
3 import torch
4
5 # load model and tokenizer
6 processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
7 model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")
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 # Batch size 1
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-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 |