ml/data/make_dataset.py

"""
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<bytes, path>`` 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)