"""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())