| """LibriSpeech audio data preparation with AudioVQEncoder targets. |
| |
| Dataset: https://huggingface.co/datasets/openslr/librispeech_asr |
| Storage: ~6 GB (FLAC compressed) |
| License: CC-BY-4.0 |
| """ |
| import torch, torchaudio |
| from dataclasses import dataclass |
| import sys, os |
| sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..")) |
| from arbitor.config import SPECIAL_VOCAB, AUDIO_SR |
|
|
|
|
| @dataclass |
| class LibriSpeechConfig: |
| subset: str = "clean" |
| batch_size: int = 4 |
| shuffle_buffer: int = 5000 |
| split: str = "train" |
| max_sec: int = 5 |
| audio_sr: int = AUDIO_SR |
|
|
|
|
| class LibriSpeechStream: |
| def __init__(self, cfg: LibriSpeechConfig): |
| self.cfg = cfg |
| self._ds = None |
| self._vq_enc = None |
| self._bos = SPECIAL_VOCAB['BOS'] |
| self._eos = SPECIAL_VOCAB['EOS'] |
| self._pad = SPECIAL_VOCAB['PAD'] |
|
|
| def _lazy_init(self): |
| if self._ds is not None: |
| return |
| from datasets import load_dataset |
| conf = "clean" if self.cfg.subset == "clean" else "other" |
| subset_map = {"train": f"train.{conf}.100", |
| "dev": f"dev.{conf}", |
| "test": f"test.{conf}"} |
| ds = load_dataset("openslr/librispeech_asr", |
| subset_map.get(self.cfg.split, "train.clean.100"), |
| split="train", streaming=True, trust_remote_code=True) |
| self._ds = ds.shuffle(buffer_size=self.cfg.shuffle_buffer, seed=42) |
|
|
| def _get_vq_encoder(self): |
| if self._vq_enc is None: |
| from arbitor.encoders.audio import AudioVQEncoder |
| self._vq_enc = AudioVQEncoder().cuda() |
| self._vq_enc.eval() |
| return self._vq_enc |
|
|
| def batches(self): |
| self._lazy_init() |
| buf = [] |
| for example in self._ds: |
| audio = example["audio"] |
| wav = torch.tensor(audio["array"], dtype=torch.float32).unsqueeze(0) |
| sr = audio["sampling_rate"] |
| if sr != self.cfg.audio_sr: |
| resample = torchaudio.transforms.Resample(sr, self.cfg.audio_sr) |
| wav = resample(wav) |
| max_samples = self.cfg.max_sec * self.cfg.audio_sr |
| if wav.shape[1] < self.cfg.audio_sr: |
| continue |
| wav = wav[:, :max_samples] |
| text = example.get("text", "") |
| raw = text.encode("utf-8") |
| text_tokens = [self._bos] + list(raw) + [self._eos] |
| text_t = torch.tensor(text_tokens, dtype=torch.long) |
|
|
| vq_enc = self._get_vq_encoder() |
| with torch.no_grad(): |
| _, vq_tokens = vq_enc(wav.unsqueeze(0).cuda()) |
| vq_tokens = vq_tokens.squeeze(0).cpu() |
| vq_t = torch.cat([text_t.new_tensor([self._bos]), vq_tokens, |
| text_t.new_tensor([self._eos])]) |
| buf.append((wav, vq_t)) |
| if len(buf) >= self.cfg.batch_size: |
| batch = buf[:self.cfg.batch_size] |
| buf = buf[self.cfg.batch_size:] |
| waves = torch.stack([b[0] for b in batch]) |
| max_len = max(b[1].numel() for b in batch) |
| targets = torch.stack([ |
| torch.cat([b[1], b[1].new_full((max_len - b[1].numel(),), self._pad)]) |
| for b in batch |
| ]) |
| yield waves, targets |
|
|
| def num_samples(self) -> int: |
| return 290_000 |
|
|
