| """Arm A data: memmapped Vocos mels + byte tokens with per-token language ids; frame-budget bucketing. |
| Beta-binomial alignment priors are computed in DataLoader workers.""" |
| import json, os, random |
| import numpy as np |
| import torch |
| from torch.utils.data import Dataset, Sampler |
|
|
| from .text import encode, LANG2ID |
| from .model import beta_binomial_prior |
|
|
| N_MELS = 100 |
|
|
|
|
| def load_meta(root, datasets, split="train", max_frames=1400, max_tokens=420): |
| items, mm = [], {} |
| for ds in datasets: |
| p = f"{root}/{ds}/meta.jsonl" |
| if not os.path.exists(p) or os.path.getsize(p) == 0: |
| continue |
| mm[ds] = np.memmap(f"{root}/{ds}/mels.f16", dtype=np.float16, mode="r").reshape(-1, N_MELS) |
| for l in open(p, encoding="utf-8"): |
| m = json.loads(l) |
| if m["split"] != split or m["nf"] > max_frames: |
| continue |
| if len(m["text"].encode("utf-8")) + 2 > max_tokens: |
| continue |
| items.append(m) |
| return items, mm |
|
|
|
|
| class TTSData(Dataset): |
| def __init__(self, items, mm, spk2id): |
| self.items, self.mm, self.spk2id = items, mm, spk2id |
|
|
| def __len__(self): |
| return len(self.items) |
|
|
| def __getitem__(self, i): |
| m = self.items[i] |
| mel = np.asarray(self.mm[m["ds"]][m["off"]: m["off"] + m["nf"]], dtype=np.float32) |
| tok, _ = encode(f"<{m['lang']}>{m['text']}</{m['lang']}>") |
| lang = [LANG2ID[m["lang"]]] * len(tok) |
| return dict(tok=tok, lang=lang, mel=mel, spk=self.spk2id[m["spk"]], id=m["id"]) |
|
|
|
|
| def collate(batch): |
| B = len(batch) |
| N = max(len(b["tok"]) for b in batch) |
| T = max(len(b["mel"]) for b in batch) |
| tok = torch.zeros(B, N, dtype=torch.long) |
| lang = torch.zeros(B, N, dtype=torch.long) |
| mel = torch.zeros(B, T, N_MELS) |
| prior = torch.full((B, T, N), -1e4) |
| tl, ml = torch.zeros(B, dtype=torch.long), torch.zeros(B, dtype=torch.long) |
| for i, b in enumerate(batch): |
| n, t = len(b["tok"]), len(b["mel"]) |
| tok[i, :n] = torch.tensor(b["tok"]); lang[i, :n] = torch.tensor(b["lang"]) |
| mel[i, :t] = torch.from_numpy(b["mel"]); tl[i], ml[i] = n, t |
| prior[i, :t, :n] = beta_binomial_prior(n, t) |
| spk = torch.tensor([b["spk"] for b in batch]) |
| return dict(tok=tok, lang=lang, tlens=tl, mel=mel, mlens=ml, spk=spk, prior=prior) |
|
|
|
|
| class FrameBucketSampler(Sampler): |
| """Batches of similar length under a padded-frame budget; dataset weights by repetition.""" |
|
|
| def __init__(self, items, frame_budget=40000, weights=None, seed=0): |
| self.items, self.budget, self.seed, self.epoch = items, frame_budget, seed, 0 |
| self.weights = weights or {} |
|
|
| def _batches(self): |
| rng = random.Random(self.seed + self.epoch) |
| idx = [] |
| for i, m in enumerate(self.items): |
| w = self.weights.get(m["ds"], 1.0) |
| idx += [i] * int(w) + ([i] if rng.random() < (w - int(w)) else []) |
| idx.sort(key=lambda i: self.items[i]["nf"] + rng.random() * 30) |
| batches, cur, mx = [], [], 0 |
| for i in idx: |
| nf = self.items[i]["nf"] |
| if cur and max(mx, nf) * (len(cur) + 1) > self.budget: |
| batches.append(cur); cur, mx = [], 0 |
| cur.append(i); mx = max(mx, nf) |
| if cur: |
| batches.append(cur) |
| rng.shuffle(batches) |
| return batches |
|
|
| def __iter__(self): |
| b = self._batches() |
| self.epoch += 1 |
| return iter(b) |
|
|
| def __len__(self): |
| return len(self._batches()) |
|
|
