spb/localvm-research Public License

1# =============================================================================2#  Project   : localvm-research3#  File      : src/localvm/quality/decision_stats.py4#  Purpose   : Shared decision-stability measurement utilities (greedy5#              trajectories, teacher-forced margins/agreement, AUROC, curves)6#  Author    : Simon-Pierre Boucher7#  Contact   : contact@spboucher.ai8#  Created   : 2026-08-129#  Modified  : 2026-08-1210#  Platform  : macOS / Apple Silicon (arm64) — MLX / Metal11#  License   : All rights reserved (research code)12# =============================================================================13"""Decision-stability measurement utilities shared by expG/expD and successors.1415First used (inlined) by experiments/micro/expG_decision_stability/benchmark.py16(commit 42c7b3d); extracted here unchanged so later experiments reuse one17implementation. expG's committed copy is kept as-is for reproducibility.18"""1920from __future__ import annotations2122import mlx.core as mx23import numpy as np242526def greedy_generate(model, tokenizer, prompt_ids: list[int], n_tokens: int) -> list[int]:27    """Deterministic greedy generation; returns generated token ids."""28    from mlx_lm.models.cache import make_prompt_cache2930    cache = make_prompt_cache(model)31    generated: list[int] = []32    inp = mx.array(list(prompt_ids))[None]33    for _ in range(n_tokens):34        logits = model(inp, cache=cache)35        nxt = int(mx.argmax(logits[0, -1]).item())36        if nxt == tokenizer.eos_token_id:37            break38        generated.append(nxt)39        inp = mx.array([[nxt]])40    return generated414243def teacher_forced_stats(model, full_ids: list[int], start: int) -> dict:44    """Forward full_ids once; per-position stats for predictions of tokens45    [start, len). Returns margin (top-1 minus top-2 logit gap), argmax, and46    float16 logprobs (float16 keeps 48 x (128, ~152k) around 2 GB)."""47    logits = model(mx.array(full_ids)[None])[0]48    sel = logits[start - 1 : len(full_ids) - 1].astype(mx.float32)49    top2 = mx.topk(sel, 2, axis=-1)50    argmax = mx.argmax(sel, axis=-1)51    logprobs = sel - mx.logsumexp(sel, axis=-1, keepdims=True)52    mx.eval(top2, argmax, logprobs)53    v = np.array(top2)54    return {55        "margin": np.abs(v[:, 1] - v[:, 0]),56        "argmax": np.array(argmax),57        "logprobs": np.array(logprobs).astype(np.float16),58    }596061def auroc(scores: np.ndarray, labels: np.ndarray) -> float:62    """Rank-based AUROC with tie handling (scores: higher = predicted 1)."""63    pos, neg = scores[labels == 1], scores[labels == 0]64    if len(pos) == 0 or len(neg) == 0:65        return float("nan")66    allv = np.concatenate([pos, neg])67    order = np.argsort(allv, kind="mergesort")68    ranks = np.empty(len(order))69    ranks[order] = np.arange(1, len(order) + 1)70    sorted_v = allv[order]71    i = 072    while i < len(sorted_v):73        j = i74        while j + 1 < len(sorted_v) and sorted_v[j + 1] == sorted_v[i]:75            j += 176        if j > i:77            ranks[order[i : j + 1]] = ranks[order[i : j + 1]].mean()78        i = j + 179    r_pos = ranks[: len(pos)].sum()80    return float((r_pos - len(pos) * (len(pos) + 1) / 2) / (len(pos) * len(neg)))818283def escalation_curve(margins: np.ndarray, agree: np.ndarray, points: int = 200) -> list[dict]:84    """Escalate tokens with margin < tau (escalated decision assumed exact);85    report escalated fraction vs residual disagreement."""86    qs = np.quantile(margins, np.linspace(0, 1, points))87    out, n = [], len(margins)88    for tau in qs:89        esc = margins < tau90        out.append({91            "tau": float(tau),92            "escalated_frac": float(esc.mean()),93            "residual_disagree": float(np.sum((~esc) & (agree == 0)) / n),94        })95    return out969798def kl_ref_vs(model_logprobs_f16: np.ndarray, ref_logprobs_f16: np.ndarray) -> np.ndarray:99    """Per-position KL(ref || model), computed in float32."""100    ref = ref_logprobs_f16.astype(np.float32)101    q = model_logprobs_f16.astype(np.float32)102    return np.sum(np.exp(ref) * (ref - q), axis=-1)103