| """ |
| 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 |
|
|
| |
| 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, |
| } |
|
|
