# delay_dataset.py: The frame-aligned multi-codebook dataset the depth-module models train on: codebook 0 in the # backbone stream, codebooks 1-9 per frame, the three think-block tiers, lazy JSONL loading. from dataclasses import dataclass from typing import Any, Dict, List, Optional, Sequence import torch from torch.utils.data import Dataset import json import logging import os from pathlib import Path from tokenizer import TokenRegistry logger = logging.getLogger(__name__) IGNORE_INDEX = -100 EMPTY_CODE = -1 NO_WORD = -1 PLAN_COT_MODES = ("plan", "tags+plan", "full") LEGACY_COT_MODES = ("prose", "tags", "none") COT_MODES = LEGACY_COT_MODES + PLAN_COT_MODES PLAN_SCHEDULE_MODES = ("stretch", "pack") PLAN_FRAME_RATE_HZ = 21.53 NUM_LEVELS = int(os.environ.get("MVC_NUM_CODEC_LEVELS", "8")) CODEBOOK_SIZE = 2048 def build_plan_schedule(durations: Sequence[int], n_frames: int, mode: str = "stretch") -> List[int]: if mode not in PLAN_SCHEDULE_MODES: raise ValueError( f"plan schedule mode {mode!r} not in {PLAN_SCHEDULE_MODES}") n_frames = int(n_frames) out = [NO_WORD] * max(0, n_frames) n_words = len(durations) if n_words == 0 or n_frames <= 0: return out durs = [int(d) for d in durations] if any(d < 0 for d in durs): raise ValueError(f"negative word duration in plan: {durs}") total = sum(durs) if total <= 0: raise ValueError("plan durations sum to 0 frames — the keystone token " "(§3.1) is degenerate; the sidecar is broken") if mode == "pack": pos = 0 for w, d in enumerate(durs): if pos >= n_frames: break for t in range(pos, min(n_frames, pos + d)): out[t] = w pos += d return out prev, cum = 0, 0 for w, d in enumerate(durs): cum += d boundary = n_frames if w == n_words - 1 else int(round(n_frames * cum / total)) if n_frames - prev >= n_words - w: boundary = max(prev + 1, min(boundary, n_frames - (n_words - 1 - w))) else: boundary = max(prev, min(boundary, n_frames)) for t in range(prev, boundary): out[t] = w prev = boundary return out def _plan_container(row: dict, field: str): raw = row.get(field) if raw is None: return None, False if isinstance(raw, dict): if raw.get("ok") is False: return None, True words = raw.get("words", raw.get("plan")) if words is None: raise ValueError( f"row[{field!r}] is a dict without 'words'/'plan'; keys=" f"{sorted(raw)[:8]}") raw = words if not isinstance(raw, (list, tuple)): raise ValueError( f"row[{field!r}] must be a list of word dicts (or a sidecar object " f"carrying one), got {type(raw).__name__}") return (list(raw) or None), False def normalise_plan_words(raw_words: Sequence[dict], *, n_pitch_bins: int, n_duration_bins: int, n_energy_bins: int, frame_rate: float = PLAN_FRAME_RATE_HZ, where: str = "") -> List[Dict[str, int]]: out: List[Dict[str, int]] = [] for i, w in enumerate(raw_words): tag = f"{where}word {i}" if not isinstance(w, dict): raise ValueError(f"{tag}: expected a dict, got {type(w).__name__}") def _req_int(key, hi): if key not in w or w[key] is None: raise ValueError(f"{tag}: missing {key!r} (keys={sorted(w)[:8]})") try: v = int(w[key]) except (TypeError, ValueError): raise ValueError(f"{tag}: {key}={w[key]!r} is not an integer") if not 0 <= v < hi: raise ValueError(f"{tag}: {key}={v} out of range [0,{hi})") return v dur_bin = _req_int("dur_bin", n_duration_bins) energy_bin = _req_int("energy_bin", n_energy_bins) pb = w.get("pitch_bin", None) if pb is None: pitch_bin = None else: try: pb = int(pb) except (TypeError, ValueError): raise ValueError(f"{tag}: pitch_bin={w['pitch_bin']!r} is not an integer") if pb == n_pitch_bins: pitch_bin = None elif 0 <= pb < n_pitch_bins: pitch_bin = pb else: raise ValueError( f"{tag}: pitch_bin={pb} out of range [0,{n_pitch_bins}] " f"({n_pitch_bins} = the unvoiced sentinel)") if w.get("dur_frames") is not None: frames = int(w["dur_frames"]) elif w.get("dur_s") is not None: frames = max(1, int(round(float(w["dur_s"]) * frame_rate))) else: raise ValueError( f"{tag}: needs 'dur_frames' (or 'dur_s'). The schedule (§3.2) " f"and the decode window (§3.4) both count FRAMES, so a plan " f"without them cannot be consumed — re-run " f"scripts/plan_extract.py quantise, which emits dur_frames.") if frames < 1: raise ValueError(f"{tag}: dur_frames={frames} must be >= 1") out.append({"pitch_bin": pitch_bin, "dur_bin": dur_bin, "energy_bin": energy_bin, "dur_frames": frames}) return out class DelayMimiDataset(Dataset): LAZY_THRESHOLD_BYTES = 2 * 1024 ** 3 def __init__(self, data_path: str, tokenizer, registry: TokenRegistry, max_seq_len: int = 2048, max_frames: int = 0, cfg_dropout: float = 0.0, aligned: bool = False, lazy: bool = None, stage: int = 1, reasoning_weight: float = 0.3, max_reasoning_tokens: int = 512, default_speaker_id: int = None, audio_continue_weight: float = 0.0, cot_mode: str = "prose", plan_field: str = "plan", plan_weight: float = 1.0, tag_weight: float = None, plan_schedule_mode: str = "stretch", plan_lookahead: int = 1, plan_frame_rate: float = PLAN_FRAME_RATE_HZ, min_plan_coverage: float = 0.0, plan_audit_rows: int = 2000): self.tokenizer = tokenizer self.registry = registry self.max_seq_len = max_seq_len self.cfg_dropout = cfg_dropout assert stage in (1, 2, "auto"), f"stage must be 1, 2 or 'auto', got {stage!r}" self.stage = stage if cot_mode in PLAN_COT_MODES: self.plan_mode = True else: assert cot_mode in ("prose", "tags", "none"), cot_mode self.plan_mode = False self.cot_mode = cot_mode self.reasoning_weight = reasoning_weight self.max_reasoning_tokens = max_reasoning_tokens self.plan_weight = float(plan_weight) self.tag_weight = None if tag_weight is None else float(tag_weight) if plan_schedule_mode not in PLAN_SCHEDULE_MODES: raise ValueError(f"plan_schedule_mode must be one of " f"{PLAN_SCHEDULE_MODES}, got {plan_schedule_mode!r}") self.plan_schedule_mode = plan_schedule_mode if plan_lookahead not in (0, 1): raise ValueError(f"plan_lookahead must be 0 or 1, got {plan_lookahead}") self.plan_lookahead = int(plan_lookahead) self.plan_field = plan_field self.plan_frame_rate = float(plan_frame_rate) reg_pitch = getattr(registry, "pitch_ids", []) or [] reg_dur = getattr(registry, "duration_ids", []) or [] reg_energy = getattr(registry, "energy_ids", []) or [] self._n_pitch_bins = len(reg_pitch) self._n_duration_bins = len(reg_dur) self._n_energy_bins = len(reg_energy) if self.plan_mode and not getattr(registry, "plan_enabled", False): raise RuntimeError( f"cot_mode={cot_mode!r} needs the P²-CoT plan vocabulary, but this " f"tokenizer has none (MVC_ENABLE_PLAN_TOKENS is off). Export " f"MVC_ENABLE_PLAN_TOKENS=1 before python starts (and across the " f"singularity boundary as SINGULARITYENV_MVC_ENABLE_PLAN_TOKENS) " f"and rebuild the tokenizer. Refusing to run a plan arm with no " f"plan tokens — it would train as a plain CoT arm and report " f"nothing wrong.") if self.plan_mode and stage == 1: raise ValueError( f"cot_mode={cot_mode!r} with stage=1 emits no block, so " f"no plan would ever reach the model. Use stage=2 (or 'auto' for " f"a mixed corpus).") self.n_rows_seen = 0 self.n_rows_missing_plan = 0 self.n_rows_plan_rejected = 0 self.n_rows_plan_truncated = 0 self.n_rows_reasoning_dropped = 0 self.audio_continue_weight = audio_continue_weight self.default_speaker_id = default_speaker_id self.aligned = aligned self.data_path = data_path if lazy is None: lazy = Path(data_path).stat().st_size > self.LAZY_THRESHOLD_BYTES self.lazy = lazy self._fh = None self._fh_pid = None self.data: List[dict] = [] self.offsets: List[int] = [] if self.lazy: with open(data_path, "rb") as f: pos = f.tell() for line in f: if line.strip(): self.offsets.append(pos) pos = f.tell() logger.info(f"DelayMimiDataset[LAZY]: {len(self.offsets)} rows " f"indexed from {data_path} (max_seq_len={max_seq_len})") else: with open(data_path, "r", encoding="utf-8") as f: for line in f: line = line.strip() if not line: continue d = json.loads(line) resid = d.get("residual_codes") if not resid or len(resid) != NUM_LEVELS - 1: continue if max_frames and len(d["speech_tokens"]) > max_frames: continue self.data.append(d) logger.info(f"DelayMimiDataset: {len(self.data)} rows from {data_path} " f"(max_seq_len={max_seq_len})") self.plan_audit: Optional[Dict[str, Any]] = None if self.plan_mode: self.plan_audit = self._audit_plan_coverage(int(plan_audit_rows), float(min_plan_coverage)) def _audit_plan_coverage(self, n_max: int, min_coverage: float) -> Dict[str, Any]: n_rows = len(self) limit = n_rows if n_max <= 0 else min(n_rows, n_max) scanned = with_plan = rejected = n_words = 0 for i in range(limit): d = self._row(i) scanned += 1 words, was_rejected = _plan_container(d, self.plan_field) if was_rejected: rejected += 1 elif words: with_plan += 1 n_words += len(words) coverage = (with_plan / scanned) if scanned else 0.0 audit = {"scanned": scanned, "rows": n_rows, "with_plan": with_plan, "rejected": rejected, "coverage": coverage, "sampled": limit < n_rows, "mean_words": (n_words / with_plan) if with_plan else 0.0} msg = (f"DelayMimiDataset[PLAN] cot_mode={self.cot_mode}: " f"{with_plan}/{scanned} rows carry row[{self.plan_field!r}] " f"({100 * coverage:.1f}%), {rejected} rejected by the extractor's " f"own gates, mean {audit['mean_words']:.1f} words/plan" + (f" (SAMPLED: first {scanned} of {n_rows} rows)" if audit["sampled"] else "")) if with_plan == 0: raise RuntimeError( msg + " — a plan arm with NO plans is always a bug. Point " "--data_path at a plan-bearing corpus, or use a legacy " "cot_mode.") if coverage < min_coverage: raise RuntimeError( msg + f" — below the required min_plan_coverage={min_coverage:.2f}.") (logger.warning if coverage < 0.95 else logger.info)(msg) return audit def __len__(self): return len(self.offsets) if self.lazy else len(self.data) def __getstate__(self): state = self.__dict__.copy() state["_fh"] = None state["_fh_pid"] = None return state def _row(self, idx: int) -> dict: if not self.lazy: return self.data[idx] if self._fh is None or self._fh_pid != os.getpid(): self._fh = open(self.data_path, "rb") self._fh_pid = os.getpid() self._fh.seek(self.offsets[idx]) d = json.loads(self._fh.readline()) resid = d.get("residual_codes") if not resid or len(resid) != NUM_LEVELS - 1: raise ValueError( f"malformed row {idx} in {self.data_path} (lazy mode trusts the " f"assembler — run scripts/assemble_corpus.py validation)") return d def _plan_token_groups(self, words: Sequence[Dict[str, int]]) -> List[List[int]]: reg = self.registry groups: List[List[int]] = [] for w in words: group = [reg.plan_word_sep_id] if w["pitch_bin"] is not None: group.append(reg.pitch_bin_to_id(w["pitch_bin"])) group.append(reg.duration_bin_to_id(w["dur_bin"])) group.append(reg.energy_bin_to_id(w["energy_bin"])) groups.append(group) return groups def _plan_ids(self, groups: Sequence[Sequence[int]]) -> List[int]: if not groups: return [] ids = [self.registry.plan_start_id] for group in groups: ids.extend(group) ids.append(self.registry.plan_end_id) return ids def _think_block(self, prose_ids: Sequence[int], tag_ids: Sequence[int], plan_groups: Sequence[Sequence[int]]) -> List[int]: content = list(prose_ids) + list(tag_ids) + self._plan_ids(plan_groups) if not content: return [] return [self.registry.think_start_id] + content + [self.registry.think_end_id] def __getitem__(self, idx: int) -> Dict[str, Any]: d = self._row(idx) reg = self.registry if self.cfg_dropout > 0.0 and torch.rand(1).item() < self.cfg_dropout: prompt_ids = [] else: prompt_ids = self.tokenizer.encode(d["prompt"], add_special_tokens=False) self.n_rows_seen += 1 plan_words: List[Dict[str, int]] = [] if self.plan_mode: raw_plan, plan_rejected = _plan_container(d, self.plan_field) if raw_plan: plan_words = normalise_plan_words( raw_plan, n_pitch_bins=self._n_pitch_bins, n_duration_bins=self._n_duration_bins, n_energy_bins=self._n_energy_bins, frame_rate=self.plan_frame_rate, where=f"{self.data_path} row {idx}: ") if plan_rejected: self.n_rows_plan_rejected += 1 if not plan_words: self.n_rows_missing_plan += 1 if self.stage == "auto": row_stage = 2 if str(d.get("reasoning") or "").strip() else 1 if row_stage == 1 and plan_words: row_stage = 2 else: row_stage = self.stage prose_ids: List[int] = [] tag_ids: List[int] = [] plan_groups: List[List[int]] = [] if row_stage == 2: if self.cot_mode in ("prose", "full"): rtext = (d.get("reasoning") or "").strip() for tag in ("", ""): rtext = rtext.replace(tag, "") rtext = rtext.strip() if rtext: prose_ids = self.tokenizer.encode( rtext, add_special_tokens=False)[: self.max_reasoning_tokens] if self.cot_mode == "tags" or self.cot_mode in ("tags+plan", "full"): if (not self.plan_mode) or any( d.get(k) for k in ("emotion_label", "pace", "pitch")): ttext = (f"EMO={d.get('emotion_label') or 'neutral'} " f"PACE={d.get('pace') or 'normal'} " f"PITCH={d.get('pitch') or 'normal'}") for tag in ("", ""): ttext = ttext.replace(tag, "") ttext = ttext.strip() if ttext: tag_ids = self.tokenizer.encode( (" " + ttext) if prose_ids else ttext, add_special_tokens=False)[: self.max_reasoning_tokens] if plan_words: plan_groups = self._plan_token_groups(plan_words) think_block = self._think_block(prose_ids, tag_ids, plan_groups) levels = [d["speech_tokens"]] + d["residual_codes"] T = len(levels[0]) stagger = 0 if self.aligned else NUM_LEVELS - 1 span = T + stagger fixed = 2 + len(prompt_ids) + 1 + 2 budget = self.max_seq_len - fixed if self.plan_mode and think_block and len(think_block) + span > budget: if prose_ids: prose_ids = [] self.n_rows_reasoning_dropped += 1 think_block = self._think_block(prose_ids, tag_ids, plan_groups) if len(think_block) + span > budget and tag_ids: tag_ids = [] think_block = self._think_block(prose_ids, tag_ids, plan_groups) dropped = 0 while plan_groups and len(think_block) + stagger + 1 > budget: plan_groups.pop() plan_words.pop() dropped += 1 think_block = self._think_block(prose_ids, tag_ids, plan_groups) if dropped: self.n_rows_plan_truncated += 1 overhead = 2 + len(prompt_ids) + len(think_block) + 1 + 2 max_span = self.max_seq_len - overhead if span > max_span: T = max(0, max_span - stagger) span = T + stagger levels = [row[:T] for row in levels] if T <= 0: T, span = 0, 0 frame_to_word: List[int] = [] if plan_groups: frame_to_word = build_plan_schedule( [w["dur_frames"] for w in plan_words], T, self.plan_schedule_mode) last = max(frame_to_word) if frame_to_word else NO_WORD if last < len(plan_groups) - 1: del plan_groups[last + 1:] del plan_words[last + 1:] self.n_rows_plan_truncated += 1 think_block = self._think_block(prose_ids, tag_ids, plan_groups) plan_ids = self._plan_ids(plan_groups) codes = [[EMPTY_CODE] * NUM_LEVELS for _ in range(span)] if self.aligned: for s in range(span): for k in range(NUM_LEVELS): codes[s][k] = levels[k][s] else: for s in range(span): for k in range(NUM_LEVELS): t = s - k if 0 <= t < T: codes[s][k] = levels[k][t] audio_placeholder = reg.audio_start_id input_ids = ( [reg.bos_id, reg.user_prompt_id] + prompt_ids + think_block + [reg.audio_start_id] + [audio_placeholder] * span + [reg.audio_end_id, reg.eos_id] ) prompt_end = 2 + len(prompt_ids) think_lo = prompt_end think_hi = prompt_end + len(think_block) audio_start = think_hi + 1 audio_end = audio_start + span labels = list(input_ids) for i in range(prompt_end): labels[i] = IGNORE_INDEX if self.audio_continue_weight <= 0.0: for i in range(audio_start, audio_end): labels[i] = IGNORE_INDEX loss_weights = [1.0] * len(input_ids) for i in range(think_lo, think_hi): loss_weights[i] = self.reasoning_weight tags_lo = think_lo + 1 + len(prose_ids) tags_hi = tags_lo + len(tag_ids) plan_lo = tags_hi plan_hi = plan_lo + len(plan_ids) if self.tag_weight is not None: for i in range(tags_lo, tags_hi): loss_weights[i] = self.tag_weight for i in range(plan_lo, plan_hi): loss_weights[i] = self.plan_weight if self.audio_continue_weight > 0.0: for i in range(audio_start, audio_end): loss_weights[i] = self.audio_continue_weight full_codes = [[EMPTY_CODE] * NUM_LEVELS for _ in range(len(input_ids))] for s in range(span): full_codes[audio_start + s] = codes[s] plan_schedule = None plan_slots = None plan_durations = None if plan_groups: sched = [NO_WORD] * len(input_ids) for t, w in enumerate(frame_to_word): if w == NO_WORD: continue pos = audio_start + t - self.plan_lookahead if pos >= audio_start - 1: sched[pos] = w plan_schedule = sched plan_slots = [] cursor = plan_lo + 1 for group in plan_groups: has_pitch = len(group) == 4 p_pos = cursor + 1 if has_pitch else NO_WORD d_pos = cursor + (2 if has_pitch else 1) e_pos = d_pos + 1 plan_slots.append([p_pos, d_pos, e_pos]) cursor += len(group) plan_durations = [w["dur_frames"] for w in plan_words] sid = d.get("speaker_id") if sid is None: sid = self.default_speaker_id return { "input_ids": torch.tensor(input_ids, dtype=torch.long), "labels": torch.tensor(labels, dtype=torch.long), "codes": torch.tensor(full_codes, dtype=torch.long), "loss_weights": torch.tensor(loss_weights, dtype=torch.float32), "speaker_id": sid, "prosody_vec": d.get("prosody_vec"), "audio_start": audio_start, "audio_end": audio_end, "plan_schedule": (None if plan_schedule is None else torch.tensor(plan_schedule, dtype=torch.long)), "plan_slots": (None if plan_slots is None else torch.tensor(plan_slots, dtype=torch.long)), "plan_durations": (None if plan_durations is None else torch.tensor(plan_durations, dtype=torch.long)), } @dataclass class DelayCollator: pad_token_id: int def __call__(self, batch: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]: input_ids = torch.nn.utils.rnn.pad_sequence( [b["input_ids"] for b in batch], batch_first=True, padding_value=self.pad_token_id) labels = torch.nn.utils.rnn.pad_sequence( [b["labels"] for b in batch], batch_first=True, padding_value=IGNORE_INDEX) B, S = input_ids.shape codes = torch.full((B, S, NUM_LEVELS), EMPTY_CODE, dtype=torch.long) for i, b in enumerate(batch): codes[i, : b["codes"].shape[0]] = b["codes"] out = { "input_ids": input_ids, "labels": labels, "codes": codes, "attention_mask": (input_ids != self.pad_token_id).long(), } if any("loss_weights" in b for b in batch): lw = torch.ones(B, S, dtype=torch.float32) for i, b in enumerate(batch): if "loss_weights" in b: lw[i, : b["loss_weights"].shape[0]] = b["loss_weights"] out["loss_weights"] = lw if any(b.get("speaker_id") is not None for b in batch): out["speaker_ids"] = torch.tensor( [b.get("speaker_id") or 0 for b in batch], dtype=torch.long) if any(b.get("prosody_vec") is not None for b in batch): out["prosody_vecs"] = torch.tensor( [b.get("prosody_vec") or [0.0] * 6 for b in batch], dtype=torch.float32) if any(b.get("plan_schedule") is not None for b in batch): sched = torch.full((B, S), NO_WORD, dtype=torch.long) n_words = max((int(b["plan_slots"].shape[0]) for b in batch if b.get("plan_slots") is not None), default=0) slots = torch.full((B, n_words, 3), NO_WORD, dtype=torch.long) durations = torch.zeros((B, n_words), dtype=torch.long) lengths = torch.zeros(B, dtype=torch.long) for i, b in enumerate(batch): ps = b.get("plan_schedule") if ps is None: continue sched[i, : ps.shape[0]] = ps sl = b["plan_slots"] slots[i, : sl.shape[0]] = sl du = b["plan_durations"] durations[i, : du.shape[0]] = du lengths[i] = sl.shape[0] out["plan_schedule"] = sched out["plan_slots"] = slots out["plan_durations"] = durations out["plan_lengths"] = lengths return out