app.py

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