Views
No views yet
1from datasets import load_dataset
2from transformers import Wav2Vec2ForCTC
3import torchaudio
4import torch
5
6ds = load_dataset("patrickvonplaten/librispeech_asr_dummy", "clean", split="validation")
7
8model = Wav2Vec2ForCTC.from_pretrained("patrickvonplaten/wav2vec2_tiny_random_robust")
9
10def load_audio(batch):
11 batch["samples"], _ = torchaudio.load(batch["file"])
12 return batch
13
14ds = ds.map(load_audio)
15
16input_values = torch.nn.utils.rnn.pad_sequence([torch.tensor(x[0]) for x in ds["samples"][:10]], batch_first=True)
17
18# forward
19logits = model(input_values).logits
20pred_ids = torch.argmax(logits, dim=-1)
21
22# dummy loss
23dummy_labels = pred_ids.clone()
24dummy_labels[dummy_labels == model.config.pad_token_id] = 1 # can't have CTC blank token in label
25dummy_labels = dummy_labels[:, -(dummy_labels.shape[1] // 4):] # make sure labels are shorter to avoid "inf" loss (can still happen though...)
26loss = model(input_values, labels=dummy_labels).loss