scripts/run_motivation_f2.py

"""Phase 1 motivation: consume grounding_g_t.parquet from compute_grounding.py and produce:

  F2a  : g_t histogram (overall + per data_source)            -- is the distribution bimodal?
  F2a2 : g_t by token position bucket (early / mid / late)    -- does grounding drift along
                                                                 the response?
  F2b  : low-g_t bottom-20% tokens jsonl for GPT-4o judge     -- are they really hallucinated?
  F2b+ : high-g_t top-20%   tokens jsonl                      -- sanity check on grounded side
  F2c  : alpha-sweep target logit shift                       -- visualize CD-OPD effect at
                                                                 different alpha levels

All figures saved as PNG; jsonl tables for GPT-4o judging saved alongside. CPU only.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import pandas as pd

ROOT = Path("/mnt/local-fast/opd_zt")


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--input", default=str(ROOT / "data" / "grounding_g_t.parquet"))
    p.add_argument("--outdir", default=str(ROOT / "outputs" / "motivation"))
    p.add_argument("--judge_n_per_bucket", type=int, default=200)
    p.add_argument("--context_window", type=int, default=10)
    return p.parse_args()


def fig_g_t_hist(df: pd.DataFrame, outdir: Path) -> None:
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    g = df["g_t"].to_numpy()
    # Clip to a sensible range for plotting.
    g_clip = np.clip(g, np.quantile(g, 0.001), np.quantile(g, 0.999))

    fig, axes = plt.subplots(1, 2, figsize=(12, 4))
    axes[0].hist(g_clip, bins=80, color="#444", alpha=0.85)
    axes[0].axvline(0, color="r", linewidth=1, linestyle="--")
    axes[0].set_xlabel("g_t = log p_T(a|v) - log p_T(a|v⁻)")
    axes[0].set_ylabel("token count")
    axes[0].set_title("F2a: per-token grounding score")

    for ds, sub in df.groupby("data_source"):
        gd = np.clip(sub["g_t"].to_numpy(),
                     np.quantile(g, 0.001), np.quantile(g, 0.999))
        axes[1].hist(gd, bins=60, alpha=0.55, label=f"{ds} (n={len(sub)})", density=True)
    axes[1].axvline(0, color="r", linewidth=1, linestyle="--")
    axes[1].set_xlabel("g_t")
    axes[1].set_ylabel("density")
    axes[1].set_title("by data_source")
    axes[1].legend(fontsize=8)

    fig.tight_layout()
    out = outdir / "F2a_g_t_hist.png"
    fig.savefig(out, dpi=140)
    plt.close(fig)
    print(f"[F2a] wrote {out}")


def fig_g_t_by_position(df: pd.DataFrame, outdir: Path) -> None:
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    # Bucket each token by relative position in its response.
    pos_frac = df["token_pos"].to_numpy() / np.maximum(df["response_len"].to_numpy(), 1)
    buckets = pd.cut(pos_frac, bins=[-0.001, 0.33, 0.66, 1.001], labels=["early", "mid", "late"])
    df2 = df.assign(bucket=buckets)

    fig, ax = plt.subplots(figsize=(7, 4))
    data = [df2[df2["bucket"] == b]["g_t"].to_numpy() for b in ["early", "mid", "late"]]
    bp = ax.boxplot(
        data,
        labels=["early (0-33%)", "mid (33-66%)", "late (66-100%)"],
        showfliers=False,
        patch_artist=True,
    )
    for patch, color in zip(bp["boxes"], ["#5b9bd5", "#ed7d31", "#a5a5a5"], strict=False):
        patch.set_facecolor(color)
    ax.axhline(0, color="r", linewidth=1, linestyle="--")
    ax.set_ylabel("g_t")
    ax.set_title("F2a2: grounding by response-position bucket")
    fig.tight_layout()
    out = outdir / "F2a2_g_t_by_position.png"
    fig.savefig(out, dpi=140)
    plt.close(fig)
    print(f"[F2a2] wrote {out}")


def dump_judge_jsonl(
    df: pd.DataFrame, outdir: Path, n_per_bucket: int, context_window: int
) -> None:
    """Pick top/bottom 20% by g_t, sample n_per_bucket from each, build (token + context)
    rows for an LLM judge to label as hallucinated/grounded."""
    quantile_lo = df["g_t"].quantile(0.20)
    quantile_hi = df["g_t"].quantile(0.80)

    low = df[df["g_t"] <= quantile_lo].sample(
        n=min(n_per_bucket, len(df[df["g_t"] <= quantile_lo])), random_state=0
    )
    high = df[df["g_t"] >= quantile_hi].sample(
        n=min(n_per_bucket, len(df[df["g_t"] >= quantile_hi])), random_state=1
    )

    def build_record(row: pd.Series, group: str) -> dict:
        # Get response context window around the target token. Recover by grouping the
        # parent df by sample_idx and slicing token_pos.
        same = df[df["sample_idx"] == row["sample_idx"]].sort_values("token_pos")
        center = int(row["token_pos"])
        lo, hi = max(0, center - context_window), min(len(same), center + context_window + 1)
        ctx = same.iloc[lo:hi]
        before = "".join(ctx[ctx["token_pos"] < center]["token_text"].tolist())
        target = str(row["token_text"])
        after = "".join(ctx[ctx["token_pos"] > center]["token_text"].tolist())
        return {
            "sample_idx": int(row["sample_idx"]),
            "token_pos": center,
            "g_t": float(row["g_t"]),
            "group": group,
            "context_before": before,
            "target_token": target,
            "context_after": after,
            "data_source": row["data_source"],
        }

    records = [build_record(r, "low_g_t") for _, r in low.iterrows()]
    records += [build_record(r, "high_g_t") for _, r in high.iterrows()]
    out = outdir / "F2b_judge_input.jsonl"
    with out.open("w") as f:
        for rec in records:
            f.write(json.dumps(rec) + "\n")
    print(f"[F2b] wrote {len(records)} records -> {out}")
    print(f"       low g_t cutoff: <= {quantile_lo:.3f}    high g_t cutoff: >= {quantile_hi:.3f}")


def fig_alpha_sweep(df: pd.DataFrame, outdir: Path) -> None:
    """Visualise the CD-OPD target logit shift  (1+alpha)*logp_v - alpha*logp_v_perturbed -
    logp_v   versus  g_t,  for several alpha levels.  Shows how CD-OPD damps low-g_t tokens
    and boosts high-g_t tokens."""
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    g = df["g_t"].to_numpy()

    fig, ax = plt.subplots(figsize=(7, 4))
    for alpha in (0.0, 0.5, 1.0, 1.5, 2.0):
        shift = alpha * (-g)  # since (1+alpha)*lv - alpha*lvp - lv = alpha*(lv - lvp - 0) but
        # we want (1+alpha)*lv - alpha*lvp - lv = alpha*(lv - lvp) = alpha * g_t.
        # The correction wrt vanilla target is +alpha*g_t.  So plot that.
        shift = alpha * g
        # x: g_t  y: target logit shift relative to vanilla OPD (alpha=0)
        order = np.argsort(g)
        ax.plot(g[order], shift[order], label=f"α={alpha}", linewidth=1.3)
    ax.axhline(0, color="k", linewidth=0.6)
    ax.axvline(0, color="r", linewidth=0.6, linestyle="--")
    ax.set_xlabel("g_t")
    ax.set_ylabel("CD-OPD logit shift vs vanilla   (α · g_t)")
    ax.set_title("F2c: CD-OPD target shift as a function of α")
    ax.legend()
    fig.tight_layout()
    out = outdir / "F2c_alpha_sweep.png"
    fig.savefig(out, dpi=140)
    plt.close(fig)
    print(f"[F2c] wrote {out}")


def summary_stats(df: pd.DataFrame, outdir: Path) -> None:
    g = df["g_t"]
    report = {
        "n_tokens": int(len(df)),
        "n_samples": int(df["sample_idx"].nunique()),
        "g_t_mean": float(g.mean()),
        "g_t_std": float(g.std()),
        "g_t_p01": float(g.quantile(0.01)),
        "g_t_p10": float(g.quantile(0.10)),
        "g_t_p25": float(g.quantile(0.25)),
        "g_t_p50": float(g.quantile(0.50)),
        "g_t_p75": float(g.quantile(0.75)),
        "g_t_p90": float(g.quantile(0.90)),
        "g_t_p99": float(g.quantile(0.99)),
        "frac_g_t_le_0": float((g <= 0).mean()),
        "frac_g_t_gt_2": float((g > 2.0).mean()),
        "per_data_source": {
            ds: {
                "n": int(len(sub)),
                "mean": float(sub["g_t"].mean()),
                "p25": float(sub["g_t"].quantile(0.25)),
                "p75": float(sub["g_t"].quantile(0.75)),
            }
            for ds, sub in df.groupby("data_source")
        },
    }
    out = outdir / "summary.json"
    out.write_text(json.dumps(report, indent=2))
    print(f"[summary] wrote {out}")
    print(json.dumps(report, indent=2))


def main() -> int:
    args = parse_args()
    outdir = Path(args.outdir)
    outdir.mkdir(parents=True, exist_ok=True)
    df = pd.read_parquet(args.input)
    print(f"[load] {len(df)} token rows from {args.input}")

    summary_stats(df, outdir)
    fig_g_t_hist(df, outdir)
    fig_g_t_by_position(df, outdir)
    dump_judge_jsonl(df, outdir, args.judge_n_per_bucket, args.context_window)
    fig_alpha_sweep(df, outdir)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())