app.py

"""Code2LoRA timeline demo.

Pick a held-out repository, move through its commit history, generate the
Code2LoRA-Evo adapter at any commit, and compare the adapted model against the
plain base model on the test assertions that commit introduced, or on your own
code prompt.
"""

from __future__ import annotations

import gzip
import html
import json
import os
import threading
import time
from collections import OrderedDict
from functools import lru_cache
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import gradio as gr
import numpy as np
import torch
from huggingface_hub import hf_hub_download
from matplotlib.figure import Figure
from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList

from c2l.core import (
    Code2LoRAHead,
    clear_all_lora_weights,
    discover_module_types_and_dims,
    get_module_specs,
    inject_lora_weights,
    replace_with_lora,
)
from c2l.metrics import compute_metrics, strip_fim_tokens

BASE_MODEL = "Qwen/Qwen2.5-Coder-1.5B"
CKPT_REPO, CKPT_FILE = "code2lora/code2lora-gru", "code2lora_gru.pt"
TARGET_MODULES = ["q_proj", "k_proj", "v_proj", "o_proj", "up_proj", "gate_proj", "down_proj"]
DATA = Path(os.environ.get("C2L_DATA_DIR", Path(__file__).parent / "data"))
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32
LIVE_MAX_INPUT = 4096 if DEVICE == "cuda" else 1024
PAPER_URL = "https://arxiv.org/abs/2606.06492"
SITE_URL = "https://code2lora.app"

torch.set_grad_enabled(False)


# ---------------------------------------------------------------------------
# Data
# ---------------------------------------------------------------------------

INDEX: List[dict] = json.loads((DATA / "index.json").read_text())
INDEX_BY_SAFE = {r["safe"]: r for r in INDEX}


@lru_cache(maxsize=6)
def load_repo(safe: str) -> dict:
    with gzip.open(DATA / f"{safe}.json.gz", "rt", encoding="utf-8") as f:
        d = json.load(f)
    d["ctx"] = np.load(DATA / f"{safe}.ctx.npy")
    d["tt_commits"] = [k for k, c in enumerate(d["commits"]) if c.get("time_travel")]
    d["live_only"] = INDEX_BY_SAFE[safe].get("c2l_em") is None
    d["jump_targets"] = d["tt_commits"] or [k for k, g in enumerate(d["qnas"]) if g]
    return d


def repo_label(r: dict) -> str:
    return f"{r['repo_id']}  ·  {r['n_commits']} commits"


# ---------------------------------------------------------------------------
# Model
# ---------------------------------------------------------------------------

class NewlineStop(StoppingCriteria):
    def __init__(self, tok, start: int):
        self.tok, self.start = tok, start

    def __call__(self, input_ids, scores, **kwargs):
        done = "\n" in self.tok.decode(input_ids[0, self.start:], skip_special_tokens=True)
        return torch.full((input_ids.shape[0],), done, dtype=torch.bool, device=input_ids.device)


class Engine:
    def __init__(self) -> None:
        t0 = time.time()
        self.tok = AutoTokenizer.from_pretrained(BASE_MODEL)
        if self.tok.pad_token is None:
            self.tok.pad_token = self.tok.eos_token
        self.bos = self.tok.bos_token_id if self.tok.bos_token_id is not None else self.tok.eos_token_id
        self.model = AutoModelForCausalLM.from_pretrained(BASE_MODEL, dtype=DTYPE).to(DEVICE).eval()

        path = hf_hub_download(CKPT_REPO, CKPT_FILE)
        try:
            state = torch.load(path, map_location="cpu", weights_only=False, mmap=True)
        except Exception:
            state = torch.load(path, map_location="cpu", weights_only=False)
        hc = state["head_config"]
        self.specs = get_module_specs(self.model, TARGET_MODULES)
        type_dims = discover_module_types_and_dims(self.specs)
        self.rank = int(hc.get("rank", 16))
        self.head = Code2LoRAHead(input_dim=int(hc["input_dim"]), type_dims=type_dims,
                                  hidden_dim=int(hc["hidden_dim"]), rank=self.rank)
        self.head.load_state_dict(state["head_state"])
        self.head.to(DEVICE).eval()
        alpha = float((state.get("args") or {}).get("alpha", 32.0))
        del state
        replace_with_lora(self.model, self.specs, rank=self.rank, alpha=alpha)
        self.n_adapter_params = sum(self.rank * (i + o) for i, o in type_dims.values())
        self.lock = threading.Lock()
        self.adapters: "OrderedDict[Tuple[str, int], dict]" = OrderedDict()
        print(f"[engine] ready on {DEVICE} in {time.time() - t0:.0f}s", flush=True)

    def adapter(self, safe: str, k: int) -> Tuple[dict, float]:
        key = (safe, k)
        if key in self.adapters:
            self.adapters.move_to_end(key)
            return self.adapters[key], 0.0
        ctx = torch.from_numpy(load_repo(safe)["ctx"][k:k + 1]).to(DEVICE)
        t0 = time.perf_counter()
        out = self.head(ctx)
        if DEVICE == "cuda":
            torch.cuda.synchronize()
        ms = (time.perf_counter() - t0) * 1000
        self.adapters[key] = out
        while len(self.adapters) > 32:
            self.adapters.popitem(last=False)
        return out, ms

    def _ids(self, prompt: str, had_bos: bool, max_input: int) -> List[int]:
        ids = ([self.bos] if had_bos and self.bos is not None else []) + \
            self.tok.encode(prompt, add_special_tokens=False)
        return ids[-max_input:]

    def generate(self, prompt: str, adapter: Optional[dict], *, had_bos: bool = True,
                 max_input: int = LIVE_MAX_INPUT, max_new: int = 64,
                 stop_newline: bool = True) -> Tuple[str, float, int]:
        ids = self._ids(prompt, had_bos, max_input)
        x = torch.tensor([ids], device=DEVICE)
        stop = StoppingCriteriaList([NewlineStop(self.tok, len(ids))]) if stop_newline else None
        with self.lock:
            t0 = time.perf_counter()
            if adapter is None:
                clear_all_lora_weights(self.model, self.specs)
            else:
                inject_lora_weights(self.model, self.specs, adapter, batch_index=0)
            try:
                out = self.model.generate(
                    input_ids=x, attention_mask=torch.ones_like(x), max_new_tokens=int(max_new),
                    do_sample=False, pad_token_id=self.tok.pad_token_id,
                    eos_token_id=self.tok.eos_token_id, stopping_criteria=stop, use_cache=True)
            finally:
                clear_all_lora_weights(self.model, self.specs)
            sec = time.perf_counter() - t0
        text = self.tok.decode(out[0, len(ids):].tolist(), skip_special_tokens=True)
        return text, sec, len(ids)


ENGINE = Engine()


# ---------------------------------------------------------------------------
# Rendering helpers
# ---------------------------------------------------------------------------

def esc(s) -> str:
    return html.escape(str(s if s is not None else ""))


def short_date(ts: Optional[str]) -> str:
    return (ts or "")[:10]


def verdict(em: Optional[float]) -> str:
    if em is None:
        return ""
    return '<span class="ok">✓ exact match</span>' if em == 1.0 else '<span class="bad">✗ wrong</span>'


def pct(x: Optional[float]) -> str:
    return "–" if x is None else f"{100 * x:.0f}%"


def repo_header(safe: str) -> str:
    r = INDEX_BY_SAFE[safe]
    stars = f"★ {r['stars']:,} · " if r.get("stars") else ""
    desc = f"<div class='desc'>{esc(r.get('description'))}</div>" if r.get("description") else ""
    if r.get("c2l_em") is None:
        em = ("<div class='repo-em'>Live mode: cached GPU scores for this repository are still being computed. "
              "Press <b>Adapt</b> and both models answer live on this Space.</div>")
    else:
        em = (f"<div class='repo-em'>Exact match across this repo's history: base model <b>{pct(r['base_em'])}</b>"
              f" → Code2LoRA <b class='accent'>{pct(r['c2l_em'])}</b></div>")
    return f"""
<div class="repo-card">
  <div class="repo-title"><a href="https://github.com/{esc(r['repo_id'])}" target="_blank">{esc(r['repo_id'])}</a>
    <span class="pill held">held out — never seen in training</span></div>
  {desc}
  <div class="repo-stats">{stars}{r['n_commits']} commits · {short_date(r['first_date'])} → {short_date(r['last_date'])}
    · {r['n_qnas']:,} test assertions</div>
  {em}
</div>"""


def commit_card(safe: str, k: int, adapted_k: Optional[int]) -> str:
    d = load_repo(safe)
    c = d["commits"][k]
    n = len(d["commits"])
    url = f"https://github.com/{d['repo_id']}/commit/{c['sha']}"
    tt = "<span class='pill star'>✨ adapter learns something new here</span>" if c.get("time_travel") else ""
    if adapted_k == k:
        status = "<span class='pill live'>● adapter active for this commit</span>"
    elif adapted_k is not None:
        status = f"<span class='pill warn'>adapter is from commit {adapted_k + 1}: press Adapt to update</span>"
    else:
        status = "<span class='pill idle'>no adapter yet: press Adapt</span>"
    return f"""
<div class="commit-card">
  <div class="commit-top"><b>Commit {k + 1} / {n}</b> · <a href="{url}" target="_blank"><code>{c['sha'][:10]}</code></a>
    · {short_date(c.get('timestamp'))} · {c.get('n_new_assertions', 0)} new test assertions {tt}</div>
  <div class="commit-status">{status}</div>
</div>"""


def timeline_plot(safe: str, k: int) -> Figure:
    d = load_repo(safe)
    cs = d["commits"]
    x = np.arange(1, len(cs) + 1)

    def rolling(key: str) -> np.ndarray:
        v = np.array([c[key] if c.get(key) is not None else np.nan for c in cs], dtype=float)
        w = max(3, len(v) // 12)
        out = np.full_like(v, np.nan)
        for i in range(len(v)):
            seg = v[max(0, i - w + 1): i + 1]
            seg = seg[~np.isnan(seg)]
            out[i] = seg.mean() * 100 if len(seg) else np.nan
        return out

    fig = Figure(figsize=(9, 2.4), dpi=110)
    ax = fig.add_subplot(111)
    if d["live_only"]:
        ax.bar(x, [len(g) for g in d["qnas"]], color="#c5cbf5", width=0.9)
        ax.set_ylabel("Test assertions\nadded")
    else:
        ax.plot(x, rolling("base_em"), color="#9aa3b5", lw=1.6, label="Base model")
        ax.plot(x, rolling("c2l_em"), color="#4f5bd5", lw=2.2, label="Code2LoRA (adapter at each commit)")
        for t in d["tt_commits"]:
            ax.axvline(t + 1, color="#22c3a6", alpha=0.18, lw=1)
        ax.set_ylim(0, 100)
        ax.set_ylabel("Exact match %\n(rolling)")
        ax.legend(loc="lower right", fontsize=8, frameon=False)
    ax.axvline(k + 1, color="#ff6b8b", lw=2)
    ax.set_xlim(0.5, max(2, len(cs)) + 0.5)
    ax.set_xlabel("Commit")
    for s in ("top", "right"):
        ax.spines[s].set_visible(False)
    fig.tight_layout()
    return fig


def assertion_choices(safe: str, k: int) -> List[Tuple[str, int]]:
    qs = load_repo(safe)["qnas"][k]
    out = []
    for i, q in enumerate(qs):
        star = "✨ " if q.get("prev_em") == 0.0 else ""
        out.append((f"{star}{Path(q['test_file']).name}:{q['lineno']}  {q.get('test_function') or ''}", i))
    return out


def context_tail(prompt: str, n_lines: int = 28) -> str:
    return "\n".join(prompt.rstrip("\n").split("\n")[-n_lines:])


def pred_row(label: str, cls: str, pred: Optional[str], em: Optional[float], note: str = "") -> str:
    if pred is None:
        return ""
    shown = strip_fim_tokens(pred).split("\n")[0]
    return f"""<tr class="{cls}"><td class="who">{label}</td><td><code>{esc(shown) or '∅'}</code></td>
<td>{verdict(em)}</td><td class="note">{note}</td></tr>"""


def assertion_panel(safe: str, k: int, i: Optional[int], adapted_k: Optional[int]) -> Tuple[str, str]:
    qs = load_repo(safe)["qnas"][k]
    if i is None or i >= len(qs):
        return "", "<div class='hint'>Pick an assertion above.</div>"
    q = qs[i]
    truth = f"""<tr class="truth"><td class="who">Ground truth</td><td><code>{esc(q['target'])}</code></td><td></td>
<td class="note">what the developer actually wrote</td></tr>"""
    if q.get("base_pred") is None:
        hint = ("Both models are answering live below…" if adapted_k == k else
                "Press <b>Adapt at this commit</b>: the adapter is generated and both models answer live.")
        return context_tail(q["prompt"]), f"<table class='preds'>{truth}</table><div class='hint'>{hint}</div>"
    rows = [truth,
            pred_row("Base model", "base", q["base_pred"], q["base_em"], "Qwen2.5-Coder-1.5B, no adapter")]
    if adapted_k == k:
        if q.get("prev_pred") is not None:
            rows.append(pred_row(f"Adapter @ commit {k}", "prev", q["prev_pred"], q["prev_em"],
                                 "one commit earlier"))
        rows.append(pred_row(f"Code2LoRA @ commit {k + 1}", "c2l", q["c2l_pred"], q["c2l_em"],
                             "adapter generated from history up to this commit"))
        foot = ("<div class='hint'>Cached results were computed on a GPU with the exact adapter you just "
                f"generated ({q['n_prompt_tokens']:,}-token context). Press <b>Run live</b> to regenerate on this Space.</div>")
    else:
        rows.append("<tr class='c2l'><td class='who'>Code2LoRA</td><td colspan='3' class='hint'>"
                    "Press <b>Adapt at this commit</b> to generate the adapter and reveal its answer.</td></tr>")
        foot = ""
    return context_tail(q["prompt"]), f"<table class='preds'>{''.join(rows)}</table>{foot}"


# ---------------------------------------------------------------------------
# Event handlers
# ---------------------------------------------------------------------------

def default_commit(safe: str) -> int:
    d = load_repo(safe)
    if not d["jump_targets"]:
        return len(d["commits"]) - 1
    return d["jump_targets"][-1] if d["live_only"] else d["jump_targets"][0]


def on_repo(safe: str):
    d = load_repo(safe)
    k = default_commit(safe)
    choices = assertion_choices(safe, k)
    first = choices[0][1] if choices else None
    ctx, preds = assertion_panel(safe, k, first, None)
    jump = "Jump to next commit with tests" if d["live_only"] else "Jump to next ✨ commit"
    return (gr.Button(value=jump), repo_header(safe),
            gr.Slider(minimum=1, maximum=len(d["commits"]), value=k + 1, step=1),
            commit_card(safe, k, None), timeline_plot(safe, k),
            d["commits"][k].get("diff_preview") or "",
            gr.Radio(choices=choices, value=first), ctx, preds, None, "", "")


def on_commit(safe: str, k1: int, adapted: Optional[Tuple[str, int]]):
    k = int(k1) - 1
    adapted_k = adapted[1] if adapted and adapted[0] == safe else None
    d = load_repo(safe)
    choices = assertion_choices(safe, k)
    first = next((v for lbl, v in choices if lbl.startswith("✨")), choices[0][1] if choices else None)
    ctx, preds = assertion_panel(safe, k, first, adapted_k)
    return (commit_card(safe, k, adapted_k), timeline_plot(safe, k),
            d["commits"][k].get("diff_preview") or "",
            gr.Radio(choices=choices, value=first), ctx, preds, "")


def step(safe: str, k1: int, delta: int) -> int:
    n = len(load_repo(safe)["commits"])
    return int(min(max(1, int(k1) + delta), n))


def next_star(safe: str, k1: int) -> int:
    tt = load_repo(safe)["jump_targets"]
    k = int(k1) - 1
    later = [t for t in tt if t > k]
    return (later[0] if later else (tt[0] if tt else k)) + 1


def on_adapt(safe: str, k1: int, i: Optional[int]):
    k = int(k1) - 1
    out, ms = ENGINE.adapter(safe, k)
    norms = {t: float(out["B"][t][0].norm() * out["A"][t][0].norm()) for t in sorted(out["A"])}
    top = max(norms.values()) or 1.0
    bars = "".join(f"<div class='bar'><span>{t}</span><i style='width:{100 * v / top:.0f}%'></i></div>"
                   for t, v in norms.items())
    timing = "served from cache" if ms == 0.0 else f"generated in {ms:.0f} ms on {DEVICE.upper()}"
    status = f"""
<div class="adapt-card">
  <div class="adapt-head">⚡ Adapter for commit {k + 1} {timing}</div>
  <div class="adapt-sub">{ENGINE.n_adapter_params:,} LoRA parameters (rank {ENGINE.rank}) for 7 projection types,
    shared across all 28 layers · no training, no retrieved context</div>
  <div class="bars">{bars}</div>
</div>"""
    ctx, preds = assertion_panel(safe, k, i, k)
    return status, (safe, k), commit_card(safe, k, k), ctx, preds


def on_pick(safe: str, k1: int, i: Optional[int], adapted: Optional[Tuple[str, int]]):
    k = int(k1) - 1
    adapted_k = adapted[1] if adapted and adapted[0] == safe else None
    ctx, preds = assertion_panel(safe, k, i, adapted_k)
    return ctx, preds, ""


def live_table(rows: List[Tuple[str, str, str, Optional[dict], float]], target: Optional[str]) -> str:
    body = []
    for label, cls, pred, m, sec in rows:
        shown = (m["pred_clean"] if m else strip_fim_tokens(pred).split("\n")[0])
        body.append(f"""<tr class="{cls}"><td class="who">{label}</td><td><code>{esc(shown) or '∅'}</code></td>
<td>{verdict(float(m['exact_match'])) if m else ''}</td><td class="note">{sec:.1f}s</td></tr>""")
    truth = (f"<tr class='truth'><td class='who'>Ground truth</td><td><code>{esc(target)}</code></td>"
             f"<td></td><td></td></tr>") if target else ""
    return f"<div class='live-head'>Live on this Space ({DEVICE.upper()})</div><table class='preds'>{truth}{''.join(body)}</table>"


def on_live(safe: str, k1: int, i: Optional[int], adapted: Optional[Tuple[str, int]],
            progress=gr.Progress()):
    k = int(k1) - 1
    if not adapted or adapted != (safe, k):
        return "<div class='hint warn'>Press <b>Adapt at this commit</b> first.</div>"
    if i is None:
        return "<div class='hint warn'>Pick an assertion first.</div>"
    q = load_repo(safe)["qnas"][k][i]
    adapter, _ = ENGINE.adapter(safe, k)
    progress(0.1, desc="Base model…")
    b, bs, n = ENGINE.generate(q["prompt"], None, had_bos=q.get("had_bos", True))
    progress(0.55, desc="Code2LoRA adapter…")
    a, as_, _ = ENGINE.generate(q["prompt"], adapter, had_bos=q.get("had_bos", True))
    note = "" if n >= q["n_prompt_tokens"] else (
        f"<div class='hint'>CPU mode uses the last {n:,} of {q['n_prompt_tokens']:,} context tokens, "
        "so live answers can differ from the cached GPU run.</div>")
    return live_table([("Base model", "base", b, compute_metrics(b, q["target"]), bs),
                       (f"Code2LoRA @ commit {k + 1}", "c2l", a, compute_metrics(a, q["target"]), as_)],
                      q["target"]) + note


def on_auto_live(safe: str, k1: int, i: Optional[int], adapted: Optional[Tuple[str, int]],
                 progress=gr.Progress()):
    if not load_repo(safe)["live_only"] or adapted != (safe, int(k1) - 1) or i is None:
        return gr.update()
    return on_live(safe, k1, i, adapted, progress)


def on_prefill(safe: str, k1: int, i: Optional[int]) -> str:
    qs = load_repo(safe)["qnas"][int(k1) - 1]
    if i is None or i >= len(qs):
        return ""
    return context_tail(qs[i]["prompt"], 40)


def on_custom(safe: str, k1: int, prompt: str, max_new: int,
              adapted: Optional[Tuple[str, int]], progress=gr.Progress()):
    k = int(k1) - 1
    if not prompt.strip():
        return "<div class='hint warn'>Type or paste a code prefix first.</div>"
    if not adapted or adapted != (safe, k):
        return "<div class='hint warn'>Press <b>Adapt at this commit</b> first.</div>"
    adapter, _ = ENGINE.adapter(safe, k)
    progress(0.1, desc="Base model…")
    b, bs, _ = ENGINE.generate(prompt, None, max_new=max_new, stop_newline=False)
    progress(0.55, desc="Code2LoRA adapter…")
    a, as_, _ = ENGINE.generate(prompt, adapter, max_new=max_new, stop_newline=False)
    rows = []
    for label, cls, text, sec in (("Base model", "base", b, bs), (f"Code2LoRA @ commit {k + 1}", "c2l", a, as_)):
        rows.append(f"<div class='gen {cls}'><div class='who'>{label} · {sec:.1f}s</div>"
                    f"<pre>{esc(strip_fim_tokens(text))}</pre></div>")
    return f"<div class='live-head'>Live on this Space ({DEVICE.upper()})</div><div class='gens'>{''.join(rows)}</div>"


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

CSS = """
.gradio-container {max-width: 1180px !important}
#hero h1 {font-size: 2rem; margin-bottom: 4px; letter-spacing: -0.02em}
#hero p {color: #5a6275; margin: 0}
.repo-card, .commit-card, .adapt-card {border: 1px solid #e3e6ef; border-radius: 12px; padding: 14px 16px; background: #fff}
.repo-title {font-size: 1.15rem; font-weight: 700}
.repo-title a, .commit-top a {color: #2b3487; text-decoration: none}
.desc {color: #5a6275; margin: 4px 0}
.repo-stats, .commit-status {color: #6b7287; font-size: 0.9rem; margin-top: 4px}
.repo-em {margin-top: 8px}
.accent {color: #4f5bd5}
.pill {display: inline-block; font-size: 0.75rem; font-weight: 600; padding: 2px 9px; border-radius: 999px; margin-left: 6px; vertical-align: middle}
.pill.held {background: #eef0ff; color: #3d4ee0}
.pill.star {background: #e2f8f3; color: #0d8a75}
.pill.live {background: #e2f8f3; color: #0d8a75; margin-left: 0}
.pill.warn {background: #fff4e0; color: #a35f00; margin-left: 0}
.pill.idle {background: #f1f2f6; color: #6b7287; margin-left: 0}
.adapt-card {background: linear-gradient(135deg, #f4f5ff, #effbf8)}
.adapt-head {font-weight: 700; font-size: 1.05rem}
.adapt-sub {color: #5a6275; font-size: 0.88rem; margin: 2px 0 10px}
.bars .bar {display: flex; align-items: center; gap: 8px; font-size: 0.78rem; color: #5a6275; margin: 3px 0}
.bars .bar span {width: 76px; font-family: monospace}
.bars .bar i {height: 8px; border-radius: 4px; background: linear-gradient(90deg, #6d7cff, #22d3b6)}
table.preds {width: 100%; border-collapse: collapse; margin-top: 6px}
table.preds td {padding: 8px 10px; border-bottom: 1px solid #eef0f5; vertical-align: top}
table.preds td.who {width: 150px; font-weight: 600; white-space: nowrap}
table.preds td.note {width: 130px; color: #8a91a5; font-size: 0.8rem}
table.preds tr.truth {background: #f7f8fb}
table.preds tr.c2l {background: #f4f5ff}
table.preds tr.prev {background: #fffaf2}
table.preds code {font-size: 0.92rem; white-space: pre-wrap; word-break: break-word}
.ok {color: #0d8a75; font-weight: 700; white-space: nowrap}
.bad {color: #d6455e; font-weight: 700; white-space: nowrap}
.hint {color: #6b7287; font-size: 0.86rem; margin-top: 8px}
.hint.warn {color: #a35f00}
.live-head {font-weight: 700; margin: 6px 0 2px}
.gens {display: grid; grid-template-columns: 1fr 1fr; gap: 12px}
.gen {border: 1px solid #e3e6ef; border-radius: 10px; padding: 10px}
.gen.c2l {background: #f4f5ff}
.gen pre {white-space: pre-wrap; font-size: 0.85rem; margin: 6px 0 0}
"""

INTRO = f"""
<div id="hero">
<h1>Code2LoRA: watch an adapter learn a repository, commit by commit</h1>
<p>Pick a repository the model never saw during training, move through its git history, and generate the
repository-specific LoRA adapter at any point in time. Then compare the adapted model with the plain base model
on the test assertions that commit introduced, or on your own code.
<a href="{PAPER_URL}" target="_blank">Paper</a> · <a href="{SITE_URL}" target="_blank">code2lora.app</a> ·
<a href="https://pypi.org/project/code2lora/" target="_blank">pip install code2lora</a></p>
</div>
"""

HOW = """
**How it works.** For every commit, the production-code diff was embedded once with a frozen
Qwen3-Embedding-0.6B and folded into a GRU state (precomputed for this demo, 8 KB per commit).
**Adapt** runs the Code2LoRA hypernetwork head live on that state and emits a rank-16 LoRA for all attention and
MLP projections of the frozen Qwen2.5-Coder-1.5B. The base model never sees retrieved context: everything it
knows about the repository comes from the adapter.

**What you see.** Each commit's test assertions are the ones the developer added in that commit, scored
exactly as in the paper (first ≤ 8 per commit, exact match after normalization). ✨ marks commits where the
adapter gets an assertion right that the adapter from one commit earlier got wrong.
"""

choices = [(repo_label(r), r["safe"]) for r in INDEX]
start = max(INDEX, key=lambda r: ((r["c2l_em"] or 0) - (r["base_em"] or 0), r.get("stars") or 0))["safe"]

with gr.Blocks(css=CSS, title="Code2LoRA timeline demo", theme=gr.themes.Soft(primary_hue="indigo")) as demo:
    adapted_state = gr.State(None)
    gr.HTML(INTRO)
    with gr.Row():
        repo_dd = gr.Dropdown(choices=choices, value=start, label="Repository (held out from training)", scale=3)
    repo_html = gr.HTML()

    with gr.Group():
        commit_slider = gr.Slider(1, 2, value=1, step=1, label="Commit")
        with gr.Row():
            prev_btn = gr.Button("◀ Previous", size="sm")
            next_btn = gr.Button("Next ▶", size="sm")
            star_btn = gr.Button("Jump to next ✨ commit", size="sm")
            adapt_btn = gr.Button("⚡ Adapt at this commit", variant="primary", size="sm")
    commit_html = gr.HTML()
    plot = gr.Plot(show_label=False)
    adapt_html = gr.HTML()

    with gr.Tabs():
        with gr.Tab("Test assertions from this commit"):
            pick = gr.Radio(label="Assertions added in this commit", choices=[])
            with gr.Row():
                with gr.Column(scale=5):
                    ctx_code = gr.Code(label="Test file, right before the assertion's expected value",
                                       language="python", interactive=False, lines=18)
                with gr.Column(scale=6):
                    preds_html = gr.HTML()
                    live_btn = gr.Button("Run live on this Space")
                    live_html = gr.HTML()
        with gr.Tab("Your own prompt"):
            gr.Markdown("Paste a code prefix (for example a test ending in `assert result ==`). "
                        "Both models continue it; the adapter is the one generated for the selected commit.")
            prefill_btn = gr.Button("Prefill with the selected assertion's context", size="sm")
            custom_in = gr.Code(language="python", lines=14, label="Code prefix")
            max_new = gr.Slider(8, 128, value=48, step=8, label="Max new tokens")
            custom_btn = gr.Button("Generate with base model and Code2LoRA", variant="primary")
            custom_out = gr.HTML()
        with gr.Tab("Commit diff"):
            diff_code = gr.Code(label="Production-code diff (first 6,000 characters)", interactive=False, lines=24)
        with gr.Tab("About"):
            gr.Markdown(HOW + ("\n\n*Running on CPU: live generation uses the last "
                               f"{LIVE_MAX_INPUT:,} context tokens and takes roughly 10–40 s.*"
                               if DEVICE == "cpu" else ""))

    repo_outputs = [star_btn, repo_html, commit_slider, commit_html, plot, diff_code, pick, ctx_code, preds_html,
                    adapted_state, adapt_html, live_html]
    commit_outputs = [commit_html, plot, diff_code, pick, ctx_code, preds_html, live_html]

    demo.load(on_repo, repo_dd, repo_outputs)
    repo_dd.change(on_repo, repo_dd, repo_outputs)
    commit_slider.release(on_commit, [repo_dd, commit_slider, adapted_state], commit_outputs)
    prev_btn.click(lambda s, k: step(s, k, -1), [repo_dd, commit_slider], commit_slider).then(
        on_commit, [repo_dd, commit_slider, adapted_state], commit_outputs)
    next_btn.click(lambda s, k: step(s, k, 1), [repo_dd, commit_slider], commit_slider).then(
        on_commit, [repo_dd, commit_slider, adapted_state], commit_outputs)
    star_btn.click(next_star, [repo_dd, commit_slider], commit_slider).then(
        on_commit, [repo_dd, commit_slider, adapted_state], commit_outputs)
    live_inputs = [repo_dd, commit_slider, pick, adapted_state]
    adapt_btn.click(on_adapt, [repo_dd, commit_slider, pick],
                    [adapt_html, adapted_state, commit_html, ctx_code, preds_html], concurrency_limit=1).then(
        on_auto_live, live_inputs, live_html, concurrency_limit=1)
    pick.change(on_pick, live_inputs, [ctx_code, preds_html, live_html]).then(
        on_auto_live, live_inputs, live_html, concurrency_limit=1)
    live_btn.click(on_live, [repo_dd, commit_slider, pick, adapted_state], live_html, concurrency_limit=1)
    prefill_btn.click(on_prefill, [repo_dd, commit_slider, pick], custom_in)
    custom_btn.click(on_custom, [repo_dd, commit_slider, custom_in, max_new, adapted_state], custom_out,
                     concurrency_limit=1)

if __name__ == "__main__":
    demo.queue(max_size=32).launch()