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