"""Gradio app: In-Context Elicitation of J-Space Contents. Two tabs (§7): 1. **Single prompt** — pick an eval prompt and ``n``, run the report path, and see the model's ranked concepts beside the J-lens ground truth, with scores. 2. **Scaling curve** — run the grid over ``n in {0,1,3,10}`` + format-only across the eval set and plot overlap-vs-n with the chance-floor baseline. Both GPU paths (ground truth, report) run on ZeroGPU inside the ``@spaces.GPU`` functions in ``jspace.groundtruth`` / ``jspace.report``. This module is pure UI + orchestration and imports cleanly on CPU. """ from __future__ import annotations import gradio as gr import pandas as pd from jspace import experiment as exp from jspace import groundtruth as gt # noqa: F401 (kept for symmetry / API parity) from jspace import model as model_mod from jspace import report as rp from jspace.dataset import make_split from jspace.experiment import CONDITIONS, RunConfig # noqa: F401 # On the Space, load gemma + the lens once at startup (outside any GPU time-box) # so the @spaces.GPU paths just reuse the cached bundle. No-op on CPU/local. model_mod.preload() SPLIT = make_split() EVAL_BY_LABEL = {f"{it.name} — {it.prompt[:60]}…": it for it in SPLIT.eval} N_CHOICES = [0, 1, 3, 10] # ------------------------------------------------------------- single prompt def run_single(eval_label: str, n: int, format_only: bool): from jspace.dataset import sample_examples item = EVAL_BY_LABEL[eval_label] # Only compute the ground-truth labels this call actually needs: the eval # item, plus the specific pool items sampled as examples for this n. (Cached # across calls, so repeated clicks are cheap.) chosen = sample_examples(SPLIT.pool, int(n), seed=SPLIT.seed, exclude=item.name) labels = exp.ensure_labels([item] + chosen) truth = labels[item.name]["tokens"] gt_layer = labels[item.name]["layer"] gt_pos = labels[item.name]["position"] # Few-shot examples for this n, from the pool. examples = [[ex.prompt, labels[ex.name]["tokens"]] for ex in chosen] rep = rp.report_gpu(item.prompt, examples, format_only=bool(format_only)) from jspace import scoring sc = scoring.score(rep["concepts"], truth, raw_report_text=rep.get("raw", "")) truth_md = "\n".join(f"{i}. `{t}`" for i, t in enumerate(truth, 1)) report_md = "\n".join(f"{i}. {c}" for i, c in enumerate(rep["concepts"], 1)) or "_(empty)_" match_md = ( "\n".join(f"- {r} ≈ `{t}` (sim {s:.2f})" for r, t, s in sc.matches) or "_(no matches above threshold)_" ) scores_md = ( f"**overlap@k** = {sc.overlap_at_k:.2f} ({sc.n_matched}/{sc.k}) | " f"**Spearman** = {sc.spearman:.2f} | " f"**refusal** = {'yes' if sc.is_refusal else 'no'}\n\n" f"_Ground truth: J-lens top-{sc.k} at layer {gt_layer}, position {gt_pos}._" ) return report_md, truth_md, match_md, scores_md # --------------------------------------------------------------- scaling run def run_curve(eval_size: int, progress=gr.Progress()): cfg = RunConfig(eval_size=int(eval_size)) progress(0.3, desc="running grid in a single GPU call…") # Fused: one ZeroGPU attach for the whole grid (avoids the rapid-init # "No CUDA GPUs available" failure from many tiny @spaces.GPU calls). out = exp.run_experiment_fused(config=cfg) curve = out["curve"] msgs: list[str] = [] rows = [ { "condition": c["label"], "n": c["n"], "mean_overlap@k": round(c["mean_overlap"], 3), "mean_spearman": round(c["mean_spearman"], 3), "refusal_rate": round(c["refusal_rate"], 3), "items": c["n_items"], } for c in curve["conditions"] ] table = pd.DataFrame(rows) # Plot only the numeric-n conditions (exclude format-only) for the curve. plot_rows = [ {"n": c["n"], "mean_overlap@k": c["mean_overlap"]} for c in curve["conditions"] if not c["format_only"] ] plot_df = pd.DataFrame(sorted(plot_rows, key=lambda r: r["n"])) summary = ( f"**Chance floor** overlap@k = {curve['baseline_floor_overlap']:.3f} | " f"**Ceiling** (lens self-consistency) = {curve['baseline_ceiling_overlap']:.2f}\n\n" + "\n".join(f"- {m}" for m in msgs[-6:]) ) return table, plot_df, summary def build_ui() -> gr.Blocks: with gr.Blocks(title="ICL Elicitation of J-Space") as demo: gr.Markdown( "# In-Context Elicitation of J-Space Contents\n" "Can `gemma-3-12b-it` report the concepts in its own J-space using " "few-shot prompting alone? Ground truth is the Jacobian-lens top-k " "readout at a fixed (middle layer, last content token). " "See PRD for the full design." ) with gr.Tab("Single prompt"): with gr.Row(): eval_dd = gr.Dropdown( choices=list(EVAL_BY_LABEL), label="Held-out eval prompt", value=next(iter(EVAL_BY_LABEL)), ) n_dd = gr.Dropdown(choices=N_CHOICES, value=3, label="n (in-context examples)") fmt_cb = gr.Checkbox(label="format-only (no worked examples)", value=False) run_btn = gr.Button("Run", variant="primary") with gr.Row(): report_out = gr.Markdown(label="Model report (ranked)") truth_out = gr.Markdown(label="J-lens ground truth (top-k)") match_out = gr.Markdown(label="Matches") scores_out = gr.Markdown(label="Scores") run_btn.click( run_single, inputs=[eval_dd, n_dd, fmt_cb], outputs=[report_out, truth_out, match_out, scores_out], ) with gr.Tab("Scaling curve"): gr.Markdown( "Runs `n ∈ {0,1,3,10}` + format-only over the eval set. " "Results are cached and resumable across ZeroGPU quota windows." ) size_sl = gr.Slider( 2, len(SPLIT.eval), value=min(6, len(SPLIT.eval)), step=1, label="eval prompts to run (smaller = fits free ZeroGPU quota)", ) curve_btn = gr.Button("Run / refresh scaling curve", variant="primary") curve_summary = gr.Markdown() curve_plot = gr.LinePlot( x="n", y="mean_overlap@k", title="Report–vs–ground-truth overlap vs n", ) curve_table = gr.Dataframe(label="Per-condition summary") curve_btn.click( run_curve, inputs=[size_sl], outputs=[curve_table, curve_plot, curve_summary], ) return demo demo = build_ui() if __name__ == "__main__": demo.launch()