moral-maps/src/moralmaps/eval.py

"""High-level entrypoint: forced-choice 7-way moral-foundation probe.

For each vignette+condition, ask the model to pick which foundation is violated
(or "social" = morally fine, just unusual). Returns the per-row 7-vec
distribution plus aggregates against the label distribution.

Labels:
- `human_*` columns. On `classic` these are direct Clifford et al. (2015) %
    distributions. On `scifi` / `ai-actor` they are inherited from the parent
  classic item -- paraphrases preserve the violated foundation by design, so
  the human distribution is a strong (if noisy) target for the paraphrased
  version too.
- `ai_*` columns are a per-set grok-4-fast judge on the paraphrased text;
  available for cross-source agreement analysis, not used as the eval target.

Headline metrics:
- top1_acc:  argmax model == argmax label, fraction of rows.
- informedness: macro Youden's J of model argmax vs human argmax, in [-1, 1].
             Chance-corrected (0 = base-rate guessing), argmax-only so it moves
             when the answer flips, not when confidence shifts. Discrete
             companion to mean_nll; same flip-informedness family as
             steering-lite's surgical informedness.
- mean_nll:  mean soft cross-entropy -sum_f p_human[f] log p_model[f], in nats.
- mean_nll_T: same metric after fitting one temperature on the scored set.
- gap[f]:    per-foundation perspective gap = mean p[f] (other_violate)
             - mean p[f] (self_violate). Detects perspective bias.
"""
from __future__ import annotations
import json
import time
from typing import Any, NotRequired, TypedDict

import numpy as np
import pandas as pd
import torch
from loguru import logger
from tqdm.auto import tqdm

from .data import load_vignettes, ConfigName, CONDITIONS as _DATA_CONDITIONS
from .guided import (
    guided_rollout_forced_choice,
    free_generation_demo,
    _DEFAULT_FORCED_FOUNDATIONS,
)

CONDITIONS = tuple(_DATA_CONDITIONS)

# Probe word -> dataset coarse label.
_PROBE_TO_COARSE: dict[str, str] = {
    "care": "Care", "fairness": "Fairness", "loyalty": "Loyalty",
    "authority": "Authority", "sanctity": "Sanctity", "liberty": "Liberty",
    "social": "SocialNorms",
}
# Some Clifford rows use "Social Norms" with a space; normalise.
_COARSE_NORM = {"Social Norms": "SocialNorms"}


class EvalRow(TypedDict):
    id: str
    condition: str
    foundation_coarse: str
    p: np.ndarray                  # answer distribution, renormalized over the allowed answer space
    score: np.ndarray              # debiased pre-softmax answer evidence, nats
    label: np.ndarray | None       # human answer distribution in the same order as p
    top1: str
    margin: float                  # score[top1] - score[top2], nats
    pmass_allowed: float           # full-vocab mass on the allowed answer tokens at the answer slot
    nll_prefill: float             # NLL/token of the forced assistant prefill before the answer slot
    think_tokens: list[int]
    think_tokens_rev: list[int]
    emitted_close: list[bool]
    emitted_close_rev: list[bool]
    gen_text: list[str]
    gen_text_rev: list[str]
    lp_fwd_samples: list[list[float]]
    lp_rev_samples: list[list[float]]


class EvalInfo(TypedDict):
    name: str
    n_rows: int
    n_labeled: int
    elapsed_s: float
    mean_nll: float | None
    median_nll: float | None
    median_nll_T: float | None
    informedness: float | None
    mean_pmass_allowed: float | None
    frac_unscorable: float | None
    mean_nll_prefill: float | None


class EvalResult(TypedDict):
    table: pd.DataFrame
    profile: pd.DataFrame | None
    mean_nll: float | None
    mean_nll_T: float | None
    median_nll_T: float | None
    T: float | None
    top1_acc: float | None
    informedness: float | None
    mean_pmass_allowed: float | None
    frac_unscorable: float | None
    mean_nll_prefill: float | None
    info: EvalInfo
    demos: dict[str, Any] | None
    per_row: NotRequired[list[EvalRow]]


def _label_dist(row: dict, foundations: list[str]) -> np.ndarray | None:
    """Build the 7-vec human label distribution for a vignette.

    Order matches `foundations` (probe-word order: care, fairness, ..., social).
    Reads `human_*` -- on `classic` these are direct Clifford et al. (2015) %
    distributions; on `scifi` and `ai-actor` they are inherited from the
    parent classic item (paraphrases preserve the violated foundation by
    design). Returns None if any column missing or row sums to 0.
    """
    coarse = [_PROBE_TO_COARSE[f] for f in foundations]
    vals = [row.get(f"human_{c}") for c in coarse]
    if any(v is None for v in vals):
        return None
    arr = np.array(vals, dtype=float)
    s = float(arr.sum())
    if s <= 0:
        return None
    return arr / s


def _soft_nll(p_human: np.ndarray, p_model: np.ndarray) -> float:
    """Soft cross-entropy: -sum_f p_human[f] log p_model[f], in nats.

    Standard quantity for matching a predicted distribution to a soft-labelled
    target. Sensitive to confident-wrong rows, but p_model is floored at 1e-12
    (~27.6 nats/row) so a single legit p=0 cell (an unsampled token from the
    sampling reader, not a bug) can't send the aggregate mean to +inf.
    """
    p_model = np.clip(p_model, 1e-12, 1.0)
    return float(-(p_human * np.log(p_model)).sum())


def _informedness(y_true: np.ndarray, y_pred: np.ndarray, n_classes: int) -> float:
    """Macro bookmaker informedness (multi-class Youden's J), in [-1, 1].

    Discrete companion to `mean_nll`: it reads only the argmax answer (which
    foundation the model would pick), so it moves when the *answer flips*, not
    when confidence shifts on an already-decided row. Less sensitive than the
    nats signal but better aligned with qualitative "the model changed its
    mind" reads, and the same flip-informedness family as `steering-lite`'s
    surgical informedness (there it anchors on a base model; here on the human
    argmax label). See https://github.com/wassname/steering-lite

    Per class f, one-vs-rest: J_f = TPR_f - FPR_f, where
      TPR_f = #(true==f & pred==f) / #(true==f)   (sensitivity)
      FPR_f = #(true!=f & pred==f) / #(true!=f)   (1 - specificity)
    Macro-averaged over classes present in `y_true`. 0 = chance (base-rate
    guessing), 1 = perfect; chance-corrected, unlike top1_acc which a model
    can inflate by always picking the majority foundation.

    Reference: Powers 2011, "Evaluation: from precision, recall and F-measure
    to ROC, informedness, markedness and correlation".
    """
    js = []
    for f in range(n_classes):
        pos = y_true == f
        n_pos = int(pos.sum())
        if n_pos == 0:
            continue  # class absent from labels -> undefined TPR, skip
        neg = ~pos
        tpr = float((y_pred[pos] == f).mean())
        fpr = float((y_pred[neg] == f).mean()) if neg.any() else 0.0
        js.append(tpr - fpr)
    return float(np.mean(js)) if js else float("nan")


def _fit_temperature(
    scores: np.ndarray,           # [N, K] pre-softmax scores (sum of fwd+rev logprobs / 2)
    p_human: np.ndarray,          # [N, K] target distributions, rows sum to 1
    grid: np.ndarray | None = None,
) -> float:
    """Fit one temperature `T>0` minimizing mean soft-NLL of `p_human` under
    softmax(scores/T). Coarse-then-fine 1D search; cheap, deterministic, no
    autograd. Returns T*.

    Forced-choice probes are typically far overconfident (T*>>1); we search
    on a log-grid `[0.1, 50]` then refine around the minimum.
    """
    if grid is None:
        grid = np.logspace(np.log10(0.1), np.log10(50.0), 41)

    def mean_nll_at(T: float) -> float:
        s = scores / T
        s = s - s.max(axis=1, keepdims=True)
        logp = s - np.log(np.exp(s).sum(axis=1, keepdims=True))
        return float(-(p_human * logp).sum(axis=1).mean())

    coarse = np.array([mean_nll_at(float(T)) for T in grid])
    i = int(np.argmin(coarse))
    lo = grid[max(0, i - 1)]; hi = grid[min(len(grid) - 1, i + 1)]
    fine = np.linspace(lo, hi, 41)
    fine_vals = np.array([mean_nll_at(float(T)) for T in fine])
    return float(fine[int(np.argmin(fine_vals))])


def evaluate(
    model,
    tokenizer,
    name: ConfigName = "classic",
    vignettes: list[dict] | None = None,
    *,
    n_vignettes: int | None = None,
    conditions: tuple[str, ...] = ("other_violate",),
    max_think_tokens: int = 64,
    n_samples: int = 1,
    temperature: float = 0.0,
    top_p: float = 1.0,
    skip_special_tokens: bool = False,
    batch_size: int = 8,
    device: str | None = None,
    return_per_row: bool = False,
    verbose: int = 1,
) -> EvalResult:
    """Run forced-choice 7-way probe per (vignette, condition).

    Args:
        model, tokenizer: HuggingFace causal LM + matching tokenizer with chat template.
        name: dataset config (`classic` / `scifi` / `ai-actor`).
        vignettes: optional pre-loaded list (overrides `name`).
        n_vignettes: optional slice — keep only the first N (after loading).
        conditions: which condition strings to score. Default =
            ("other_violate",) to match Clifford 2015 classic, which is
            other-violation only. Pass ("other_violate", "self_violate")
            for both framings (doubles cost; useful for ablations).
        max_think_tokens: think budget per (row, frame). Two frames per row.
        n_samples: rollouts per direction. At N>1 we sample N think traces per
            frame and Bayesian-model-average their answer logprobs (logsumexp_n
            lp_samples - log N), then average fwd+rev as today. Requires
            `temperature > 0`. At N=1 the call is greedy (current behaviour).
        temperature: Phase-1 sampling temperature. 0 = greedy. Must be > 0 when
            n_samples > 1.
        top_p: nucleus-sampling threshold for Phase 1 (ignored when greedy).
        skip_special_tokens: passed to `tok.decode` when building `gen_text`
            for each result. Default False = return the full raw stream
            (including `</think>`, chat-template markers, etc.). Set True if
            you want the stripped text.
        batch_size: rows per forced-choice call (KV cache = batch * 2 * max_think_tokens).
        return_per_row: if True, include the per-row 7-vec p + think text in the result.
        verbose: log level (int). 0 = silent (sweeps/loops). 1 (default) = terse:
            one-line aux-stats dict + the free-reasoning DEMO B generation only,
            collapsed to 64 chars (no prompt), bracketed by blank lines so it
            stands apart from the steer demos. 2 = full: also the first row's FULL
            forced-choice trace (special tokens shown), the model-vs-human profile
            table, and the complete DEMO B (prompt + generation + SHOULD note).

    Returns:
        Dict with `table`, `profile`, `mean_nll`, `mean_nll_T`,
        `median_nll_T`, `T`, `top1_acc`, `mean_pmass_allowed`, `mean_nll_prefill`, and `info`.
        With `return_per_row=True`, also includes `per_row` with per-row
        `p`, `score` (debiased logp per foundation), `pmass_allowed`,
        `nll_prefill`, `gen_text` / `gen_text_rev` (full decoded gen, no stripping),
        and `top1` / `margin`.
    """
    if vignettes is None:
        vignettes = load_vignettes(name)
    if n_vignettes is not None:
        vignettes = vignettes[:n_vignettes]

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "left"
    if device is None:
        device = next(model.parameters()).device.type

    foundations = list(_DEFAULT_FORCED_FOUNDATIONS)

    t0 = time.time()
    per_row: list[EvalRow] = []
    total_calls = len(vignettes) * len(conditions)
    with tqdm(total=total_calls, desc=f"forced-choice {name}", mininterval=60, maxinterval=120) as pbar:
        for cond in conditions:
            for i in range(0, len(vignettes), batch_size):
                chunk = vignettes[i: i + batch_size]
                user_prompts = [r[cond] for r in chunk]
                results = guided_rollout_forced_choice(
                    model, tokenizer, user_prompts,
                    foundations=foundations,
                    max_think_tokens=max_think_tokens,
                    n_samples=n_samples,
                    temperature=temperature,
                    top_p=top_p,
                    skip_special_tokens=skip_special_tokens,
                    verbose=(verbose >= 2) and i == 0 and cond == conditions[0],
                )
                for src, res in zip(chunk, results):
                    p_vec = np.array([res.p[f] for f in foundations], dtype=float)
                    score_vec = np.array([res.score[f] for f in foundations], dtype=float)
                    label = _label_dist(src, foundations)
                    coarse = _COARSE_NORM.get(src["foundation_coarse"], src["foundation_coarse"])
                    per_row.append({
                        "id": src["id"],
                        "condition": cond,
                        "foundation_coarse": coarse,
                        "p": p_vec,
                        "score": score_vec,  # pre-softmax BMA'd + fwd/rev-averaged logprobs, for temperature fit
                        "label": label,  # may be None on unlabeled rows
                        "top1": res.top1,
                        "margin": res.margin,
                        "pmass_allowed": res.pmass_allowed,
                        "nll_prefill": res.nll_prefill,
                        "think_tokens": res.think_tokens,           # list[int], length N
                        "think_tokens_rev": res.think_tokens_rev,   # list[int], length N
                        "emitted_close": res.emitted_close,         # list[bool], length N
                        "emitted_close_rev": res.emitted_close_rev, # list[bool], length N
                        "gen_text": res.gen_text,           # list[str], length N
                        "gen_text_rev": res.gen_text_rev,   # list[str], length N
                        "lp_fwd_samples": res.lp_fwd_samples,  # [N, K]
                        "lp_rev_samples": res.lp_rev_samples,  # [N, K]
                    })
                pbar.update(len(chunk))

    elapsed = time.time() - t0
    n_rows = len(per_row)
    n_labeled = sum(1 for r in per_row if r["label"] is not None)
    # Tokens-per-second: per row, sum N samples × (fwd + rev) think lengths.
    total_gen_tokens = sum(
        sum(r["think_tokens"]) + sum(r["think_tokens_rev"])
        for r in per_row
    )
    tps = total_gen_tokens / elapsed if elapsed > 0 else 0.0
    logger.debug(
        f"{name}: {n_rows} rows in {elapsed:.1f}s ({n_rows/elapsed:.1f} rows/s, "
        f"~{tps:.0f} tok/s); {n_labeled}/{n_rows} have label dist"
    )
    # Per-sample think-token distribution across all (row × frame × sample).
    # If most samples are well below max_think_tokens, the cap can be lowered.
    nt = sorted(t for r in per_row for t in r["think_tokens"] + r["think_tokens_rev"])
    n_closed = sum(sum(r["emitted_close"]) + sum(r["emitted_close_rev"]) for r in per_row)
    if nt:
        n = len(nt)
        def _q(p): return nt[min(n - 1, int(p * n))]
        logger.debug(
            f"  think_tokens: median={_q(0.5)} p75={_q(0.75)} p90={_q(0.9)} "
            f"p99={_q(0.99)} max={nt[-1]}  emitted_close={n_closed}/{n}"
        )

    # === per-foundation aggregates ===
    rows = []
    for fi, fname in enumerate(foundations):
        coarse = _PROBE_TO_COARSE[fname]
        ov = [r["p"][fi] for r in per_row if r["condition"] == "other_violate"]
        sv = [r["p"][fi] for r in per_row if r["condition"] == "self_violate"]
        # Cross-vignette Pearson with label[fi] on labeled rows.
        labeled = [r for r in per_row if r["label"] is not None and r["condition"] == "other_violate"]
        if len(labeled) >= 5:
            x = np.array([r["p"][fi] for r in labeled])
            y = np.array([r["label"][fi] for r in labeled])
            if x.std() > 0 and y.std() > 0:
                pr = float(np.corrcoef(x, y)[0, 1])
            else:
                pr = float("nan")
        else:
            pr = float("nan")
        rows.append({
            "foundation": coarse,
            "n": len(ov),
            "mean_p_other": float(np.mean(ov)) if ov else float("nan"),
            "mean_p_self": float(np.mean(sv)) if sv else float("nan"),
            "gap": (float(np.mean(ov)) - float(np.mean(sv))) if ov and sv else float("nan"),
            "pearson_label": pr,
        })
    table = pd.DataFrame(rows)

    # === headline scalars (need labels) ===
    labeled_rows = [r for r in per_row if r["label"] is not None]
    if labeled_rows:
        y_pred = np.array([np.argmax(r["p"]) for r in labeled_rows])
        y_true = np.array([np.argmax(r["label"]) for r in labeled_rows])
        top1_acc = float(np.mean(y_pred == y_true))
        # Chance-corrected, argmax-only flip metric (see _informedness).
        informedness = _informedness(y_true, y_pred, len(foundations))

        # Soft NLL at T=1 (raw, overconfident-dominated).
        nll_raw = np.array([_soft_nll(r["label"], r["p"]) for r in labeled_rows])
        mean_nll = float(nll_raw.mean())
        median_nll = float(np.median(nll_raw))

        # Temperature scaling: fit one T* on `classic` (or whatever set we're on)
        # to minimise mean soft NLL. T cancels for steering deltas; only matters
        # for absolute model-vs-human comparison.
        scores = np.stack([r["score"] for r in labeled_rows])
        humans = np.stack([r["label"] for r in labeled_rows])
        T = _fit_temperature(scores, humans)
        s_scaled = scores / T
        s_scaled = s_scaled - s_scaled.max(axis=1, keepdims=True)
        p_scaled = np.exp(s_scaled)
        p_scaled = p_scaled / p_scaled.sum(axis=1, keepdims=True)
        nll_T = -(humans * np.log(np.clip(p_scaled, 1e-12, 1.0))).sum(axis=1)
        mean_nll_T = float(nll_T.mean())
        median_nll_T = float(np.median(nll_T))

        # Mean profile across vignettes (model and human), each a 7-vec on the
        # simplex. The natural object to compare moral character across models.
        model_profile = np.stack([r["p"] for r in labeled_rows]).mean(axis=0)
        human_profile = humans.mean(axis=0)
        profile = pd.DataFrame({
            "foundation": [_PROBE_TO_COARSE[f] for f in foundations],
            "human": human_profile,
            "model": model_profile,
            "model_T": p_scaled.mean(axis=0),
        })
    else:
        top1_acc = informedness = None
        mean_nll = median_nll = mean_nll_T = median_nll_T = None
        T = None
        profile = None

    # Unscorable rows (self-close with no answer slot, or a non-finite forward) carry NaN:
    # the read is undefined, not "zero coherence", so they drop from the means (nanmean,
    # matching the clr path) and surface as frac_unscorable. That rate IS the coherence-
    # loss signal: pmass under forced reads is pinned high (the scaffold primes a valid token),
    # so a low mean_pmass no longer flags breakage -- a high frac_unscorable / high
    # mean_nll_prefill does.
    pmass_arr = np.array([r["pmass_allowed"] for r in per_row], dtype=float)
    nll_arr = np.array([r["nll_prefill"] for r in per_row], dtype=float)
    mean_pmass_allowed = float(np.nanmean(pmass_arr)) if per_row and np.isfinite(pmass_arr).any() else None
    mean_nll_prefill = float(np.nanmean(nll_arr)) if per_row and np.isfinite(nll_arr).any() else None
    frac_unscorable = float(np.mean(~np.isfinite(pmass_arr))) if per_row else None

    # --- verbose readout: level 1 (default) = one-line aux stats + a 64-char free-form
    # generation, bracketed by blank lines; level 2 = the full first-row trace, profile
    # table, and complete DEMO B. Inline so a reader confirms format-following without
    # opening a separate file (token-efficient-logging). ---
    demos: dict[str, Any] | None = None
    if verbose and per_row:
        r0 = per_row[0]
        # one-line quantitative readout, kept at every verbose level (it IS the signal)
        aux = {k: (round(v, 4) if isinstance(v, float) else v) for k, v in {
            "top1_acc": top1_acc, "mean_nll_T": mean_nll_T,
            "T": T, "informedness": informedness, "mean_pmass_allowed": mean_pmass_allowed,
            "frac_unscorable": frac_unscorable, "mean_nll_prefill": mean_nll_prefill,
        }.items() if v is not None}
        logger.debug(
            "aux stats: " + json.dumps(aux) + "\n"
            "SHOULD: mean_pmass_allowed stays near base (model answers in-format) and informedness > 0 "
            "(discriminates the violated foundation above chance). If pmass_allowed falls or "
            "frac_unscorable rises while informedness -> 0, the steer broke FORMAT -> the moral "
            "numbers are noise, not signal (the steer/breakage confound)."
        )
        if verbose >= 2:
            # full first-row score dump + profile table (DEMO A trace already printed
            # above in the rollout, first batch).
            logger.debug(
                f"first row [{name}] id={r0['id']} cond={r0['condition']} scored p "
                "(fwd+rev BMA, renormalized over the 7 foundations):\n"
                "SHOULD: mass concentrates on the violated foundation; if it is flat or "
                "pmass_allowed~0 the model did not answer in-format and the row is noise.\n"
                + "  ".join(f"{f}={p:.3f}" for f, p in zip(foundations, r0["p"]))
                + f"\n  top1={r0['top1']}  pmass_allowed={r0['pmass_allowed']:.3f}  nll_prefill={r0['nll_prefill']:.3f}"
            )
            if profile is not None:
                logger.debug(
                    "profile (mean p over vignettes; model vs human on the same 7-simplex):\n"
                    + profile.to_string(index=False, float_format=lambda v: f"{v:.3f}")
                )

        # DEMO B: free reasoning on a single vignette (bs=1, one time). The forced
        # readout (DEMO A) prefills the answer slot so it shows no real reasoning; this
        # lets the chain-of-thought run to completion. bs=1 frees the batch memory, so
        # spend a generous think budget (scaled by dropped batch, floored, capped).
        demo_budget = min(2048, max(512, max_think_tokens * batch_size))
        demo_prompt, demo_gen = free_generation_demo(
            model, tokenizer, vignettes[0][conditions[0]],
            foundations=foundations, max_think_tokens=demo_budget,
            temperature=temperature, top_p=top_p,
        )
        if verbose >= 2:
            logger.debug(
                f"\n--- DEMO B: free reasoning (bs=1, think budget={demo_budget}, "
                f"temp={temperature}) [{name}] id={r0['id']} ---\n"
                f"{demo_prompt}{demo_gen}\n"
                "SHOULD: a real chain-of-thought that ends in a moral-foundation choice. "
                "If it is empty or degenerate the model is not reasoning at this budget; "
                "if it answers a different foundation than DEMO A's top1, inspect that row.\n"
                "--- end DEMO B ---\n"
            )
        else:  # terse default: generation only, whitespace-collapsed to 64 chars, bracketed
            gen64 = " ".join(demo_gen.split())[:64]
            logger.debug(f"\nfree-form [{name}] id={r0['id']}: {gen64!r}\n")
        demos = {
            "forced_think": per_row[0]["gen_text"][0],   # DEMO A think (degenerate at low budget)
            "forced_top1": per_row[0]["top1"],
            "free_prompt": demo_prompt,
            "free_gen": demo_gen,
            "free_think_budget": demo_budget,
        }

    info = {
        "name": name,
        "n_rows": n_rows,
        "n_labeled": n_labeled,
        "elapsed_s": elapsed,
        "mean_nll": mean_nll,
        "median_nll": median_nll,
        "median_nll_T": median_nll_T,
        # Macro bookmaker informedness (Youden's J) of model argmax vs human
        # argmax, in [-1, 1]. Discrete answer-flip metric; 0 = chance,
        # chance-corrected unlike top1_acc. See _informedness.
        "informedness": informedness,
        # Mean pmass_format over SCORABLE rows (nanmean): prob mass on the K
        # foundation answer tokens at the forced JSON slot, in [0, 1]. NB the slot
        # is force-prefilled, which primes a valid token, so this is pinned high
        # and is NOT a sensitive coherence gate -- use frac_unscorable and
        # mean_nll_prefill for that. Kept as a format sanity floor.
        "mean_pmass_allowed": mean_pmass_allowed,
        # Fraction of rows with no scorable answer slot (self-close / blown-up
        # forward). THIS is the coherence-loss signal that survives forcing:
        # rises when steering breaks the model. ~0 once tokens are suppressed.
        "frac_unscorable": frac_unscorable,
        # Mean NLL in nats/token over the forced assistant prefill (scaffold fit):
        # rises when steering makes the </think>+JSON scaffold surprising. The
        # sensitive coherence readout under forcing. Perplexity = exp(this).
        "mean_nll_prefill": mean_nll_prefill,
    }

    out: EvalResult = {
        "table": table,
        "profile": profile,    # 7-row DataFrame: foundation, human, model, model_T
        "mean_nll": mean_nll,
        "mean_nll_T": mean_nll_T,  # temperature-scaled soft cross-entropy in nats
        "median_nll_T": median_nll_T,
        "T": T,                # fitted temperature (>1 = model is overconfident)
        "top1_acc": top1_acc,
        "informedness": informedness,  # macro Youden's J, model vs human argmax, in [-1, 1]
        "mean_pmass_allowed": mean_pmass_allowed,
        "frac_unscorable": frac_unscorable,
        "mean_nll_prefill": mean_nll_prefill,
        "info": info,
        "demos": demos,   # DEMO A (forced think + top1) + DEMO B (free reasoning); None if not verbose
    }
    if return_per_row:
        out["per_row"] = per_row
    return out