File size: 3,353 Bytes
ce3c8df | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | """
LibriSpeech dataset wrapper + collate function.
torchaudio.datasets.LIBRISPEECH reads <root>/LibriSpeech/<split>/... and
returns waveforms shaped (1, L).
"""
from typing import cast
import torch
from torch.utils.data import ConcatDataset, Dataset
import torchaudio
from src.tokenizer import ASRTokenizer, PAD_ID
# LIBRISPEECH.__getitem__'s declared return type; ConcatDataset erases it.
LibriSpeechItem = tuple[torch.Tensor, int, str, int, int, int]
class LibriSpeechASR(Dataset):
def __init__(self, root: str, splits, download: bool = False):
if isinstance(splits, str):
splits = [splits]
self.datasets = [
torchaudio.datasets.LIBRISPEECH(root=root, url=split, download=download)
for split in splits
]
self._concat = ConcatDataset(self.datasets)
def __len__(self):
return len(self._concat)
def __getitem__(self, idx):
waveform, sample_rate, transcript, _speaker_id, _chapter_id, _utterance_id = cast(
LibriSpeechItem, self._concat[idx]
)
assert sample_rate == 16000, f"expected 16kHz audio, got {sample_rate}"
return waveform.squeeze(0), transcript
class ASRCollate:
"""Pads a batch and produces the tensors forward_train/forward_eval need:
waveforms, CTC targets, and decoder teacher-forcing input/target pairs."""
def __init__(self, tokenizer: ASRTokenizer):
self.tokenizer = tokenizer
def __call__(self, batch):
waveforms, transcripts = zip(*batch)
wave_lengths = torch.tensor([w.size(0) for w in waveforms], dtype=torch.long)
max_wave_len = int(wave_lengths.max().item())
padded_waveforms = torch.zeros(len(waveforms), max_wave_len, dtype=torch.float32)
for i, w in enumerate(waveforms):
padded_waveforms[i, : w.size(0)] = w
ctc_ids_list = [self.tokenizer.encode(t) for t in transcripts]
ctc_target_lengths = torch.tensor([len(ids) for ids in ctc_ids_list], dtype=torch.long)
max_ctc_len = max(1, int(ctc_target_lengths.max().item()))
ctc_targets = torch.full((len(batch), max_ctc_len), PAD_ID, dtype=torch.long)
for i, ids in enumerate(ctc_ids_list):
if ids:
ctc_targets[i, : len(ids)] = torch.tensor(ids, dtype=torch.long)
dec_pairs = [self.tokenizer.decoder_input_and_target(t) for t in transcripts]
decoder_input_lengths = torch.tensor([len(p[0]) for p in dec_pairs], dtype=torch.long)
max_dec_len = int(decoder_input_lengths.max().item())
decoder_input = torch.full((len(batch), max_dec_len), PAD_ID, dtype=torch.long)
decoder_target = torch.full((len(batch), max_dec_len), PAD_ID, dtype=torch.long)
for i, (inp, tgt) in enumerate(dec_pairs):
decoder_input[i, : len(inp)] = torch.tensor(inp, dtype=torch.long)
decoder_target[i, : len(tgt)] = torch.tensor(tgt, dtype=torch.long)
return {
"waveforms": padded_waveforms,
"wave_lengths": wave_lengths,
"ctc_targets": ctc_targets,
"ctc_target_lengths": ctc_target_lengths,
"decoder_input": decoder_input,
"decoder_input_lengths": decoder_input_lengths,
"decoder_target": decoder_target,
"transcripts": transcripts,
}
|