app.py

"""Gradio demo for the vineyard plotting-code LoRA adapters.

Pick an adapter from the dropdown, pick one of the 12 synthetic vineyard
DataFrames, ask for a chart. The Space generates the plotting code, then runs it
in a subprocess sandbox and shows you the chart it actually produced -- the same
execution check `pipeline/evaluate.py` scores offline, so a demo can't pass off
plausible-looking code that doesn't run.

Runs on ZeroGPU when the Space has it, and falls back to CPU otherwise.
"""

import json
import os
import subprocess
import sys
import tempfile
from pathlib import Path

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from execute import clean_code_string
from schemas import GENERATORS, df_preview

# ZeroGPU: only the decorated function gets a real GPU. Outside it, PyTorch runs
# in a CUDA emulation mode -- which is why weights are placed on `cuda` in the
# parent process (see _load) and the decorated function only does the forward
# pass. The decorated call inherits the parent's state, so the model cache below
# survives; loading inside the decorator instead would re-download and re-place
# the weights on every single request and burn the daily GPU quota doing it.
#
# On a CPU Space the import fails and everything below runs unchanged on CPU.
def _duration(label, *_args, **_kwargs):
    """Ask for only as much GPU time as the chosen model plausibly needs."""
    return 45 if "0.5B" in label else 90


try:
    import spaces

    GPU = spaces.GPU(duration=_duration)
    ZERO_GPU = True
except Exception:  # noqa: BLE001
    def GPU(fn):
        return fn

    ZERO_GPU = False

HERE = Path(__file__).resolve().parent
GEN_BY_NAME = {g.__name__: g for g in GENERATORS}
SCHEMAS = sorted(GEN_BY_NAME)

# MUST match pipeline/build_dataset.py SYSTEM_PROMPT exactly. Every training
# example carried this system turn; changing a word moves the prompt off the
# distribution the adapters were tuned on.
SYSTEM_PROMPT = (
    "You write Python plotting code for vineyard / viticulture data. "
    "A pandas DataFrame `df` is already loaded. Return only a ```python "
    "code block using matplotlib or seaborn."
)

# label -> (base repo on the Hub, adapter dir in this repo, blurb for the UI)
MODELS = {
    "Qwen2.5-Coder-0.5B · bf16": (
        "Qwen/Qwen2.5-Coder-0.5B-Instruct",
        "models/qwen2.5-coder-0.5b-plotter-lora",
        "Trained in bf16. Best scorecard of the set: 6/10 semantically correct, "
        "0/10 invented columns, 10/10 `tight_layout()`. Fastest to load.",
    ),
    "Qwen2.5-Coder-1.5B · bf16 (best checkpoint)": (
        "unsloth/Qwen2.5-Coder-1.5B-Instruct",
        "models/qwen2.5-coder-1.5b-plotter-lora-bf16-best",
        "The checkpoint selected on eval loss rather than the final step. "
        "Start here if you want the 1.5B.",
    ),
    "Qwen2.5-Coder-1.5B · bf16 (final step)": (
        "unsloth/Qwen2.5-Coder-1.5B-Instruct",
        "models/qwen2.5-coder-1.5b-plotter-lora-bf16",
        "Same run, saved at the last step instead of the best one.",
    ),
    "Qwen2.5-Coder-1.5B · 4-bit NF4 run": (
        "Qwen/Qwen2.5-Coder-1.5B-Instruct",
        "models/qwen2.5-coder-1.5b-plotter-lora",
        "The original 1.5B, trained under 4-bit NF4 because bf16 would not fit "
        "the training card. Overfit at step 186 (eval loss 0.599 vs 0.480 at "
        "step 100) — kept for comparison. Served here in bf16.",
    ),
    "Phi-3.5-mini-instruct · 3.8B": (
        "microsoft/Phi-3.5-mini-instruct",
        "models/phi35-mini-instruct-lora",
        "Largest base in the set. Slowest to load on a cold Space.",
    ),
    "LFM2-2.6B": (
        "LiquidAI/LFM2-2.6B",
        "models/lfm-2.6b-lora",
        "Liquid AI's hybrid-conv base — a different architecture family from the "
        "others.",
    ),
}
DEFAULT_MODEL = "Qwen2.5-Coder-0.5B · bf16"

# One model resident at a time: the big bases will not co-exist in Space RAM.
_LOADED: dict = {"key": None, "model": None, "tokenizer": None}


def _load(label: str, use_adapter: bool):
    """Base model + (optionally) the LoRA adapter, cached across calls."""
    key = (label, use_adapter)
    if _LOADED["key"] == key:
        return _LOADED["model"], _LOADED["tokenizer"]

    repo, adapter_rel, _ = MODELS[label]
    adapter = HERE / adapter_rel

    # Evict first -- loading the new weights before freeing the old ones is what
    # actually OOMs the box.
    _LOADED.update(key=None, model=None, tokenizer=None)
    try:
        torch.cuda.empty_cache()
    except Exception:  # noqa: BLE001 -- no-op under ZeroGPU's CUDA emulation
        pass

    # Prefer the tokenizer saved beside the adapter: it carries the exact chat
    # template training used, so prompts here are formatted identically.
    tok_src = str(adapter) if (adapter / "tokenizer_config.json").exists() else repo
    tokenizer = AutoTokenizer.from_pretrained(tok_src)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # bf16 on any GPU we can land on; CPU has no usable bf16 matmul path here.
    dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32

    model = AutoModelForCausalLM.from_pretrained(repo, dtype=dtype)
    if use_adapter:
        from peft import PeftModel

        model = PeftModel.from_pretrained(model, str(adapter))
    model.to("cuda" if torch.cuda.is_available() else "cpu")
    model.eval()

    _LOADED.update(key=key, model=model, tokenizer=tokenizer)
    return model, tokenizer


@GPU
def _generate(label, preview, request, max_new_tokens, temperature):
    """Forward pass only. _load() must already have run in the parent process."""
    model, tokenizer = _LOADED["model"], _LOADED["tokenizer"]

    # Same three-part shape build_dataset.py wrote, minus the assistant answer.
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {
            "role": "user",
            "content": f"### DataFrame Preview:\n{preview}\n\n"
                       f"### Visualization Request:\n{request}",
        },
    ]
    prompt = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            do_sample=temperature > 0,
            temperature=temperature if temperature > 0 else None,
            pad_token_id=tokenizer.pad_token_id,
        )
    return tokenizer.decode(
        out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True
    ).strip()


def _render(schema: str, code: str, seed: int):
    """Execute the generated code out-of-process; return (ok, error, png path)."""
    png = Path(tempfile.mkdtemp(prefix="chart_")) / "chart.png"
    payload = json.dumps(
        {"schema": schema, "code": code, "seed": int(seed), "png": str(png)}
    )
    try:
        proc = subprocess.run(
            [sys.executable, str(HERE / "sandbox_runner.py")],
            input=payload,
            capture_output=True,
            text=True,
            timeout=45,
            cwd=str(HERE),
            env={**os.environ, "MPLBACKEND": "Agg"},
        )
    except subprocess.TimeoutExpired:
        return False, "timed out after 45s (likely an unbounded loop)", None

    line = (proc.stdout or "").strip().splitlines()
    if not line:
        return False, f"sandbox produced no result. stderr: {(proc.stderr or '')[-400:]}", None
    try:
        res = json.loads(line[-1])
    except json.JSONDecodeError:
        return False, f"unreadable sandbox output: {line[-1][:300]}", None
    return res["ok"], res["error"], res["png"]


def preview_for(schema: str, seed: int) -> str:
    return df_preview(GEN_BY_NAME[schema](seed=int(seed)))


def model_blurb(label: str) -> str:
    repo, adapter, note = MODELS[label]
    return f"**Base:** `{repo}` · **Adapter:** `{adapter}`\n\n{note}"


def run(schema, seed, request, label, use_adapter, max_new_tokens, temperature,
        progress=gr.Progress()):
    request = (request or "").strip()
    if not request:
        return "", "Type a visualization request first.", None

    preview = preview_for(schema, seed)

    progress(0.1, desc=f"Loading {label}…")
    try:
        _load(label, use_adapter)
    except Exception as e:  # noqa: BLE001
        return "", f"❌ Could not load {label} — {type(e).__name__}: {e}", None

    progress(0.4, desc="Generating…")
    try:
        reply = _generate(label, preview, request, max_new_tokens, temperature)
    except Exception as e:  # noqa: BLE001
        return "", f"❌ Generation failed — {type(e).__name__}: {e}", None

    code = clean_code_string(reply) or reply

    progress(0.8, desc="Executing in the sandbox…")
    ok, err, png = _render(schema, code, seed)
    status = "✅ Code executed and produced a chart." if ok else f"❌ Did not run — {err}"
    return code, status, png


EXAMPLES = [
    ["block_vintage_df", "Show average yield per acre by block as a bar chart."],
    ["barrel_aging_df", "Vanillin concentration vs months in barrel, coloured by barrel type."],
    ["cellar_ferment_df", "Plot fermentation temperature over time for each tank."],
    ["irrigation_sensor_df", "Distribution of soil moisture readings by sensor depth."],
    ["wine_sales_df", "Monthly revenue trend broken down by channel."],
    ["berry_chem_df", "Relationship between brix and titratable acidity."],
]

CSS = """
.small-note { font-size: 0.85rem; opacity: 0.75; }
"""

with gr.Blocks(title="Vineyard Plotting-Code Models") as demo:
    gr.Markdown(
        "# 🍇 Vineyard plotting-code models\n"
        "Six LoRA adapters fine-tuned to write matplotlib/seaborn code against "
        "synthetic vineyard DataFrames. Pick one, pick a DataFrame, ask for a "
        "chart. **The generated code is then executed in a sandbox** and the "
        "chart below is what it actually produced — not a mock-up."
    )

    with gr.Row():
        with gr.Column(scale=1):
            model = gr.Dropdown(
                sorted(MODELS), value=DEFAULT_MODEL, label="Model"
            )
            blurb = gr.Markdown(model_blurb(DEFAULT_MODEL), elem_classes="small-note")
            use_adapter = gr.Checkbox(
                value=True,
                label="Use the fine-tuned adapter",
                info="Uncheck to run the untuned base model for comparison.",
            )
            schema = gr.Dropdown(
                SCHEMAS, value="block_vintage_df", label="DataFrame"
            )
            seed = gr.Number(value=0, precision=0, label="Seed")
            request = gr.Textbox(
                label="Visualization request",
                placeholder="Show average yield per acre by block as a bar chart.",
                lines=3,
            )
            with gr.Accordion("Decoding", open=False):
                max_new_tokens = gr.Slider(64, 768, value=512, step=32,
                                           label="Max new tokens")
                temperature = gr.Slider(0.0, 1.0, value=0.2, step=0.05,
                                        label="Temperature (0 = greedy)")
            go = gr.Button("Generate chart", variant="primary")

        with gr.Column(scale=1):
            status = gr.Markdown("")
            chart = gr.Image(label="Rendered chart", type="filepath", height=380)
            code_out = gr.Code(label="Generated code", language="python")

    with gr.Accordion("DataFrame the model is shown", open=False):
        preview = gr.Textbox(
            value=preview_for("block_vintage_df", 0),
            label="Preview (exactly the string passed to the model)",
            lines=10, max_lines=14,
        )

    gr.Examples(EXAMPLES, inputs=[schema, request], label="Try one")

    gr.Markdown(
        f"Running on **{'ZeroGPU' if ZERO_GPU else 'CPU'}**. The first request for "
        "a given model downloads its base weights and is slower than the rest.\n\n"
        "Generated code runs with restricted builtins (no `os`, `open`, `eval`, "
        "`subprocess`) in a separate short-lived process — a proportionate guard "
        "against hallucinated code, not a hardened sandbox.",
        elem_classes="small-note",
    )

    model.change(model_blurb, model, blurb)
    for comp in (schema, seed):
        comp.change(preview_for, [schema, seed], preview)

    go.click(
        run,
        [schema, seed, request, model, use_adapter, max_new_tokens, temperature],
        [code_out, status, chart],
    )

if __name__ == "__main__":
    demo.queue(max_size=12).launch(css=CSS)