""" ml/data/make_dataset.py - Build the unified training/eval dataset ================================================================== Merges gold real corpora (UCLASS, SEP-28k, LibriStutter, L2-ARCTIC, CMU-ARCTIC) plus user recordings into a single HuggingFace ``Dataset`` with a clean consistent schema and canonical label set. Believability / anti-leak: the train/validation/test split is made by **speaker**, never by random clip. The model must classify speakers it has never listened to during training — the honest proof of generalization. Audio decode bypass: remote corpora store audio as embedded ``bytes`` inside a pyarrow ``struct`` column. datasets' own ``Audio`` feature decoding routes through ``torchcodec`` (an ffmpeg binding that is DLL-broken on Windows + torch 2.6). So we NEVER touch ``ds[i]`` (which triggers feature decoding); we read the raw pyarrow ``ArrowTable`` columns directly and decode the `bytes` ourselves with soundfile + resample to 16 kHz. ``torchcodec`` is never imported. Output: data/metadata/dataset/ serialized HF dataset data/metadata/dataset.json provenance + counts Usage: python -m ml.data.download_corpora --only stutter_event uclass python -m ml.data.make_dataset # all corpora python -m ml.data.make_dataset --corpora uclass --seed 7 """ from __future__ import annotations import argparse import io import json import re from pathlib import Path from collections import defaultdict from typing import Optional from datasets import Dataset, load_dataset, Audio, Features, Value, Sequence import numpy as np from ml.data.corpora_config import CORPUS_REGISTRY, LABEL_INDEX from ml.data.download_corpora import CORPUS_LOAD OUT_DIR = Path("data/metadata") TARGET_SR = 16000 # Column layout of the unified dataset on disk. 'audio_array' is stored as a # plain variable-length float32 array (NOT an HF Audio feature): many remote # corpora carry path=None, and datasets 5.x crashes when save_to_disk tries to # embed an Audio feature whose path is None. Keeping the raw 16k waveform as a # Sequence is schema-trivial and round-trips reliably. SCHEMA = Features({ "id": Value("string"), "split": Value("string"), "corpus": Value("string"), "speaker_id": Value("string"), "audio_array": Sequence(Value("float32"), length=-1), "text": Value("string"), "label": Value("string"), }) # Canonical UCLASS class-code scheme (standard stutter-fluency labeling used by # the UCLASS archive). 4 (interjection) and 7 are not unambiguous stutter # subtypes in our LABEL_INDEX, so they are dropped here and the drop is # documented in provenance — better a smaller honest set than a guessed label. UCLASS_CLASS_MAP = { "0": "fluent_control", "1": "stutter_repetition", # part-word repetition "2": "stutter_prolongation", "3": "stutter_block", "5": "stutter_repetition", # word repetition "6": "stutter_repetition", # phrase repetition # 4 (interjection) and 7 (unmapped) intentionally omitted } # ---------------- column auto-detection (schemas differ) ------------------- def _col(ds, *cands): """Return a column whose (lowercased) name contains a candidate token.""" low = {c.lower(): c for c in ds.column_names} for c in cands: if c.lower() in low: return low[c.lower()] for c in ds.column_names: cl = c.lower() if any(tok in cl for tok in cands): return c return None def _find_audio(ds): """Locate the column carrying an Audio feature; else a descriptive name.""" for c in ds.column_names: kind = getattr(ds.features[c], "__class__", None) if kind is not None and kind.__name__ == "Audio": return c return _col(ds, "audio", "wav", "file", "path", "file_path") def _speaker_from_audio_path(p: str, fallback: str) -> str: """UCLASS names interleave speaker id + age, e.g. F_0101_10y4m_1_segment_0.""" m = re.match(r"^([A-Za-z0-9]+_\d+)", Path(p).name) return m.group(1).replace("_", "-") if m else fallback def _decode_audio_dict(a) -> Optional[np.ndarray]: """Decode a raw pyarrow audio struct {bytes, path} -> 16k float32 mono.""" import soundfile as sf import librosa if isinstance(a, dict): b = a.get("bytes") if b: raw, sr = sf.read(io.BytesIO(b), dtype="float32") else: p = a.get("path") if not p or not Path(p).exists(): return None raw, sr = sf.read(str(p), dtype="float32") elif isinstance(a, str) and Path(a).exists(): raw, sr = sf.read(str(a), dtype="float32") else: return None arr = raw.mean(axis=1) if raw.ndim > 1 else raw if sr != TARGET_SR: arr = librosa.resample(arr, orig_sr=sr, target_sr=TARGET_SR) return arr.astype("float32") def _tier_to_label(v: str, reg_key: str) -> str: """SEP-28k is a DETECTION-first corpus: every clip is `yes` (stutter) or `no` (fluent), not a disfluency-subtype tier. Map those to the canonical binary labels DIRECTLY — running them through the subtype heuristics would turn stuttered clips (block/prolongation) into `fluent_control`. """ v = v.lower().strip() if reg_key == "stutter_event": return "stutter" if v in ("yes", "true", "1", "stutter") else "fluent_control" if "rep" in v or "repetition" in v or "repeat" in v: return "stutter_repetition" if "par" in v or "prolongation" in v or "prolong" in v: return "stutter_prolongation" if "block" in v: return "stutter_block" return "fluent_control" def build(corpora: Optional[list] = None, seed: int = 42): """Load configured corpora into a single HF Dataset with speaker split.""" if corpora is None: corpora = list(CORPUS_LOAD) records = [] provenance = {"built_utc": None, "seed": seed, "corpora": {}} dropped = defaultdict(int) for reg_key in corpora: _, dsid, split = CORPUS_LOAD[reg_key] meta = CORPUS_REGISTRY[reg_key] cache = Path("data/corpora") / reg_key print(f"[read] {meta['name']} ({reg_key}) <- {dsid} [{split}]") ds = load_dataset(dsid, split=split, cache_dir=str(cache)) audio_col = _find_audio(ds) if audio_col is None: print(f" [!] no audio column for {reg_key}; skipping") continue text_col = _col(ds, "transcription", "transcript", "text", "prompt", "reference", "word_sequence", "utterance", "sentence") label_col = _col(ds, "label", "category", "disfluency", "tier", "disfluency_tier", "type", "class", "onset") speaker_col = _col(ds, "speaker_id", "speaker", "spk_id", "client_id", "speaker_idx", "name") # Read RAW pyarrow columns. Never ds[i] -> no torchcodec import. tab = ds.data audios = tab.column(audio_col).to_pylist() texts = (tab.column(text_col).to_pylist() if text_col else [""] * len(ds)) labels_raw = (tab.column(label_col).to_pylist() if label_col else ["fluent_control"] * len(ds)) speakers_raw = (tab.column(speaker_col).to_pylist() if speaker_col else ["unknown"] * len(ds)) n_ok = 0 for i in range(len(ds)): arr = None if audios[i] is not None: try: arr = _decode_audio_dict(audios[i]) except Exception: arr = None # corrupt/undecodable clip -> drop silently if arr is None: dropped[reg_key] += 1 continue if len(arr) == 0: dropped[reg_key] += 1 continue # label resolution per corpus lr = str(labels_raw[i]).lower() if reg_key == "uclass": label = UCLASS_CLASS_MAP.get(lr) if label is None: dropped[reg_key] += 1 # 4 (interjection) / 7 unknown continue elif reg_key in ("stutter_event", "libristutter"): label = _tier_to_label(lr, reg_key) else: label = lr if lr in LABEL_INDEX else "fluent_control" spk = str(speakers_raw[i]) if speaker_col else "unknown" if reg_key == "uclass" and spk == "unknown": spk = _speaker_from_audio_path(str(audios[i].get("path", "")), spk) # SEP-28k: carries no per-speaker id, only a per-clip `file` label. # Treat each distinct audio path stem (or its full embedding id) # as a separate "speaker" so SEP clips are never split randomly # across train/test — the anti-leak guard still holds. elif reg_key == "stutter_event" and spk == "unknown": aud = audios[i] if isinstance(audios[i], dict) else {} fname = str(aud.get("path") or "") spk = (Path(fname).stem if fname else f"sep:{i}") rec = { "id": f"{reg_key}:{i}", "corpus": reg_key, "speaker_id": spk, "audio_array": arr, "text": str(texts[i]), "label": label, } records.append(rec) n_ok += 1 print(f" [{reg_key}] {n_ok} rows kept out of {len(ds)}" f" ({dropped[reg_key]} dropped/undecodable)") provenance["corpora"][reg_key] = { "rows": n_ok, "total": len(ds), "dropped": dropped[reg_key], "audio_col": audio_col, "hf_id": dsid} del ds # Assemble HF Dataset. audio_array is a plain float32 column; the trainer # reads it directly (no Audio path embedding, which breaks on path=None). # 'split' is omitted here and added below via add_column (whole-speaker). no_split = {k: v for k, v in SCHEMA.items() if k != "split"} ds = Dataset.from_list(records, features=Features(no_split)) # ---------------- speaker-level held-out split ------------------------ split_col = splitter(ds, seed) ds = ds.add_column("split", split_col) counts = ds.to_pandas()["split"].value_counts().to_dict() print("\nSpeaker-held-out split (grouped by speaker):") for k in ("train", "val", "test"): n = counts.get(k, 0) if n > 0: subset = ds.filter(lambda r: r["split"] == k) lab = subset.to_pandas()["label"].value_counts().to_dict() print(f" {k:6} {n:8} rows labels={lab}") ds.save_to_disk(OUT_DIR / "dataset") provenance["n_records"] = len(ds) provenance["split_counts"] = counts (OUT_DIR / "dataset.json").write_text(json.dumps(provenance, indent=2)) print(f"\nDataset written to {OUT_DIR/'dataset'} (seed={seed}, {len(ds)} rows)") return ds def splitter(ds: Dataset, seed: int = 42): """Whole-speaker assignment to train/val/test. Returns split column.""" buckets = defaultdict(list) for i, s in enumerate(ds["speaker_id"]): buckets[s].append(i) rng = np.random.default_rng(seed) keys = list(buckets) rng.shuffle(keys) n = len(keys) tr = set(keys[: int(0.7 * n)]) va = set(keys[int(0.7 * n): int(0.85 * n)]) te = set(keys[int(0.85 * n):]) col = [] for s in ds["speaker_id"]: col.append("train" if s in tr else ("val" if s in va else "test")) return col if __name__ == "__main__": ap = argparse.ArgumentParser(description="Build unified HF dataset") ap.add_argument("--corpora", nargs="*", default=None, help="corpus keys; default all") ap.add_argument("--seed", type=int, default=42) args = ap.parse_args() build(corpora=args.corpora, seed=args.seed)