Timeline demo: precomputed GRU states, live adapter generation

README.md

@@ -1,130 +1,36 @@
 ---
-title: Code2LoRA-GRU
+title: Code2LoRA Timeline
 emoji: 🧬
-colorFrom: blue
+colorFrom: indigo
 colorTo: green
 sdk: gradio
-sdk_version: 5.9.1
+sdk_version: 5.50.0
 app_file: app.py
-pinned: false
+pinned: true
 license: mit
-short_description: Turn a repo's commit history into a LoRA adapter, live.
+short_description: Generate a repo-specific LoRA at any commit and test it.
+models:
+  - code2lora/code2lora-gru
+  - Qwen/Qwen2.5-Coder-1.5B
+datasets:
+  - code2lora/code2lora-evo
 ---
 
-# Code2LoRA-GRU — interactive demo
+# Code2LoRA: watch an adapter learn a repository, commit by commit
 
-Paste a Python repository (a **git URL** or a **local clone path**) and watch a
-trained hypernetwork turn its *commit history* into a parameter-efficient LoRA
-adapter for a frozen code model — then watch that adapter answer questions the
-base model gets wrong.
+Pick one of the held-out test repositories from
+[RepoPeftBench](https://huggingface.co/datasets/code2lora/code2lora-evo), move through its commit history,
+generate the Code2LoRA-Evo adapter at any commit, and compare the adapted
+Qwen2.5-Coder-1.5B with the plain base model on the test assertions that commit
+introduced, or on your own code prefix.
 
-This Space implements the inference pipeline from **Code2LoRA: Hypernetworks
-for Repository-Conditioned and Commit-Streaming Adapters of Code Language
-Models**.
+- **Adapt** runs the hypernetwork head live on the precomputed GRU state of the
+  selected commit and injects the resulting rank-16 LoRA into the frozen base model.
+- Cached predictions were produced on a GPU with exactly the same adapter, prompt
+  (≤ 4,096 tokens) and scoring as the paper; **Run live** regenerates them here.
+- `tools/` contains the scripts that built `data/` from the published dataset and
+  checkpoint.
 
----
-
-## What you get
-
-0. **Requirements panel** — what kind of repository works (Python + pytest /
-   unittest, real git history, assertions that grew over time).
-1. **Repository input** — paste a public git URL (cloned with full history) or
-   an absolute path to an existing clone.
-2. **Statistics** (first **80%** of the history):
-   - total first-parent commits,
-   - commits that touch test files,
-   - **commits selected for the GRU pass** (those introducing *new/changed*
-     assertions),
-   - number of assertion QnAs mined,
-   - chronological **80/10/10** train/val/test split,
-   - a **commit timeline** (dot size = new assertions, colored by split),
-   - the **assertion-type distribution**.
-2.2 **GRU processing visualization** — a PCA trajectory of the GRU hidden state
-   as it streams through the repository, plus the curve of accumulated
-   "repository knowledge" feeding the adapter.
-3. **Prediction showcase** — pick a **held-out** assertion (from the val/test
-   tail the adapter never saw), blank its right-hand side, and compare:
-   - the plain `Qwen2.5-Coder-1.5B` completion,
-   - the completion **with the Code2LoRA-GRU adapter injected**,
-   - the ground truth — with exact-match / edit-similarity badges.
-
----
-
-## Architecture (one screen)
-
-```
-repo ──► select kept commits ──► per-commit diff embeddings (Qwen3-Embedding-0.6B, 2048-d)
-                                        │
-   repo-state @ commit 0 ──► h0 ──►  ┌──┴── CommitGRU (O(1) / commit) ──► h_T
-                                        │
-                                   Code2LoRAHead ──► LoRA(A,B) per module type
-                                        │
-                          injected into frozen Qwen2.5-Coder-1.5B ──► prediction
-```
-
-* **Commit selection / diffs / assertion mining** — `c2l_demo/git_pipeline.py`,
-  `c2l_demo/assertions.py`.
-* **Embeddings** — `c2l_demo/embedding.py` (diff: 512/64 chunks → MaxPool‖MeanPool;
-  repo-state: 2048/256 chunks → mean‖max, L2-norm).
-* **GRU + adapter head + LoRA injection** — `c2l_demo/core.py`.
-* **Orchestration / generation / metrics** — `c2l_demo/engine.py`,
-  `c2l_demo/metrics.py`.
-* **Plots** — `c2l_demo/viz.py`.
-
-Every module is a faithful port of the corresponding file in the research
-project (`hypernetwork/code2lora_core.py`,
-`evaluation/run_code2lora_gru_v2_eval.py`,
-`create_dataset/build_commit_parquet_db.py`, …) so the demo scores predictions
-exactly the way the paper does.
-
----
-
-## Configuration
-
-The Space loads the published paper checkpoint by default — **no setup
-required**:
-
-* model repo: [`code2lora/code2lora-gru`](https://huggingface.co/code2lora/code2lora-gru)
-* file: `code2lora_gru.pt` (~2.85 GB — the GRU + the LoRA-generation head)
-
-To point at a different checkpoint, override via environment variables
-(Space → **Settings → Variables and secrets**):
-
-| Variable | Default | Meaning |
-|---|---|---|
-| `CODE2LORA_CKPT_REPO` | `code2lora/code2lora-gru` | HF model repo holding the checkpoint |
-| `CODE2LORA_CKPT_FILE` | `code2lora_gru.pt` | filename in that repo |
-| `CODE2LORA_CKPT` | — | *or* an absolute local path (overrides the repo) |
-
-The base model (`Qwen/Qwen2.5-Coder-1.5B`) and the encoder
-(`Qwen/Qwen3-Embedding-0.6B`) are pulled from the Hub automatically.
-
-> **Hardware.** A GPU Space (T4/A10/A100 or ZeroGPU) is strongly recommended —
-> three models (~7 GB) and live generation. It will also run on CPU, but
-> processing a repo and generating will be slow.
-
----
-
-## Deploy
-
-The checkpoint is already on the Hub, so deploying is just pushing this folder:
-
-```bash
-# Create the Space and push this directory
-huggingface-cli repo create code2lora-gru-demo --type space --space_sdk gradio
-git init && git remote add origin https://huggingface.co/spaces/code2lora/code2lora-gru-demo
-git add . && git commit -m "Code2LoRA-GRU demo" && git push origin main
-# Then pick a GPU tier in the Space settings.
-```
-
-`upload_checkpoint.py` is provided only if you ever need to publish a *new*
-checkpoint; the default one is ready to use.
-
-## Run locally
-
-```bash
-pip install -r requirements.txt
-# uses code2lora/code2lora-gru by default; or point at a local checkpoint:
-# export CODE2LORA_CKPT=/abs/path/to/gru_head.best.pt
-python app.py   # serves on http://localhost:7860
-```
+Paper: [arXiv 2606.06492](https://arxiv.org/abs/2606.06492) ·
+Product: [code2lora.app](https://code2lora.app) ·
+CLI: `pip install code2lora`

app.py

@@ -1,260 +1,429 @@
-"""Code2LoRA-GRU — interactive demo (Hugging Face Space).
+"""Code2LoRA timeline demo.
 
-Paste a Python repository (local path or git URL). The Space:
-
-1. mines its commit history and selects the commits that introduce new test
-   assertions (the "GRU pass"),
-2. shows repository statistics, restricted to the first 80% of the history,
-3. visualizes how the streaming GRU walks the repository and accumulates a
-   repository-specific adapter,
-4. lets you watch the generated LoRA change what the base code model predicts
-   for a held-out assertion.
-
-See README.md for the architecture and deployment notes.
+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 tempfile
-import traceback
+import json
+import os
+import threading
+import time
+from collections import OrderedDict
+from functools import lru_cache
 from pathlib import Path
-from typing import List, Optional
+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)
 
-from c2l_demo.engine import Engine, RepoSession
 
 # ---------------------------------------------------------------------------
-# Process-wide engine (lazy model loading happens on first use)
+# Data
 # ---------------------------------------------------------------------------
 
-ENGINE = Engine()
-WORK_DIR = Path(tempfile.gettempdir()) / "code2lora_demo_repos"
-WORK_DIR.mkdir(parents=True, exist_ok=True)
-
-EXAMPLE_REPOS = [
-    "https://github.com/psf/cachecontrol",
-    "https://github.com/john-kurkowski/tldextract",
-    "https://github.com/agronholm/anyio",
-]
-
-REQUIREMENTS_MD = """
-### What kind of repository works best?
-
-Code2LoRA-GRU learns from how a repository's **test suite** evolves, so the
-demo needs a repo that:
-
-- **is a Python project** with a `pytest` / `unittest` test suite,
-- **has real git history** (clone with full history — not a shallow/zip
-  download), because the GRU streams *commit by commit*,
-- contains **assertions that grew over time** (`assert x == ...`,
-  `self.assertEqual(...)`, `pytest.raises(...)`, `assert_allclose(...)`, …),
-- is a **moderate size** (a few hundred commits is ideal; very large
-  monorepos will be slow on CPU).
-
-Paste either a **public git URL** (it will be cloned) or an **absolute local
-path** to a clone that already exists on the Space.
-"""
+INDEX: List[dict] = json.loads((DATA / "index.json").read_text())
+INDEX_BY_SAFE = {r["safe"]: r for r in INDEX}
 
-HOW_IT_WORKS_MD = """
-**Code2LoRA-GRU** turns a repository into a *parameter-efficient adapter* for a
-frozen code model (`Qwen2.5-Coder-1.5B`) — no fine-tuning, no in-context
-retrieval at inference time.
-
-1. **Commit selection** — we keep only commits that introduce *new or changed*
-   test assertions; these form the streaming sequence.
-2. **Encoding** — each kept commit's production-code diff is embedded with a
-   frozen `Qwen3-Embedding-0.6B` (2048-d). The initial repository state seeds
-   the recurrent hidden state.
-3. **Recurrence** — a GRU walks the commits, accumulating repository knowledge
-   in its hidden state in **O(1) per commit**.
-4. **Adapter generation** — a shared-basis head maps the hidden state to LoRA
-   weights for every attention/MLP projection, which are injected into the
-   frozen base model.
-5. **Prediction** — the adapted model fills in a held-out assertion that plain
-   prompting cannot recover.
-"""
+
+@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")]
+    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, torch_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 _stat_card(label: str, value: str, sub: str = "", accent: str = "#4f8cff") -> str:
-    return f"""
-    <div style="flex:1;min-width:150px;background:#161922;border:1px solid #2a2e3a;
-                border-radius:12px;padding:14px 16px;">
-      <div style="font-size:12px;color:#8b93a7;text-transform:uppercase;
-                  letter-spacing:.05em;">{html.escape(label)}</div>
-      <div style="font-size:28px;font-weight:700;color:{accent};margin-top:4px;">
-        {html.escape(value)}</div>
-      <div style="font-size:12px;color:#8b93a7;margin-top:2px;">{html.escape(sub)}</div>
-    </div>"""
-
-
-def render_stats(session: RepoSession) -> str:
-    r = session.result
-    pct_kept = (100.0 * r.kept_commits / r.total_commits) if r.total_commits else 0.0
-    boundary = session.endpoint_index + 1
-    cards = "".join([
-        _stat_card("Total commits", f"{r.total_commits:,}",
-                   "first-parent history", "#e6e6e6"),
-        _stat_card("Touch test files", f"{r.test_touching_commits:,}",
-                   "commits changing a test", "#a78bfa"),
-        _stat_card("GRU-pass commits", f"{r.kept_commits:,}",
-                   f"{pct_kept:.0f}% introduce new assertions", "#22d3a6"),
-        _stat_card("Assertions (QnAs)", f"{r.n_assertions:,}",
-                   "prediction targets mined", "#4f8cff"),
-    ])
-    split = "".join([
-        _stat_card("Train (≤80%)", f"{r.n_train_kept:,}",
-                   f"walked: {boundary} commits", "#4f8cff"),
-        _stat_card("Val (10%)", f"{r.n_val_kept:,}", "held out", "#ffcc4d"),
-        _stat_card("Test (10%)", f"{r.n_test_kept:,}", "held out", "#ff6b6b"),
-        _stat_card("Repo-state files", f"{session.n_repo_state_files:,}",
-                   "embedded at commit 0", "#a78bfa"),
-    ])
-    return f"""
-    <div style="display:flex;gap:12px;flex-wrap:wrap;margin-bottom:12px;">{cards}</div>
-    <div style="display:flex;gap:12px;flex-wrap:wrap;">{split}</div>
-    <p style="color:#8b93a7;font-size:13px;margin-top:12px;">
-      The adapter is built by streaming the GRU over the first
-      <b style="color:#22d3a6;">80%</b> of the
-      {r.kept_commits:,} GRU-pass commits ({boundary:,} commits). The remaining
-      val/test commits are never seen by the adapter — predictions on them are
-      genuine generalization.</p>"""
-
-
-def _badge(ok: bool) -> str:
-    if ok:
-        return ('<span style="background:#0f3d2e;color:#22d3a6;border:1px solid '
-                '#22d3a6;border-radius:6px;padding:2px 8px;font-size:12px;">'
-                '✓ exact match</span>')
-    return ('<span style="background:#3d1f1f;color:#ff6b6b;border:1px solid '
-            '#ff6b6b;border-radius:6px;padding:2px 8px;font-size:12px;">'
-            '✗ miss</span>')
-
-
-def render_prediction(res: dict) -> str:
-    prefix = res["prefix"]
-    # Show only the tail of the prefix so the card stays readable.
-    tail = prefix[-1400:]
-    if len(prefix) > len(tail):
-        tail = "…\n" + tail
-    target = res["target"]
-    base = res["base_clean"] or "(empty)"
-    adapted = res["adapted_clean"] or "(empty)"
-
-    def block(title, body, color):
-        return f"""
-        <div style="flex:1;min-width:240px;background:#161922;border:1px solid #2a2e3a;
-                    border-radius:12px;padding:14px 16px;">
-          <div style="font-size:12px;color:#8b93a7;text-transform:uppercase;
-                      letter-spacing:.05em;margin-bottom:6px;">{title}</div>
-          <div style="font-family:ui-monospace,Menlo,monospace;font-size:15px;
-                      color:{color};word-break:break-word;">{html.escape(body)}</div>
-        </div>"""
-
-    cards = "".join([
-        block("Base Qwen2.5-Coder (no adapter)", base, "#ff9e6b") +
-        f'<div style="margin-top:8px;">{_badge(res["base_exact_match"])}'
-        f'<span style="color:#8b93a7;font-size:12px;margin-left:8px;">'
-        f'edit sim {res["base_edit_similarity"]:.2f}</span></div>',
-        block("+ Code2LoRA-GRU adapter", adapted, "#22d3a6") +
-        f'<div style="margin-top:8px;">{_badge(res["adapted_exact_match"])}'
-        f'<span style="color:#8b93a7;font-size:12px;margin-left:8px;">'
-        f'edit sim {res["adapted_edit_similarity"]:.2f}</span></div>',
-        block("Ground truth", target, "#e6e6e6"),
-    ])
+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 ""
     return f"""
-    <div style="background:#0f1117;border:1px solid #2a2e3a;border-radius:12px;
-                padding:14px 16px;margin-bottom:12px;">
-      <div style="font-size:12px;color:#8b93a7;text-transform:uppercase;
-                  letter-spacing:.05em;margin-bottom:6px;">
-        Test prompt (tail · {res['n_prompt_tokens']} tokens) — model must
-        complete the ▶ blank</div>
-      <pre style="white-space:pre-wrap;font-family:ui-monospace,Menlo,monospace;
-                  font-size:12.5px;color:#cfd6e4;margin:0;max-height:320px;
-                  overflow:auto;">{html.escape(tail)}<span style="color:#22d3a6;
-                  ">▶ ???</span></pre>
-    </div>
-    <div style="display:flex;gap:12px;flex-wrap:wrap;">{cards}</div>"""
-
-
-def _qna_choices(session: RepoSession) -> List:
-    """Build dropdown choices, preferring held-out (val/test) assertions."""
-    items = []
-    for c in session.result.commits:
-        for qi, q in enumerate(c.qnas):
-            held = c.in_repo_split in ("val", "test")
-            items.append((c, qi, q, held))
-    # Held-out first, then train; cap to keep the dropdown snappy.
-    items.sort(key=lambda x: (not x[3], x[0].commit_index))
-    choices = []
-    for c, qi, q, held in items[:300]:
-        tag = "★ held-out" if held else "train"
-        fn = q.test_function or "?"
-        snippet = " ".join(q.prefix.strip().splitlines()[-1:])[-48:]
-        label = f"[{tag}] {Path(q.test_file).name}::{fn} L{q.lineno} — …{snippet}"
-        choices.append((label, f"{c.commit_index}:{qi}"))
-    return choices
+<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>
+  <div class="repo-em">Exact match across this repo's history:
+    base model <b>{pct(r['base_em'])}</b> → Code2LoRA <b class="accent">{pct(r['c2l_em'])}</b></div>
+</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)
+    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.axvline(k + 1, color="#ff6b8b", lw=2)
+    ax.set_xlim(1, max(2, len(cs)))
+    ax.set_ylim(0, 100)
+    ax.set_ylabel("Exact match %\n(rolling)")
+    ax.set_xlabel("Commit")
+    for s in ("top", "right"):
+        ax.spines[s].set_visible(False)
+    ax.legend(loc="lower right", fontsize=8, frameon=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]
+    rows = [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>""",
+            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}"
 
 
 # ---------------------------------------------------------------------------
-# Callbacks
+# Event handlers
 # ---------------------------------------------------------------------------
 
-def process_repo_cb(repo_input: str, progress=gr.Progress()):
-    from c2l_demo import viz
-
-    def _cb(msg, frac):
-        try:
-            progress(min(max(frac, 0.0), 1.0), desc=msg)
-        except Exception:
-            pass
-
-    try:
-        progress(0.0, desc="Resolving repository…")
-        repo_dir, repo_id = ENGINE.resolve_repo(repo_input, WORK_DIR, progress=_cb)
-        session = ENGINE.process_and_walk(repo_dir, repo_id, walk_fraction=0.8,
-                                          progress=_cb)
-    except Exception as e:
-        err = (f"<div style='color:#ff6b6b;background:#1a1014;border:1px solid "
-               f"#ff6b6b;border-radius:10px;padding:14px;'><b>Could not process "
-               f"this repository.</b><br><br><pre style='white-space:pre-wrap;"
-               f"color:#ffb3b3;font-size:12px;'>{html.escape(str(e))}</pre></div>")
-        return (err, None, None, None,
-                gr.update(choices=[], value=None), None, "")
-
-    stats = render_stats(session)
-    fig_timeline = viz.plot_commit_timeline(session.result, session.endpoint_index)
-    fig_types = viz.plot_assertion_type_breakdown(session.result)
-    fig_gru = viz.plot_gru_processing(session)
-    choices = _qna_choices(session)
-    dd = gr.update(choices=choices, value=(choices[0][1] if choices else None))
-    ready = ("<div style='color:#22d3a6;'>Adapter ready — pick a held-out "
-             "assertion below and run the prediction.</div>")
-    return stats, fig_timeline, fig_types, fig_gru, dd, session, ready
-
-
-def predict_cb(session: Optional[RepoSession], selected: Optional[str]):
-    if session is None:
-        return ("<div style='color:#ff6b6b;'>Process a repository first.</div>")
-    if not selected:
-        return ("<div style='color:#ff6b6b;'>Pick an assertion to predict.</div>")
-    try:
-        ci_str, qi_str = selected.split(":")
-        ci, qi = int(ci_str), int(qi_str)
-        commit = next(c for c in session.result.commits if c.commit_index == ci)
-        q = commit.qnas[qi]
-        res = ENGINE.predict(session, q.prefix, q.target)
-        return render_prediction(res)
-    except Exception as e:
-        return (f"<div style='color:#ff6b6b;'><b>Prediction failed.</b><pre "
-                f"style='white-space:pre-wrap;font-size:12px;'>"
-                f"{html.escape(traceback.format_exc())}</pre></div>")
+def default_commit(safe: str) -> int:
+    d = load_repo(safe)
+    return d["tt_commits"][0] if d["tt_commits"] else len(d["commits"]) - 1
+
+
+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)
+    return (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)["tt_commits"]
+    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_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>"
 
 
 # ---------------------------------------------------------------------------
@@ -262,92 +431,137 @@ def predict_cb(session: Optional[RepoSession], selected: Optional[str]):
 # ---------------------------------------------------------------------------
 
 CSS = """
-.gradio-container {max-width: 1180px !important;}
-#hero {background: linear-gradient(135deg,#1b2a4a 0%,#0f1117 60%);
-       border:1px solid #2a2e3a;border-radius:16px;padding:24px 28px;margin-bottom:8px;}
-#hero h1 {font-size:30px;margin:0 0 6px 0;background:linear-gradient(90deg,#4f8cff,#22d3a6);
-          -webkit-background-clip:text;-webkit-text-fill-color:transparent;}
-.section-title {font-size:20px;font-weight:700;margin-top:6px;}
+.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: 170px; font-weight: 600; white-space: nowrap}
+table.preds td.note {color: #8a91a5; font-size: 0.82rem}
+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>
+"""
 
-def build_demo():
-    with gr.Blocks(title="Code2LoRA-GRU Demo", theme=gr.themes.Soft(
-            primary_hue="blue", neutral_hue="slate"), css=CSS) as demo:
-        session_state = gr.State(value=None)
+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.
 
-        gr.HTML("""
-        <div id="hero">
-          <h1>Code2LoRA-GRU</h1>
-          <div style="color:#cfd6e4;font-size:15px;max-width:820px;">
-            Turn any Python repository into a parameter-efficient adapter for a
-            frozen code LLM — by <b>streaming its commit history through a GRU</b>
-            and generating LoRA weights on the fly. No fine-tuning, no retrieval
-            at inference time.
-          </div>
-        </div>""")
+**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.
+"""
 
-        with gr.Accordion("How it works", open=False):
-            gr.Markdown(HOW_IT_WORKS_MD)
+choices = [(repo_label(r), r["safe"]) for r in INDEX]
+start = max(INDEX, key=lambda r: (r["c2l_em"] - r["base_em"]))["safe"]
 
-        # 0) Instructions / requirements
-        with gr.Accordion("0 · Repository requirements (read me first)", open=True):
-            gr.Markdown(REQUIREMENTS_MD)
+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()
 
-        # 1) Input
-        gr.Markdown("## 1 · Point at a repository", elem_classes="section-title")
-        with gr.Row():
-            repo_input = gr.Textbox(
-                label="Git URL or local path",
-                placeholder="https://github.com/owner/repo  or  /abs/path/to/clone",
-                scale=4)
-            process_btn = gr.Button("Process repository", variant="primary", scale=1)
-        gr.Examples(EXAMPLE_REPOS, inputs=repo_input, label="Example repositories")
-        status = gr.HTML()
-
-        # 2) Statistics
-        gr.Markdown("## 2 · Repository statistics (first 80% of history)",
-                    elem_classes="section-title")
-        stats_html = gr.HTML()
-        with gr.Row():
-            timeline_plot = gr.Plot(label="Commit timeline · dot size = new assertions")
+    with gr.Group():
+        commit_slider = gr.Slider(1, 2, value=1, step=1, label="Commit")
         with gr.Row():
-            types_plot = gr.Plot(label="Assertion-type distribution")
-
-        # 2.2) GRU processing
-        gr.Markdown("## 2.2 · How the GRU processes the repository",
-                    elem_classes="section-title")
-        gru_plot = gr.Plot(label="Hidden-state trajectory + accumulated repository knowledge")
-
-        # 3) Prediction
-        gr.Markdown("## 3 · Watch the adapter predict a held-out assertion",
-                    elem_classes="section-title")
-        gr.Markdown(
-            "Pick an assertion (★ = a *held-out* val/test commit the adapter "
-            "never saw). We blank its right-hand side and ask the model to fill "
-            "it in — first the plain base model, then with the repository "
-            "adapter injected.")
-        with gr.Row():
-            qna_dropdown = gr.Dropdown(label="Held-out assertion to complete",
-                                       choices=[], scale=4)
-            predict_btn = gr.Button("Predict", variant="primary", scale=1)
-        prediction_html = gr.HTML()
-
-        gr.Markdown(
-            "<div style='color:#8b93a7;font-size:12px;margin-top:16px;'>"
-            "Base model: <code>Qwen/Qwen2.5-Coder-1.5B</code> · Encoder: "
-            "<code>Qwen/Qwen3-Embedding-0.6B</code> · Adapter: Code2LoRA-GRU "
-            "(rank-16 LoRA generated by a commit-streaming GRU).</div>")
-
-        process_btn.click(
-            process_repo_cb, inputs=[repo_input],
-            outputs=[stats_html, timeline_plot, types_plot, gru_plot,
-                     qna_dropdown, session_state, status])
-        predict_btn.click(
-            predict_cb, inputs=[session_state, qna_dropdown],
-            outputs=[prediction_html])
-    return demo
-
+            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 = [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)
+    adapt_btn.click(on_adapt, [repo_dd, commit_slider, pick],
+                    [adapt_html, adapted_state, commit_html, ctx_code, preds_html], concurrency_limit=1)
+    pick.change(on_pick, [repo_dd, commit_slider, pick, adapted_state], [ctx_code, preds_html, live_html])
+    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__":
-    build_demo().queue().launch(server_name="0.0.0.0", server_port=7860)
+    demo.queue(max_size=32).launch()

{c2l_demo → c2l}/metrics.py

@@ -1,23 +1,10 @@
-"""Evaluation metrics for the demo.
+"""Assertion-completion scoring, identical to RepoPeftBench's evaluation/metrics.py
+(exact match and edit similarity; CodeBLEU omitted)."""
 
-Faithful port of ``evaluation/metrics.py`` (exact match / edit similarity /
-postprocessing) so the demo scores predictions exactly the way the paper does.
-The optional ``codebleu`` dependency is not required; a tokenized BLEU
-fallback is used when it is absent.
-"""
-
-from __future__ import annotations
-
-import re
 from difflib import SequenceMatcher
 
 FIM_TOKENS = ("<|fim_prefix|>", "<|fim_suffix|>", "<|fim_middle|>")
 
-_CODE_TOKEN_RE = re.compile(
-    r"[a-zA-Z_]\w*|0[xXoObB][\da-fA-F_]+|\d[\d_]*\.?\d*[eE]?[\d_]*"
-    r"|\"(?:[^\"\\]|\\.)*\"|'(?:[^'\\]|\\.)*'|[^\s]"
-)
-
 
 def strip_fim_tokens(s: str) -> str:
     for tok in FIM_TOKENS:
@@ -43,42 +30,6 @@ def edit_similarity(pred: str, ref: str) -> float:
     return SequenceMatcher(None, pred, ref).ratio()
 
 
-def _tokenize_code(s: str):
-    return _CODE_TOKEN_RE.findall(s)
-
-
-def _ngram_bleu(pred_tokens, ref_tokens, max_n: int = 4) -> float:
-    from collections import Counter
-    import math
-
-    if not pred_tokens or not ref_tokens:
-        return 0.0
-    brevity = min(1.0, len(pred_tokens) / len(ref_tokens))
-    log_bleu = 0.0
-    for n in range(1, max_n + 1):
-        pred_ngrams = Counter(
-            tuple(pred_tokens[i:i + n]) for i in range(len(pred_tokens) - n + 1))
-        ref_ngrams = Counter(
-            tuple(ref_tokens[i:i + n]) for i in range(len(ref_tokens) - n + 1))
-        matches = sum(min(pred_ngrams[ng], ref_ngrams[ng]) for ng in pred_ngrams)
-        total = max(len(pred_tokens) - n + 1, 1)
-        precision = (matches + 1) / (total + 1)
-        log_bleu += math.log(precision) / max_n
-    return brevity * math.exp(log_bleu)
-
-
-def code_bleu_score(pred: str, ref: str, lang: str = "python") -> float:
-    if not pred.strip() or not ref.strip():
-        return 0.0
-    if pred.strip() == ref.strip():
-        return 1.0
-    try:
-        from codebleu import calc_codebleu
-        return calc_codebleu([ref], [pred], lang=lang)["codebleu"]
-    except Exception:
-        return _ngram_bleu(_tokenize_code(pred), _tokenize_code(ref))
-
-
 def postprocess_prediction(pred: str, target: str) -> str:
     pred = strip_fim_tokens(pred)
     if "\n" not in target and "\n" in pred:
@@ -97,20 +48,7 @@ def compute_metrics(pred: str, ref: str) -> dict:
     pred_clean = postprocess_prediction(pred, ref)
     ref_clean = strip_comments(ref)
     return {
+        "pred_clean": pred_clean,
         "exact_match": exact_match(pred_clean, ref_clean),
         "edit_similarity": edit_similarity(pred_clean, ref_clean),
-        "code_bleu": code_bleu_score(pred_clean, ref_clean),
-        "pred_clean": pred_clean,
-        "ref_clean": ref_clean,
     }
-
-
-__all__ = [
-    "compute_metrics",
-    "postprocess_prediction",
-    "exact_match",
-    "edit_similarity",
-    "code_bleu_score",
-    "normalize_for_match",
-    "strip_comments",
-]

c2l_demo/__init__.py

@@ -1,17 +0,0 @@
-"""Code2LoRA-GRU interactive demo package.
-
-Self-contained re-packaging of the Code2LoRA-GRU inference pipeline so it can
-run as a standalone Hugging Face Space without depending on the research
-monorepo. Every module here is a faithful port of the corresponding file in
-the `RepoPeftData` project; see each module docstring for provenance.
-"""
-
-__all__ = [
-    "core",
-    "assertions",
-    "git_pipeline",
-    "embedding",
-    "metrics",
-    "engine",
-    "viz",
-]

c2l_demo/assertions.py

@@ -1,360 +0,0 @@
-"""Assertion extraction from Python test files.
-
-Faithful port of the AST-based assertion parsers in
-``create_dataset/create_qnas.py`` plus the full-file-prefix extractor
-``extract_assertions_from_source`` from
-``create_dataset/build_commit_parquet_db.py``.
-
-A QnA pair is one assertion turned into (prefix, target):
-
-* ``prefix`` -- the full test-file content up to (but not including) the
-  assertion's right-hand side, plus the in-line prefix (e.g. ``assert x == ``).
-* ``target`` -- the right-hand side the model must predict (e.g. ``42``).
-"""
-
-from __future__ import annotations
-
-import ast
-import re
-import warnings
-from dataclasses import dataclass
-from typing import List, Optional
-
-SELF_ASSERT_METHODS = {
-    "assertEqual", "assertNotEqual", "assertTrue", "assertFalse",
-    "assertIs", "assertIsNot", "assertIsNone", "assertIsNotNone",
-    "assertIn", "assertNotIn", "assertIsInstance", "assertNotIsInstance",
-    "assertRaises", "assertRaisesRegex", "assertAlmostEqual", "assertNotAlmostEqual",
-    "assertDictEqual", "assertListEqual", "assertTupleEqual", "assertSetEqual",
-    "assertSequenceEqual", "assertRegex", "assertNotRegex", "assertCountEqual",
-    "assertWarns", "assertLogs", "assertMultiLineEqual",
-    "assertGreater", "assertGreaterEqual", "assertLess", "assertLessEqual",
-}
-
-
-@dataclass
-class AssertionPair:
-    assertion_type: str
-    prefix: str
-    target: str
-    lineno: int
-    col_offset: int
-
-
-def get_source_segment(source: str, node: ast.AST) -> Optional[str]:
-    if hasattr(ast, "get_source_segment"):
-        return ast.get_source_segment(source, node)
-    lines = source.splitlines()
-    start = getattr(node, "lineno", 1) - 1
-    end = getattr(node, "end_lineno", start + 1) - 1
-    if start < 0 or end >= len(lines):
-        return None
-    return "\n".join(lines[start:end + 1])
-
-
-def flatten_to_oneliner(s: str) -> str:
-    if not s or "\n" not in s:
-        return s
-    return " ".join(s.split())
-
-
-# --- 1. plain assert -------------------------------------------------------
-
-def _parse_plain_assert(source, node, lines):
-    src = get_source_segment(source, node)
-    if not src:
-        return None
-    lineno, col = node.lineno, node.col_offset
-    test = node.test
-    if isinstance(test, ast.Compare):
-        left_src = get_source_segment(source, test.left)
-        if not left_src:
-            return None
-        op = test.ops[0]
-        comparators = test.comparators
-        if not comparators:
-            return None
-        rhs_src = get_source_segment(source, comparators[0])
-        if not rhs_src:
-            return None
-        op_map = {
-            ast.Eq: " == ", ast.NotEq: " != ", ast.Lt: " < ", ast.LtE: " <= ",
-            ast.Gt: " > ", ast.GtE: " >= ", ast.In: " in ", ast.NotIn: " not in ",
-            ast.Is: " is ", ast.IsNot: " is not ",
-        }
-        op_str = op_map.get(type(op), " ??? ")
-        prefix = f"assert {left_src}{op_str}"
-        target = rhs_src
-    else:
-        prefix = "assert "
-        target = get_source_segment(source, test) or ""
-    return AssertionPair("assert", prefix, target, lineno, col)
-
-
-# --- 2. self.assert* -------------------------------------------------------
-
-def _parse_self_assert(source, node, lines):
-    if not isinstance(node.func, ast.Attribute):
-        return None
-    if not isinstance(node.func.value, ast.Name) or node.func.value.id != "self":
-        return None
-    method = node.func.attr
-    if method not in SELF_ASSERT_METHODS:
-        return None
-    args = node.args
-    if not args:
-        return None
-    if not get_source_segment(source, node):
-        return None
-
-    if method == "assertRaises":
-        exc_src = get_source_segment(source, args[0])
-        if exc_src:
-            return AssertionPair(f"self.{method}", "self.assertRaises(",
-                                 exc_src + ")", node.lineno, node.col_offset)
-        return None
-    if method == "assertRaisesRegex":
-        if len(args) >= 2:
-            exc_src = get_source_segment(source, args[0])
-            pat_src = get_source_segment(source, args[1])
-            if exc_src and pat_src:
-                return AssertionPair(f"self.{method}", "self.assertRaisesRegex(",
-                                     f"{exc_src}, {pat_src})",
-                                     node.lineno, node.col_offset)
-        return None
-    if method in ("assertTrue", "assertFalse", "assertIsNone", "assertIsNotNone"):
-        arg_src = get_source_segment(source, args[0])
-        if arg_src:
-            return AssertionPair(f"self.{method}", f"self.{method}(",
-                                 arg_src + ")", node.lineno, node.col_offset)
-        return None
-    if len(args) >= 2:
-        last_src = get_source_segment(source, args[-1])
-        if last_src:
-            prefix_parts = [f"self.{method}("]
-            for i, arg in enumerate(args[:-1]):
-                arg_src = get_source_segment(source, arg)
-                if arg_src:
-                    prefix_parts.append(arg_src)
-                    if i < len(args) - 2:
-                        prefix_parts.append(", ")
-            prefix = "".join(prefix_parts) + ", "
-            return AssertionPair(f"self.{method}", prefix, last_src + ")",
-                                 node.lineno, node.col_offset)
-    return None
-
-
-# --- 3. pytest.raises ------------------------------------------------------
-
-def _parse_pytest_raises(source, node, lines):
-    if not isinstance(node.func, ast.Attribute):
-        return None
-    val = node.func.value
-    if not isinstance(val, ast.Name) or val.id != "pytest" or node.func.attr != "raises":
-        return None
-    args = node.args
-    if not args:
-        return None
-    exc_src = get_source_segment(source, args[0])
-    if not exc_src:
-        return None
-    kw_parts = []
-    for kw in node.keywords:
-        if kw.arg and kw.value:
-            kw_src = get_source_segment(source, kw.value)
-            if kw_src:
-                kw_parts.append(f', {kw.arg}={kw_src}')
-    target = exc_src + "".join(kw_parts) + ")"
-    return AssertionPair("pytest.raises", "pytest.raises(", target,
-                         node.lineno, node.col_offset)
-
-
-# --- 4. pytest.approx ------------------------------------------------------
-
-def _parse_pytest_approx(source, node, lines):
-    if not isinstance(node.func, ast.Attribute):
-        return None
-    val = node.func.value
-    if not isinstance(val, ast.Name) or val.id != "pytest" or node.func.attr != "approx":
-        return None
-    args = node.args
-    if not args:
-        return None
-    arg_src = get_source_segment(source, args[0])
-    if not arg_src:
-        return None
-    kw_parts = []
-    for kw in node.keywords:
-        if kw.arg and kw.value:
-            kw_src = get_source_segment(source, kw.value)
-            if kw_src:
-                kw_parts.append(f', {kw.arg}={kw_src}')
-    target = arg_src + "".join(kw_parts) + ")"
-    return AssertionPair("pytest.approx", "pytest.approx(", target,
-                         node.lineno, node.col_offset)
-
-
-# --- 5. assert_* (numpy/testing) ------------------------------------------
-
-def _parse_assert_underscore(source, node, lines):
-    func = node.func
-    if isinstance(func, ast.Attribute):
-        name = func.attr
-        if not name.startswith("assert_") or name == "assert_":
-            return None
-        if isinstance(func.value, ast.Attribute):
-            if func.value.attr == "mock" or (
-                isinstance(func.value.value, ast.Name)
-                and func.value.value.id == "unittest"
-            ):
-                return None
-    elif isinstance(func, ast.Name):
-        name = func.id
-        if not name.startswith("assert_") or name == "assert_":
-            return None
-    else:
-        return None
-    args = node.args
-    if not args:
-        return None
-    last_src = get_source_segment(source, args[-1])
-    if not last_src:
-        return None
-    prefix_parts = [f"{get_source_segment(source, func) or name}("]
-    for i, arg in enumerate(args[:-1]):
-        arg_src = get_source_segment(source, arg)
-        if arg_src:
-            prefix_parts.append(arg_src)
-            if i < len(args) - 2:
-                prefix_parts.append(", ")
-    prefix = "".join(prefix_parts) + ", "
-    return AssertionPair("assert_*", prefix, last_src + ")",
-                         node.lineno, node.col_offset)
-
-
-def _find_enclosing_test(tree, node):
-    node_start = getattr(node, "lineno", 0)
-    node_end = getattr(node, "end_lineno", node_start)
-    best = None
-    for n in ast.walk(tree):
-        if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef)):
-            if "test" in n.name.lower():
-                s, e = getattr(n, "lineno", 0), getattr(n, "end_lineno", 0)
-                if s <= node_start and node_end <= e:
-                    if best is None or (getattr(best, "lineno", 0) < s):
-                        best = n
-        if isinstance(n, ast.ClassDef) and "test" in n.name.lower():
-            s, e = getattr(n, "lineno", 0), getattr(n, "end_lineno", 0)
-            if s <= node_start and node_end <= e:
-                for item in n.body:
-                    if (isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef))
-                            and item.name.startswith("test_")):
-                        is_, ie = getattr(item, "lineno", 0), getattr(item, "end_lineno", 0)
-                        if is_ <= node_start and node_end <= ie:
-                            if best is None or getattr(best, "lineno", 0) < is_:
-                                best = item
-    return best
-
-
-# ---------------------------------------------------------------------------
-# Full-file-prefix extractor (build_commit_parquet_db.py)
-# ---------------------------------------------------------------------------
-
-@dataclass
-class ExtractedAssertion:
-    assertion_type: str
-    prefix: str
-    target: str
-    lineno: int
-    col_offset: int
-    test_function: str
-
-
-_WHITESPACE_RE = re.compile(r"\s+")
-
-
-def normalize_for_id(s: str) -> str:
-    return _WHITESPACE_RE.sub(" ", (s or "").strip())
-
-
-def _full_prefix_for(source_lines: List[str], pair: AssertionPair) -> str:
-    lineno = pair.lineno
-    col = pair.col_offset
-    if lineno < 1:
-        return pair.prefix
-    pre = source_lines[: lineno - 1]
-    last_line = source_lines[lineno - 1] if lineno - 1 < len(source_lines) else ""
-    if col < 0:
-        col = 0
-    last_line_prefix = last_line[:col] if col <= len(last_line) else last_line
-    return "".join(pre) + last_line_prefix + pair.prefix
-
-
-def extract_assertions_from_source(source: str) -> List[ExtractedAssertion]:
-    """Walk the AST and return full-file-prefix assertions for one file."""
-    if not source:
-        return []
-    try:
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore", SyntaxWarning)
-            tree = ast.parse(source)
-    except (SyntaxError, ValueError):
-        return []
-
-    lines = source.splitlines(keepends=True)
-    out: List[ExtractedAssertion] = []
-
-    for node in ast.walk(tree):
-        pair: Optional[AssertionPair] = None
-        if isinstance(node, ast.Assert):
-            pair = _parse_plain_assert(source, node, lines)
-        elif isinstance(node, ast.Call):
-            func = node.func
-            if isinstance(func, ast.Attribute):
-                val = func.value
-                val_id = getattr(val, "id", None) if isinstance(val, ast.Name) else None
-                if val_id == "self" and func.attr in SELF_ASSERT_METHODS:
-                    pair = _parse_self_assert(source, node, lines)
-                elif val_id == "pytest":
-                    if func.attr == "raises":
-                        pair = _parse_pytest_raises(source, node, lines)
-                    elif func.attr == "approx":
-                        pair = _parse_pytest_approx(source, node, lines)
-                    else:
-                        pair = _parse_assert_underscore(source, node, lines)
-                else:
-                    pair = _parse_assert_underscore(source, node, lines)
-            elif isinstance(func, ast.Name) and func.id.startswith("assert_"):
-                pair = _parse_assert_underscore(source, node, lines)
-
-        if pair is None or not pair.target.strip():
-            continue
-        target = pair.target
-        if target.lstrip().startswith(","):
-            continue
-        if "\n" in target:
-            target = flatten_to_oneliner(target)
-        if not target.strip():
-            continue
-
-        full_prefix = _full_prefix_for(lines, pair)
-        test_node = _find_enclosing_test(tree, node)
-        test_name = test_node.name if test_node else ""
-        out.append(ExtractedAssertion(
-            assertion_type=pair.assertion_type,
-            prefix=full_prefix,
-            target=target,
-            lineno=pair.lineno,
-            col_offset=pair.col_offset,
-            test_function=test_name,
-        ))
-    return out
-
-
-__all__ = [
-    "AssertionPair",
-    "ExtractedAssertion",
-    "extract_assertions_from_source",
-    "normalize_for_id",
-    "SELF_ASSERT_METHODS",
-]

c2l_demo/embedding.py

@@ -1,165 +0,0 @@
-"""Frozen Qwen3-Embedding-0.6B encoder.
-
-Reproduces the two embedding recipes used to build the Code2LoRA-GRU v2
-dataset, verbatim:
-
-* **Diff embedding** (``build_diff_embeddings_shard.py`` / ``DiffEmbedder``):
-  chunk the diff text at 512 tokens (overlap 64), masked-mean each chunk, then
-  ``concat(MaxPool, MeanPool)`` over chunks -> 2048-d (no L2 norm).
-
-* **Repo-state embedding** (``build_repo_state_embeddings_shard.py``):
-  chunk each .py file at 2048 tokens (overlap 256), masked-mean each chunk,
-  average the chunk vectors per file -> 1024-d; then over all files
-  ``concat(mean_files, max_files)`` -> 2048-d and L2-normalize.
-"""
-
-from __future__ import annotations
-
-from typing import List, Optional
-
-import numpy as np
-import torch
-
-EMBED_MODEL_NAME = "Qwen/Qwen3-Embedding-0.6B"
-
-# Diff recipe
-DIFF_CHUNK_TOKENS = 512
-DIFF_OVERLAP = 64
-DIFF_MAX_LENGTH = 512
-DIFF_MIN_WINDOW = 16
-
-# Repo-state recipe
-REPO_CHUNK_TOKENS = 2048
-REPO_OVERLAP = 256
-REPO_MIN_WINDOW = 8
-
-
-def _chunk_token_ids(token_ids: List[int], chunk_tokens: int, overlap: int,
-                     min_window: int) -> List[List[int]]:
-    if chunk_tokens <= 0 or overlap >= chunk_tokens:
-        return []
-    chunks = []
-    step = chunk_tokens - overlap
-    n = len(token_ids)
-    for start in range(0, n, step):
-        end = min(start + chunk_tokens, n)
-        window = token_ids[start:end]
-        if len(window) < min_window:
-            if end >= n:
-                break
-            continue
-        chunks.append(window)
-        if end >= n:
-            break
-    return chunks
-
-
-class Qwen3Embedder:
-    """Lazy-loading wrapper around the frozen Qwen3-Embedding-0.6B encoder."""
-
-    def __init__(self, device: str = "cpu",
-                 dtype: Optional[torch.dtype] = None,
-                 batch_size: int = 8,
-                 model_name: str = EMBED_MODEL_NAME):
-        self.device = device
-        self.dtype = dtype
-        self.batch_size = batch_size
-        self.model_name = model_name
-        self._model = None
-        self._tokenizer = None
-        self._hidden = None
-
-    def load(self) -> None:
-        if self._model is not None:
-            return
-        from transformers import AutoModel, AutoTokenizer
-        self._tokenizer = AutoTokenizer.from_pretrained(self.model_name)
-        model = AutoModel.from_pretrained(self.model_name)
-        if self.dtype is not None:
-            model = model.to(self.dtype)
-        self._model = model.to(self.device).eval()
-        self._hidden = int(self._model.config.hidden_size)
-
-    @property
-    def hidden_size(self) -> int:
-        self.load()
-        return self._hidden
-
-    @property
-    def embed_dim(self) -> int:
-        return 2 * self.hidden_size
-
-    @torch.no_grad()
-    def _embed_windows(self, texts: List[str], max_length: int) -> torch.Tensor:
-        """[N, D] fp32 (cpu) masked-mean chunk embeddings."""
-        if not texts:
-            return torch.empty((0, self.hidden_size))
-        all_vecs = []
-        for i in range(0, len(texts), self.batch_size):
-            batch = texts[i:i + self.batch_size]
-            enc = self._tokenizer(batch, padding=True, truncation=True,
-                                  max_length=max_length, return_tensors="pt")
-            enc = {k: v.to(self.device) for k, v in enc.items()}
-            out = self._model(**enc)
-            last = out.last_hidden_state
-            mask = enc["attention_mask"].unsqueeze(-1).to(last.dtype)
-            mean = (last * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
-            all_vecs.append(mean.detach().float().cpu())
-        return torch.cat(all_vecs, dim=0)
-
-    # ---- Diff embedding ----
-
-    @torch.no_grad()
-    def embed_diff(self, diff_text: str) -> np.ndarray:
-        """One diff -> [2*D] fp32 (MaxPool || MeanPool, no L2 norm)."""
-        self.load()
-        zero = np.zeros(2 * self._hidden, dtype=np.float32)
-        if not diff_text or not diff_text.strip():
-            return zero
-        ids = self._tokenizer.encode(diff_text, add_special_tokens=False)
-        windows = _chunk_token_ids(ids, DIFF_CHUNK_TOKENS, DIFF_OVERLAP, DIFF_MIN_WINDOW)
-        if not windows:
-            return zero
-        texts = [self._tokenizer.decode(w, skip_special_tokens=True) for w in windows]
-        chunk_embs = self._embed_windows(texts, DIFF_MAX_LENGTH)  # [K, D]
-        mean_pool = chunk_embs.mean(dim=0)
-        max_pool = chunk_embs.max(dim=0).values
-        return torch.cat([max_pool, mean_pool], dim=-1).numpy().astype(np.float32)
-
-    # ---- Repo-state embedding ----
-
-    @torch.no_grad()
-    def _embed_file(self, text: str) -> Optional[np.ndarray]:
-        ids = self._tokenizer.encode(text, add_special_tokens=False)
-        if not ids:
-            return None
-        windows = _chunk_token_ids(ids, REPO_CHUNK_TOKENS, REPO_OVERLAP, REPO_MIN_WINDOW)
-        if not windows:
-            return None
-        texts = [self._tokenizer.decode(w, skip_special_tokens=True) for w in windows]
-        chunk_embs = self._embed_windows(texts, REPO_CHUNK_TOKENS)  # [K, D]
-        return chunk_embs.mean(dim=0).numpy().astype(np.float32)
-
-    @torch.no_grad()
-    def embed_repo_state(self, file_texts: List[str]) -> np.ndarray:
-        """List of file contents -> [2*D] L2-normalized repo-state vector."""
-        self.load()
-        repo_dim = 2 * self._hidden
-        file_vecs = []
-        for text in file_texts:
-            if not text:
-                continue
-            vec = self._embed_file(text)
-            if vec is not None:
-                file_vecs.append(vec)
-        if not file_vecs:
-            return np.zeros(repo_dim, dtype=np.float32)
-        sub = np.stack(file_vecs).astype(np.float32)
-        mean_pool = sub.mean(axis=0)
-        max_pool = sub.max(axis=0)
-        repo_vec = np.concatenate([mean_pool, max_pool], axis=0)
-        norm = float(np.linalg.norm(repo_vec)) + 1e-12
-        return (repo_vec / norm).astype(np.float32)
-
-
-__all__ = ["Qwen3Embedder", "EMBED_MODEL_NAME"]

c2l_demo/engine.py

@@ -1,400 +0,0 @@
-"""End-to-end inference engine for the Code2LoRA-GRU demo.
-
-Ties everything together:
-
-1. resolve a repository (local path or git URL -> shallow-free clone),
-2. select the kept "GRU-pass" commits + per-repo statistics,
-3. embed each commit's production-code diff and the initial repo state with
-   the frozen Qwen3-Embedding encoder,
-4. walk the commits through the :class:`CommitGRU`, recording the hidden-state
-   trajectory,
-5. generate a repository-conditioned LoRA from the hidden state and inject it
-   into the frozen Qwen2.5-Coder base model,
-6. predict a held-out assertion both with and without the adapter.
-
-Heavy models are loaded lazily and cached on the process-wide
-:class:`Engine`. Per-repository state lives in a :class:`RepoSession`.
-"""
-
-from __future__ import annotations
-
-import os
-import re
-import shutil
-import subprocess
-import tempfile
-import time
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Callable, Dict, List, Optional, Tuple
-
-import numpy as np
-import torch
-
-from .core import (
-    Code2LoRAHead,
-    CommitGRU,
-    clear_all_lora_weights,
-    discover_module_types_and_dims,
-    get_module_specs,
-    inject_lora_weights,
-    replace_with_lora,
-)
-from .embedding import Qwen3Embedder
-from .git_pipeline import RepoProcessResult, ls_tree_py, process_repo, GitCatFileBatch
-from .metrics import compute_metrics
-
-DEFAULT_BASE_MODEL = "Qwen/Qwen2.5-Coder-1.5B"
-DEFAULT_TARGET_MODULES = [
-    "q_proj", "k_proj", "v_proj", "o_proj",
-    "up_proj", "gate_proj", "down_proj",
-]
-
-# Where to find the trained GRU+head checkpoint. Either a local path
-# (CODE2LORA_CKPT) or a HF Hub model repo (CODE2LORA_CKPT_REPO + _FILE).
-ENV_CKPT_PATH = "CODE2LORA_CKPT"
-ENV_CKPT_REPO = "CODE2LORA_CKPT_REPO"
-ENV_CKPT_FILE = "CODE2LORA_CKPT_FILE"
-# The published paper checkpoint (GRU + Code2LoRAHead). Used by default so the
-# Space works zero-config; override with the env vars above if needed.
-DEFAULT_CKPT_REPO = "code2lora/code2lora-gru"
-DEFAULT_CKPT_FILE = "code2lora_gru.pt"
-
-ProgressFn = Optional[Callable[[str, float], None]]
-
-_GIT_URL_RE = re.compile(r"^(https?://|git@|ssh://)")
-
-
-def _is_git_url(s: str) -> bool:
-    return bool(_GIT_URL_RE.match(s.strip()))
-
-
-# ---------------------------------------------------------------------------
-# Per-repository session
-# ---------------------------------------------------------------------------
-
-@dataclass
-class RepoSession:
-    repo_id: str
-    result: RepoProcessResult
-    # GRU trajectory over the walked commits:
-    walked_indices: List[int]                  # kept-commit indices walked
-    hidden_states: np.ndarray                  # [T, H] post-step h[-1]
-    ctx_norms: List[float]                     # ||output_norm(h_t)||
-    endpoint_index: int                        # kept-commit index of the adapter
-    repo_state_norm: float
-    n_repo_state_files: int
-    # cached tensors for re-prediction without re-walking:
-    _hidden_endpoint: Optional[torch.Tensor] = None  # [H] on engine device
-
-
-# ---------------------------------------------------------------------------
-# Engine
-# ---------------------------------------------------------------------------
-
-class Engine:
-    def __init__(self,
-                 base_model_name: str = DEFAULT_BASE_MODEL,
-                 target_modules: Optional[List[str]] = None,
-                 device: Optional[str] = None,
-                 max_repo_state_files: int = 400):
-        self.base_model_name = base_model_name
-        self.target_modules = target_modules or list(DEFAULT_TARGET_MODULES)
-        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
-        self.dtype = torch.bfloat16 if self.device == "cuda" else torch.float32
-        self.max_repo_state_files = max_repo_state_files
-
-        self._tokenizer = None
-        self._base_model = None
-        self._specs = None
-        self._gru: Optional[CommitGRU] = None
-        self._head: Optional[Code2LoRAHead] = None
-        self._embedder: Optional[Qwen3Embedder] = None
-        self._loaded = False
-
-    # ---- lazy loading ----
-
-    @property
-    def embedder(self) -> Qwen3Embedder:
-        if self._embedder is None:
-            self._embedder = Qwen3Embedder(
-                device=self.device,
-                dtype=(torch.float16 if self.device == "cuda" else None),
-            )
-        return self._embedder
-
-    def _resolve_ckpt(self) -> str:
-        local = os.environ.get(ENV_CKPT_PATH)
-        if local:
-            if Path(local).exists():
-                return local
-            raise FileNotFoundError(f"{ENV_CKPT_PATH}={local} does not exist.")
-        from huggingface_hub import hf_hub_download
-        repo = os.environ.get(ENV_CKPT_REPO, DEFAULT_CKPT_REPO)
-        fname = os.environ.get(ENV_CKPT_FILE, DEFAULT_CKPT_FILE)
-        return hf_hub_download(repo_id=repo, filename=fname)
-
-    def load(self, progress: ProgressFn = None) -> None:
-        if self._loaded:
-            return
-        from transformers import AutoModelForCausalLM, AutoTokenizer
-
-        def _p(msg, frac):
-            if progress:
-                progress(msg, frac)
-
-        _p("Loading checkpoint…", 0.05)
-        ckpt_path = self._resolve_ckpt()
-        state = torch.load(ckpt_path, map_location="cpu", weights_only=False)
-        head_cfg = state.get("head_config") or {}
-        gru_cfg = state.get("gru_config") or {}
-        rank = int(head_cfg.get("rank", 16))
-        alpha = float((state.get("args") or {}).get("alpha", 32.0))
-        head_hidden = int(head_cfg.get("hidden_dim", 1024))
-
-        _p(f"Loading base model {self.base_model_name}…", 0.2)
-        self._tokenizer = AutoTokenizer.from_pretrained(self.base_model_name)
-        if self._tokenizer.pad_token is None:
-            self._tokenizer.pad_token = self._tokenizer.eos_token
-        self._base_model = self._load_base_model()
-        self._base_model.eval()
-        for p in self._base_model.parameters():
-            p.requires_grad = False
-
-        _p("Wiring LoRA modules…", 0.7)
-        self._specs = get_module_specs(self._base_model, self.target_modules)
-        type_dims = discover_module_types_and_dims(self._specs)
-        input_dim = int(head_cfg.get("input_dim", gru_cfg.get("hidden_dim", 2048)))
-
-        self._head = Code2LoRAHead(input_dim=input_dim, type_dims=type_dims,
-                                   hidden_dim=head_hidden, rank=rank).to(self.device)
-        self._head.load_state_dict(state["head_state"])
-        self._head.eval()
-
-        self._gru = CommitGRU(
-            diff_input_dim=int(gru_cfg.get("diff_input_dim", 2048)),
-            repo_state_dim=int(gru_cfg.get("repo_state_dim", 2048)),
-            hidden_dim=int(gru_cfg.get("hidden_dim", 2048)),
-        ).to(self.device)
-        self._gru.load_state_dict(state["gru_state"])
-        self._gru.eval()
-
-        replace_with_lora(self._base_model, self._specs, rank=rank, alpha=alpha)
-        _p("Models ready.", 1.0)
-        self._loaded = True
-
-    def _load_base_model(self):
-        from transformers import AutoModelForCausalLM
-        last_err = None
-        for attn in ("flash_attention_2", "sdpa", "eager"):
-            try:
-                return AutoModelForCausalLM.from_pretrained(
-                    self.base_model_name,
-                    torch_dtype=self.dtype,
-                    attn_implementation=attn,
-                    device_map={"": self.device} if self.device == "cuda" else None,
-                ).to(self.device)
-            except Exception as e:  # flash-attn may be unavailable
-                last_err = e
-                continue
-        raise RuntimeError(f"Could not load base model: {last_err}")
-
-    # ---- repository resolution ----
-
-    def resolve_repo(self, repo_input: str, work_dir: Path,
-                     progress: ProgressFn = None) -> Tuple[Path, str]:
-        """Return (local_repo_dir, repo_id). Clones if a URL is given."""
-        repo_input = repo_input.strip()
-        if not repo_input:
-            raise ValueError("Please provide a repository path or git URL.")
-        if _is_git_url(repo_input):
-            repo_id = _repo_id_from_url(repo_input)
-            dest = work_dir / repo_id.replace("/", "__")
-            if dest.exists():
-                shutil.rmtree(dest, ignore_errors=True)
-            if progress:
-                progress(f"Cloning {repo_input}…", 0.0)
-            res = subprocess.run(
-                ["git", "clone", "--no-single-branch", repo_input, str(dest)],
-                capture_output=True, text=True, timeout=900)
-            if res.returncode != 0 or not (dest / ".git").exists():
-                raise RuntimeError(f"git clone failed:\n{res.stderr[-2000:]}")
-            return dest, repo_id
-        path = Path(repo_input).expanduser()
-        if not (path / ".git").exists():
-            raise ValueError(f"{path} is not a git repository (no .git found).")
-        return path, _repo_id_from_path(path)
-
-    # ---- processing + GRU walk ----
-
-    def process_and_walk(self, repo_dir: Path, repo_id: str,
-                         walk_fraction: float = 0.8,
-                         progress: ProgressFn = None) -> RepoSession:
-        """Full pipeline: select commits, embed, walk the GRU, cache the
-        adapter context at the requested walk endpoint."""
-        self.load(progress=progress)
-
-        def _p(msg, frac):
-            if progress:
-                progress(msg, frac)
-
-        result = process_repo(repo_dir, repo_id=repo_id, progress=progress)
-        if result.kept_commits == 0:
-            raise RuntimeError(
-                "No commits introduced new assertions in test files — this "
-                "repository has no usable Code2LoRA history. Pick a repo with "
-                "a pytest/unittest suite that grew over time.")
-
-        # Endpoint: the last train-split commit, i.e. the chronological 80%
-        # boundary defined by the dataset's 80/10/10 split. The val/test tail
-        # is never walked, so predictions on it are genuine generalization.
-        n_kept = result.kept_commits
-        endpoint = result.split_boundary_index
-        if endpoint < 0:
-            endpoint = max(0, min(n_kept - 1, int(round(walk_fraction * n_kept)) - 1))
-
-        # ---- repo-state embedding @ commit 0 ----
-        _p("Embedding initial repository state…", 0.0)
-        repo_state_0, n_files = self._embed_repo_state(repo_dir, result.commits[0].sha)
-
-        # ---- walk the GRU, embedding diffs on the way ----
-        gru = self._gru
-        device = self.device
-        h = gru.init_hidden(torch.from_numpy(repo_state_0).to(device).unsqueeze(0))
-
-        hidden_states: List[np.ndarray] = []
-        ctx_norms: List[float] = []
-        walked_indices: List[int] = []
-        hidden_endpoint: Optional[torch.Tensor] = None
-
-        with torch.no_grad():
-            for ci in range(endpoint + 1):
-                commit = result.commits[ci]
-                _p(f"GRU streaming commit {ci + 1}/{endpoint + 1}…",
-                   (ci + 1) / max(endpoint + 1, 1))
-                diff_emb = self.embedder.embed_diff(commit.production_code_diff)
-                diff_t = torch.from_numpy(diff_emb).to(device).unsqueeze(0)
-                h = gru.step(diff_t, h)
-                h_last = h[-1].squeeze(0)                       # [H]
-                ctx = gru.output_norm(h[-1])                    # [1, H]
-                hidden_states.append(h_last.float().cpu().numpy())
-                ctx_norms.append(float(ctx.float().norm().cpu()))
-                walked_indices.append(ci)
-                if ci == endpoint:
-                    hidden_endpoint = h_last.detach()
-
-        return RepoSession(
-            repo_id=repo_id, result=result, walked_indices=walked_indices,
-            hidden_states=np.stack(hidden_states) if hidden_states else np.zeros((0, gru.hidden_dim)),
-            ctx_norms=ctx_norms, endpoint_index=endpoint,
-            repo_state_norm=float(np.linalg.norm(repo_state_0)),
-            n_repo_state_files=n_files, _hidden_endpoint=hidden_endpoint)
-
-    def _embed_repo_state(self, repo_dir: Path, commit_sha: str) -> Tuple[np.ndarray, int]:
-        files = ls_tree_py(repo_dir, commit_sha)
-        if len(files) > self.max_repo_state_files:
-            # Keep the largest files (most signal) under the cap.
-            files = sorted(files, key=lambda x: -x[2])[: self.max_repo_state_files]
-        texts: List[str] = []
-        with GitCatFileBatch(repo_dir) as cat:
-            for blob, _path, _size in files:
-                txt = cat.read_blob_text(blob)
-                if txt:
-                    texts.append(txt)
-        vec = self.embedder.embed_repo_state(texts)
-        return vec, len(texts)
-
-    # ---- adapter generation ----
-
-    def _ctx_at(self, session: RepoSession, commit_index: Optional[int]) -> torch.Tensor:
-        """output_norm(h_t) for the requested commit (defaults to endpoint)."""
-        if commit_index is None or commit_index == session.endpoint_index:
-            h_last = session._hidden_endpoint
-        else:
-            idx = max(0, min(commit_index, len(session.hidden_states) - 1))
-            h_last = torch.from_numpy(session.hidden_states[idx]).to(self.device)
-        h_last = h_last.unsqueeze(0)                  # [1, H]
-        return self._gru.output_norm(h_last)          # [1, H]
-
-    # ---- prediction ----
-
-    @torch.no_grad()
-    def predict(self, session: RepoSession, prefix: str, target: str,
-                commit_index: Optional[int] = None,
-                max_input_tokens: int = 4096,
-                max_new_tokens: int = 32) -> Dict:
-        """Predict the assertion RHS with the base model and with the adapter."""
-        tok = self._tokenizer
-        bos = getattr(tok, "bos_token_id", None)
-        if bos is None:
-            bos = getattr(tok, "eos_token_id", None)
-        ids = tok.encode(prefix, add_special_tokens=False)
-        if bos is not None:
-            ids = [bos] + ids
-        if len(ids) > max_input_tokens:
-            ids = ids[-max_input_tokens:]
-
-        # 1) base model (no adapter)
-        clear_all_lora_weights(self._base_model, self._specs)
-        base_pred = self._generate(ids, max_new_tokens)
-
-        # 2) adapted model
-        ctx = self._ctx_at(session, commit_index)
-        head_out = self._head(ctx)
-        inject_lora_weights(self._base_model, self._specs, head_out, batch_index=0)
-        adapted_pred = self._generate(ids, max_new_tokens)
-        clear_all_lora_weights(self._base_model, self._specs)
-
-        base_m = compute_metrics(base_pred, target)
-        adapted_m = compute_metrics(adapted_pred, target)
-        return {
-            "prefix": prefix,
-            "target": target,
-            "base_prediction": base_pred,
-            "adapted_prediction": adapted_pred,
-            "base_clean": base_m["pred_clean"],
-            "adapted_clean": adapted_m["pred_clean"],
-            "base_exact_match": bool(base_m["exact_match"]),
-            "adapted_exact_match": bool(adapted_m["exact_match"]),
-            "base_edit_similarity": float(base_m["edit_similarity"]),
-            "adapted_edit_similarity": float(adapted_m["edit_similarity"]),
-            "n_prompt_tokens": len(ids),
-        }
-
-    @torch.no_grad()
-    def _generate(self, input_ids: List[int], max_new_tokens: int) -> str:
-        device = self.device
-        ids = torch.tensor([input_ids], dtype=torch.long, device=device)
-        attn = torch.ones_like(ids)
-        out = self._base_model.generate(
-            input_ids=ids, attention_mask=attn,
-            max_new_tokens=int(max_new_tokens), do_sample=False,
-            pad_token_id=self._tokenizer.pad_token_id,
-            eos_token_id=self._tokenizer.eos_token_id, use_cache=True)
-        gen = out[0, ids.shape[1]:].tolist()
-        return self._tokenizer.decode(gen, skip_special_tokens=True)
-
-
-# ---------------------------------------------------------------------------
-# helpers
-# ---------------------------------------------------------------------------
-
-def _repo_id_from_url(url: str) -> str:
-    s = url.strip()
-    s = re.sub(r"\.git$", "", s)
-    s = re.sub(r"^git@[^:]+:", "", s)
-    s = re.sub(r"^(https?|ssh)://[^/]+/", "", s)
-    parts = [p for p in s.split("/") if p]
-    if len(parts) >= 2:
-        return f"{parts[-2]}/{parts[-1]}"
-    return parts[-1] if parts else "repo"
-
-
-def _repo_id_from_path(path: Path) -> str:
-    parts = path.resolve().parts
-    if len(parts) >= 2:
-        return f"{parts[-2]}/{parts[-1]}"
-    return path.name
-
-
-__all__ = ["Engine", "RepoSession", "DEFAULT_BASE_MODEL"]

c2l_demo/git_pipeline.py

@@ -1,498 +0,0 @@
-"""Git → commit-sequence pipeline.
-
-Faithful port of the repository-processing logic from
-``create_dataset/build_commit_assertion_db.py`` and
-``create_dataset/build_commit_parquet_db.py``:
-
-* stream first-parent commits oldest-first,
-* keep only commits that introduce *new or changed* assertions in test files
-  (this is the "GRU pass" subset),
-* build the ``production_code_diff`` (non-test source) between consecutive
-  kept commits,
-* split the kept commits chronologically 80/10/10 into train/val/test,
-* emit one QnA per new assertion.
-
-Everything runs against a local clone with full history.
-"""
-
-from __future__ import annotations
-
-import re
-import subprocess
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Callable, Dict, Iterator, List, Optional, Set, Tuple
-
-from .assertions import (
-    ExtractedAssertion,
-    extract_assertions_from_source,
-    normalize_for_id,
-)
-
-# ---------------------------------------------------------------------------
-# Constants (must match the dataset builders)
-# ---------------------------------------------------------------------------
-
-SOURCE_EXTS = {".py", ".md", ".rst"}
-SKIP_DIRS = {
-    ".git", "__pycache__", ".venv", "venv", "env",
-    "node_modules", "dist", "build", ".tox", ".mypy_cache",
-    "TEST_HYPERNET",
-}
-EMPTY_TREE_SHA = "4b825dc642cb6eb9a060e54bf8d69288fbee4904"
-_NULL_SHA = "0000000000000000000000000000000000000000"
-
-
-# ---------------------------------------------------------------------------
-# Path filtering
-# ---------------------------------------------------------------------------
-
-def _is_test_path(path: str) -> bool:
-    for part in Path(path).parts:
-        if "test" in part.lower():
-            return True
-    return False
-
-
-def _is_python_test_file(path: str) -> bool:
-    p = Path(path)
-    if p.suffix.lower() != ".py":
-        return False
-    if any(part in SKIP_DIRS for part in p.parts):
-        return False
-    return _is_test_path(path)
-
-
-def _should_skip_diff_path(path: str) -> bool:
-    p = Path(path)
-    if p.suffix.lower() not in SOURCE_EXTS:
-        return True
-    if any(part in SKIP_DIRS for part in p.parts):
-        return True
-    if _is_test_path(path):
-        return True
-    return False
-
-
-# ---------------------------------------------------------------------------
-# Git helpers
-# ---------------------------------------------------------------------------
-
-def _run_git(repo_dir: Path, args: List[str], timeout: int = 300) -> Optional[str]:
-    try:
-        result = subprocess.run(
-            ["git", "-C", str(repo_dir)] + args,
-            capture_output=True, timeout=timeout,
-        )
-        if result.returncode != 0:
-            return None
-        return result.stdout.decode("utf-8", errors="replace")
-    except (subprocess.TimeoutExpired, FileNotFoundError, OSError):
-        return None
-
-
-@dataclass
-class RawFileChange:
-    status: str
-    new_blob: str
-    path: str
-    old_path: Optional[str] = None
-
-
-@dataclass
-class RawCommit:
-    sha: str
-    parent_sha: Optional[str]
-    timestamp: str
-    changes: List[RawFileChange]
-
-
-def stream_commit_changes(repo_dir: Path) -> Iterator[RawCommit]:
-    """Yield one ``RawCommit`` per first-parent commit, oldest-first."""
-    proc = subprocess.Popen(
-        [
-            "git", "-C", str(repo_dir),
-            "log", "--first-parent", "--reverse",
-            "--raw", "--no-abbrev", "-M",
-            "--format=COMMITHDR %H %P %aI",
-        ],
-        stdout=subprocess.PIPE,
-        stderr=subprocess.PIPE,
-        bufsize=1 << 20,
-    )
-    assert proc.stdout is not None
-    cur: Optional[RawCommit] = None
-    try:
-        for raw in proc.stdout:
-            line = raw.decode("utf-8", errors="replace").rstrip("\n")
-            if not line:
-                continue
-            if line.startswith("COMMITHDR "):
-                if cur is not None:
-                    yield cur
-                tokens = line.split()
-                if len(tokens) < 3:
-                    cur = None
-                    continue
-                sha = tokens[1]
-                timestamp = tokens[-1]
-                parents = tokens[2:-1]
-                parent_sha = parents[0] if parents else None
-                cur = RawCommit(sha=sha, parent_sha=parent_sha,
-                                timestamp=timestamp, changes=[])
-                continue
-            if cur is None or not line.startswith(":"):
-                continue
-            try:
-                meta, _, tail = line.partition("\t")
-                if not tail:
-                    continue
-                fields = meta.split()
-                if len(fields) < 5:
-                    continue
-                new_blob = fields[3].lstrip(":")
-                status = fields[4]
-                parts = tail.split("\t")
-                if status.startswith("R") or status.startswith("C"):
-                    old_path = parts[0]
-                    new_path = parts[1] if len(parts) > 1 else parts[0]
-                    cur.changes.append(RawFileChange(
-                        status=status, new_blob=new_blob,
-                        path=new_path, old_path=old_path))
-                else:
-                    cur.changes.append(RawFileChange(
-                        status=status, new_blob=new_blob, path=parts[0]))
-            except Exception:
-                continue
-        if cur is not None:
-            yield cur
-    finally:
-        try:
-            proc.stdout.close()
-        finally:
-            proc.wait(timeout=5)
-
-
-class GitCatFileBatch:
-    """Persistent ``git cat-file --batch`` for cheap blob reads."""
-
-    def __init__(self, repo_dir: Path):
-        self._proc = subprocess.Popen(
-            ["git", "-C", str(repo_dir), "cat-file", "--batch"],
-            stdin=subprocess.PIPE, stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE, bufsize=0,
-        )
-
-    def read_blob(self, blob_sha: str) -> Optional[bytes]:
-        if not blob_sha or blob_sha == _NULL_SHA:
-            return None
-        p = self._proc
-        assert p.stdin is not None and p.stdout is not None
-        try:
-            p.stdin.write((blob_sha + "\n").encode("ascii"))
-            p.stdin.flush()
-        except (BrokenPipeError, OSError):
-            return None
-        header = p.stdout.readline()
-        if not header:
-            return None
-        parts = header.decode("utf-8", errors="replace").strip().split()
-        if len(parts) < 3 or parts[1] != "blob":
-            return None
-        try:
-            size = int(parts[2])
-        except ValueError:
-            return None
-        body = b""
-        remaining = size
-        while remaining > 0:
-            chunk = p.stdout.read(remaining)
-            if not chunk:
-                break
-            body += chunk
-            remaining -= len(chunk)
-        try:
-            p.stdout.read(1)
-        except Exception:
-            pass
-        return body
-
-    def read_blob_text(self, blob_sha: str) -> Optional[str]:
-        data = self.read_blob(blob_sha)
-        return None if data is None else data.decode("utf-8", errors="replace")
-
-    def close(self) -> None:
-        try:
-            if self._proc.stdin is not None:
-                self._proc.stdin.close()
-        except Exception:
-            pass
-        try:
-            self._proc.wait(timeout=5)
-        except Exception:
-            try:
-                self._proc.kill()
-            except Exception:
-                pass
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, *exc):
-        self.close()
-
-
-def get_diff(repo_dir: Path, parent_sha: Optional[str], child_sha: str,
-             timeout: int = 600) -> str:
-    base = parent_sha if parent_sha else EMPTY_TREE_SHA
-    out = _run_git(repo_dir, ["diff", "--no-color", "-U3", base, child_sha],
-                   timeout=timeout)
-    return out if out else ""
-
-
-_DIFF_HEADER_RE = re.compile(r"^diff --git a/(.+?) b/(.+?)(?:\n|$)", re.MULTILINE)
-
-
-def filter_diff(raw_diff: str) -> str:
-    if not raw_diff:
-        return ""
-    chunks = re.split(r"(?=^diff --git )", raw_diff, flags=re.MULTILINE)
-    kept: List[str] = []
-    for chunk in chunks:
-        if not chunk.startswith("diff --git "):
-            continue
-        m = _DIFF_HEADER_RE.match(chunk)
-        if not m:
-            continue
-        if _should_skip_diff_path(m.group(2)):
-            continue
-        kept.append(chunk)
-    return "".join(kept)
-
-
-def ls_tree_py(repo_dir: Path, commit_sha: str,
-               max_file_bytes: int = 2_000_000) -> List[Tuple[str, str, int]]:
-    """Return [(blob_sha, path, size)] for tracked .py files at a commit."""
-    out = _run_git(repo_dir, ["ls-tree", "-r", "-l", commit_sha], timeout=120)
-    if not out:
-        return []
-    rows: List[Tuple[str, str, int]] = []
-    for ln in out.splitlines():
-        try:
-            meta, fpath = ln.split("\t", 1)
-            _mode, ftype, blob, size_str = meta.split()
-        except ValueError:
-            continue
-        if ftype != "blob" or not fpath.endswith(".py"):
-            continue
-        try:
-            size = int(size_str)
-        except ValueError:
-            continue
-        if size <= 0 or size > max_file_bytes:
-            continue
-        if any(part in SKIP_DIRS for part in Path(fpath).parts):
-            continue
-        rows.append((blob, fpath, size))
-    return rows
-
-
-# ---------------------------------------------------------------------------
-# Chronological split
-# ---------------------------------------------------------------------------
-
-def chronological_splits(n_kept: int) -> List[str]:
-    """80/10/10 chronological split of ``n_kept`` kept commits."""
-    if n_kept <= 0:
-        return []
-    if n_kept < 3:
-        return ["train"] * n_kept
-    n_train = max(1, int(n_kept * 0.8))
-    n_val = max(1, int(n_kept * 0.1))
-    n_test = n_kept - n_train - n_val
-    if n_test < 1:
-        n_val = max(1, n_val - 1) if n_val > 1 else n_val
-        n_test = n_kept - n_train - n_val
-        if n_test < 1:
-            n_train = max(1, n_train - 1)
-            n_test = n_kept - n_train - n_val
-    return ["train"] * n_train + ["val"] * n_val + ["test"] * n_test
-
-
-# ---------------------------------------------------------------------------
-# Per-repo processing
-# ---------------------------------------------------------------------------
-
-@dataclass
-class QnaItem:
-    test_file: str
-    test_function: str
-    assertion_type: str
-    prefix: str
-    target: str
-    lineno: int
-    col_offset: int
-
-
-@dataclass
-class KeptCommit:
-    commit_index: int          # index among kept commits (0-based)
-    original_index: int        # index in the full first-parent history
-    sha: str
-    timestamp: str
-    in_repo_split: str         # train / val / test
-    production_code_diff: str
-    n_new_assertions: int
-    qnas: List[QnaItem] = field(default_factory=list)
-
-
-@dataclass
-class RepoProcessResult:
-    repo_id: str
-    total_commits: int            # all first-parent commits
-    test_touching_commits: int    # commits that changed a test file
-    kept_commits: int             # commits selected for the GRU pass
-    n_assertions: int             # total QnA pairs emitted
-    n_train_kept: int
-    n_val_kept: int
-    n_test_kept: int
-    commits: List[KeptCommit]
-
-    @property
-    def split_boundary_index(self) -> int:
-        """Index (within kept commits) of the last train commit (the 80% point)."""
-        last_train = -1
-        for c in self.commits:
-            if c.in_repo_split == "train":
-                last_train = c.commit_index
-        return last_train
-
-
-ProgressFn = Optional[Callable[[str, float], None]]
-
-
-def process_repo(repo_dir: Path, repo_id: str = "",
-                 diff_timeout: int = 600,
-                 progress: ProgressFn = None) -> RepoProcessResult:
-    """Process a local clone into the kept-commit sequence + QnAs + stats."""
-    prev_assert_keys: Dict[str, Set[Tuple[str, str, str]]] = {}
-    kept: List[Dict] = []
-    total_commits = 0
-    test_touching = 0
-
-    def _report(msg: str, frac: float) -> None:
-        if progress is not None:
-            progress(msg, frac)
-
-    # ---- Phase 1: detect kept commits ----
-    with GitCatFileBatch(repo_dir) as cat:
-        for ci, raw in enumerate(stream_commit_changes(repo_dir)):
-            total_commits += 1
-            if ci % 50 == 0:
-                _report(f"Scanning commit history… ({total_commits} commits)", 0.0)
-
-            candidate_changes: List[RawFileChange] = []
-            for ch in raw.changes:
-                if ch.old_path and ch.path != ch.old_path:
-                    if ch.old_path in prev_assert_keys:
-                        prev_assert_keys[ch.path] = prev_assert_keys.pop(ch.old_path)
-                if ch.status.startswith("D"):
-                    prev_assert_keys.pop(ch.path, None)
-                    continue
-                if _is_python_test_file(ch.path):
-                    candidate_changes.append(ch)
-
-            if not candidate_changes:
-                continue
-            test_touching += 1
-
-            per_file_new: List[Tuple[str, List[ExtractedAssertion]]] = []
-            for ch in candidate_changes:
-                content = cat.read_blob_text(ch.new_blob)
-                if content is None:
-                    prev_assert_keys.pop(ch.path, None)
-                    continue
-                extracted = extract_assertions_from_source(content)
-                current_keys: Set[Tuple[str, str, str]] = set()
-                for ex in extracted:
-                    current_keys.add((
-                        ex.assertion_type,
-                        normalize_for_id(ex.prefix),
-                        normalize_for_id(ex.target),
-                    ))
-                prev_keys = prev_assert_keys.get(ch.path, set())
-                new_keys = current_keys - prev_keys
-                if new_keys:
-                    new_extracted = [
-                        ex for ex in extracted
-                        if (ex.assertion_type, normalize_for_id(ex.prefix),
-                            normalize_for_id(ex.target)) in new_keys
-                    ]
-                    if new_extracted:
-                        per_file_new.append((ch.path, new_extracted))
-                prev_assert_keys[ch.path] = current_keys
-
-            if not per_file_new:
-                continue
-            kept.append({
-                "original_idx": ci,
-                "sha": raw.sha,
-                "timestamp": raw.timestamp,
-                "per_file_new": per_file_new,
-            })
-
-    if not kept:
-        return RepoProcessResult(
-            repo_id=repo_id, total_commits=total_commits,
-            test_touching_commits=test_touching, kept_commits=0,
-            n_assertions=0, n_train_kept=0, n_val_kept=0, n_test_kept=0,
-            commits=[])
-
-    # ---- Phase 2: build diffs + chronological split ----
-    splits = chronological_splits(len(kept))
-    commits: List[KeptCommit] = []
-    prev_kept_sha: Optional[str] = None
-    n_assertions = 0
-    for ki, kc in enumerate(kept):
-        _report(f"Building production-code diffs… ({ki + 1}/{len(kept)} kept commits)",
-                (ki + 1) / max(len(kept), 1))
-        sha = kc["sha"]
-        raw_diff = get_diff(repo_dir, prev_kept_sha, sha, timeout=diff_timeout)
-        filtered = filter_diff(raw_diff) if raw_diff else ""
-
-        qnas: List[QnaItem] = []
-        for tfile, extracted in kc["per_file_new"]:
-            for ex in extracted:
-                qnas.append(QnaItem(
-                    test_file=tfile, test_function=ex.test_function,
-                    assertion_type=ex.assertion_type, prefix=ex.prefix,
-                    target=ex.target, lineno=ex.lineno, col_offset=ex.col_offset))
-        n_assertions += len(qnas)
-
-        commits.append(KeptCommit(
-            commit_index=ki, original_index=kc["original_idx"], sha=sha,
-            timestamp=kc.get("timestamp", ""), in_repo_split=splits[ki],
-            production_code_diff=filtered, n_new_assertions=len(qnas), qnas=qnas))
-        prev_kept_sha = sha
-
-    n_train = sum(1 for s in splits if s == "train")
-    n_val = sum(1 for s in splits if s == "val")
-    n_test = sum(1 for s in splits if s == "test")
-    return RepoProcessResult(
-        repo_id=repo_id, total_commits=total_commits,
-        test_touching_commits=test_touching, kept_commits=len(kept),
-        n_assertions=n_assertions, n_train_kept=n_train, n_val_kept=n_val,
-        n_test_kept=n_test, commits=commits)
-
-
-__all__ = [
-    "process_repo",
-    "RepoProcessResult",
-    "KeptCommit",
-    "QnaItem",
-    "ls_tree_py",
-    "GitCatFileBatch",
-    "get_diff",
-    "filter_diff",
-    "chronological_splits",
-]

c2l_demo/viz.py

@@ -1,192 +0,0 @@
-"""Plotting helpers for the demo (matplotlib, Agg backend).
-
-Three figures:
-
-* :func:`plot_commit_timeline` -- the full commit history with the kept
-  "GRU-pass" commits highlighted, dot size proportional to the number of new
-  assertions, colored by the chronological 80/10/10 split.
-* :func:`plot_gru_processing` -- the GRU hidden-state trajectory (2-D PCA) as
-  it streams through the repository, plus the context-norm curve, marking the
-  adapter (80%) point.
-* :func:`plot_assertion_type_breakdown` -- distribution of assertion kinds.
-"""
-
-from __future__ import annotations
-
-from datetime import datetime
-from typing import List, Optional
-
-import matplotlib
-
-matplotlib.use("Agg")
-import matplotlib.dates as mdates  # noqa: E402
-import matplotlib.pyplot as plt  # noqa: E402
-import numpy as np  # noqa: E402
-
-# Modern, conference-grade palette
-_BG = "#0f1117"
-_FG = "#e6e6e6"
-_GRID = "#2a2e3a"
-_TRAIN = "#4f8cff"
-_VAL = "#ffcc4d"
-_TEST = "#ff6b6b"
-_ACCENT = "#22d3a6"
-
-_SPLIT_COLORS = {"train": _TRAIN, "val": _VAL, "test": _TEST}
-
-
-def _style(ax):
-    ax.set_facecolor(_BG)
-    for spine in ax.spines.values():
-        spine.set_color(_GRID)
-    ax.tick_params(colors=_FG, labelsize=9)
-    ax.yaxis.label.set_color(_FG)
-    ax.xaxis.label.set_color(_FG)
-    ax.title.set_color(_FG)
-    ax.grid(True, color=_GRID, linewidth=0.6, alpha=0.6)
-
-
-def _parse_ts(ts: str) -> Optional[datetime]:
-    if not ts:
-        return None
-    try:
-        return datetime.fromisoformat(ts.replace("Z", "+00:00"))
-    except ValueError:
-        return None
-
-
-def plot_commit_timeline(result, endpoint_index: int):
-    """Scatter of kept commits over real time; size = #new assertions."""
-    commits = result.commits
-    fig, ax = plt.subplots(figsize=(10, 4.2), dpi=120)
-    fig.patch.set_facecolor(_BG)
-    _style(ax)
-
-    xs, ys, sizes, colors = [], [], [], []
-    parsed = [_parse_ts(c.timestamp) for c in commits]
-    distinct_dates = {p for p in parsed if p is not None}
-    use_dates = (all(p is not None for p in parsed) and len(distinct_dates) > 1)
-    for i, c in enumerate(commits):
-        x = parsed[i] if use_dates else c.commit_index
-        xs.append(x)
-        ys.append(c.n_new_assertions)
-        sizes.append(20 + 22 * np.sqrt(max(c.n_new_assertions, 0)))
-        colors.append(_SPLIT_COLORS.get(c.in_repo_split, _FG))
-
-    ax.scatter(xs, ys, s=sizes, c=colors, alpha=0.85, edgecolors="black",
-               linewidths=0.4, zorder=3)
-
-    # Mark the adapter (80%) boundary.
-    if 0 <= endpoint_index < len(commits):
-        bx = xs[endpoint_index]
-        ax.axvline(bx, color=_ACCENT, linestyle="--", linewidth=1.6, alpha=0.9,
-                   zorder=2)
-        ax.text(bx, ax.get_ylim()[1] * 0.92, "  adapter @ 80%",
-                color=_ACCENT, fontsize=9, va="top", ha="left")
-
-    if use_dates:
-        ax.xaxis.set_major_formatter(mdates.DateFormatter("%Y-%m"))
-        fig.autofmt_xdate(rotation=25)
-        ax.set_xlabel("commit date")
-    else:
-        ax.set_xlabel("kept-commit index")
-    ax.set_ylabel("new assertions in commit")
-    ax.set_title(f"{result.repo_id}  ·  {result.kept_commits} GRU-pass commits "
-                 f"· {result.n_assertions} assertions", fontsize=11)
-
-    handles = [
-        plt.Line2D([0], [0], marker="o", color=_BG, markerfacecolor=col,
-                   markersize=8, label=lab, markeredgecolor="black")
-        for lab, col in (("train (≤80%)", _TRAIN), ("val", _VAL), ("test", _TEST))
-    ]
-    leg = ax.legend(handles=handles, loc="upper left", fontsize=8,
-                    facecolor=_BG, edgecolor=_GRID, labelcolor=_FG)
-    leg.get_frame().set_alpha(0.7)
-    fig.tight_layout()
-    return fig
-
-
-def plot_gru_processing(session):
-    """2-D PCA trajectory of the GRU hidden state + context-norm curve."""
-    H = session.hidden_states  # [T, H]
-    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10.5, 4.3), dpi=120,
-                                   gridspec_kw={"width_ratios": [1.25, 1]})
-    fig.patch.set_facecolor(_BG)
-    _style(ax1)
-    _style(ax2)
-
-    T = H.shape[0]
-    if T >= 2:
-        Hc = H - H.mean(axis=0, keepdims=True)
-        try:
-            _, _, vt = np.linalg.svd(Hc, full_matrices=False)
-            proj = Hc @ vt[:2].T  # [T, 2]
-        except np.linalg.LinAlgError:
-            proj = Hc[:, :2]
-        order = np.arange(T)
-        ax1.plot(proj[:, 0], proj[:, 1], color=_GRID, linewidth=1.0, zorder=1,
-                 alpha=0.8)
-        sc = ax1.scatter(proj[:, 0], proj[:, 1], c=order, cmap="viridis",
-                         s=34, zorder=3, edgecolors="black", linewidths=0.3)
-        ax1.scatter(proj[0, 0], proj[0, 1], marker="o", s=130,
-                    facecolors="none", edgecolors=_FG, linewidths=1.6,
-                    zorder=4, label="commit 0")
-        ax1.scatter(proj[-1, 0], proj[-1, 1], marker="*", s=240, color=_ACCENT,
-                    edgecolors="black", linewidths=0.5, zorder=5,
-                    label="adapter point")
-        cbar = fig.colorbar(sc, ax=ax1, fraction=0.046, pad=0.04)
-        cbar.set_label("commit order", color=_FG, fontsize=8)
-        cbar.ax.tick_params(colors=_FG, labelsize=8)
-        ax1.legend(loc="best", fontsize=8, facecolor=_BG, edgecolor=_GRID,
-                   labelcolor=_FG)
-    else:
-        ax1.text(0.5, 0.5, "Not enough commits\nfor a trajectory",
-                 color=_FG, ha="center", va="center", transform=ax1.transAxes)
-    ax1.set_title("GRU hidden state walking the repository (PCA)", fontsize=10.5)
-    ax1.set_xlabel("PC 1")
-    ax1.set_ylabel("PC 2")
-
-    # Context-norm curve.
-    if session.ctx_norms:
-        ax2.plot(session.walked_indices, session.ctx_norms, color=_ACCENT,
-                 linewidth=1.8)
-        ax2.fill_between(session.walked_indices, session.ctx_norms,
-                         color=_ACCENT, alpha=0.12)
-    ax2.set_title("repository knowledge accumulating", fontsize=10.5)
-    ax2.set_xlabel("kept-commit index")
-    ax2.set_ylabel("‖adapter context‖")
-    fig.tight_layout()
-    return fig
-
-
-def plot_assertion_type_breakdown(result):
-    """Horizontal bar chart of assertion kinds across all QnAs."""
-    counts: dict = {}
-    for c in result.commits:
-        for q in c.qnas:
-            counts[q.assertion_type] = counts.get(q.assertion_type, 0) + 1
-    fig, ax = plt.subplots(figsize=(10.5, max(2.2, 0.42 * max(len(counts), 1))),
-                           dpi=120)
-    fig.patch.set_facecolor(_BG)
-    _style(ax)
-    if not counts:
-        ax.text(0.5, 0.5, "No assertions", color=_FG, ha="center",
-                va="center", transform=ax.transAxes)
-        return fig
-    items = sorted(counts.items(), key=lambda kv: kv[1])
-    labels = [k for k, _ in items]
-    vals = [v for _, v in items]
-    ax.barh(labels, vals, color=_TRAIN, edgecolor="black", linewidth=0.4)
-    for i, v in enumerate(vals):
-        ax.text(v, i, f" {v}", color=_FG, va="center", fontsize=8)
-    ax.set_title("assertion types extracted from the test suite", fontsize=10.5)
-    ax.set_xlabel("count")
-    fig.tight_layout()
-    return fig
-
-
-__all__ = [
-    "plot_commit_timeline",
-    "plot_gru_processing",
-    "plot_assertion_type_breakdown",
-]

requirements.txt

@@ -1,9 +1,7 @@
-gradio>=4.44,<6
+--extra-index-url https://download.pytorch.org/whl/cpu
 torch>=2.3
-transformers>=4.45
+transformers>=4.45,<5
 accelerate>=0.34
-huggingface_hub>=0.25
+huggingface_hub>=0.28
 numpy>=1.24
 matplotlib>=3.7
-# optional, improves CodeBLEU fidelity if present:
-# codebleu>=0.7

tools/build_space_data.py

@@ -0,0 +1,92 @@
+#!/usr/bin/env python3
+"""Package the precompute cache into the small files the Space reads.
+
+    python hf_space_demo/tools/build_space_data.py \
+        --cache-dir /scratch/lhotsko/C2L_SPACE_CACHE --out-dir hf_space_demo/data
+
+Writes ``index.json`` plus, per selected repository, ``<safe>.json.gz``
+(commits + assertions + cached predictions) and ``<safe>.ctx.npy``
+(float32 [n_commits, 2048] GRU contexts).
+"""
+
+from __future__ import annotations
+
+import argparse
+import gzip
+import json
+import pickle
+import urllib.request
+from pathlib import Path
+
+import numpy as np
+
+
+def github_meta(repo_id: str) -> dict:
+    try:
+        with urllib.request.urlopen(f"https://api.github.com/repos/{repo_id}", timeout=10) as r:
+            d = json.load(r)
+        return {"description": d.get("description") or "", "stars": d.get("stargazers_count")}
+    except Exception:
+        return {"description": "", "stars": None}
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--cache-dir", required=True)
+    ap.add_argument("--out-dir", required=True)
+    ap.add_argument("--min-commits", type=int, default=10)
+    ap.add_argument("--max-commits", type=int, default=400)
+    ap.add_argument("--max-repos", type=int, default=24)
+    args = ap.parse_args()
+
+    out = Path(args.out_dir)
+    out.mkdir(parents=True, exist_ok=True)
+    candidates = []
+    for p in sorted(Path(args.cache_dir).glob("*.pkl")):
+        with open(p, "rb") as f:
+            d = pickle.load(f)
+        n_c, qn = len(d["commits"]), d["qnas"]
+        if not (args.min_commits <= n_c <= args.max_commits) or not qn:
+            continue
+        base = sum(q["base_em"] for q in qn) / len(qn)
+        c2l = sum(q["c2l_em"] for q in qn) / len(qn)
+        tt = sum(1 for q in qn if q.get("prev_em") == 0.0)
+        candidates.append((c2l - base, tt, p, d, base, c2l))
+
+    candidates.sort(key=lambda x: (x[0], x[1]), reverse=True)
+    index = []
+    for gain, tt, p, d, base, c2l in candidates[: args.max_repos]:
+        safe = p.stem
+        commits = d["commits"]
+        per_commit = [[] for _ in commits]
+        for q in d["qnas"]:
+            per_commit[q["k"]].append(q)
+        for k, c in enumerate(commits):
+            qs = per_commit[k]
+            c["n_qnas"] = len(qs)
+            c["base_em"] = sum(q["base_em"] for q in qs) / len(qs) if qs else None
+            c["c2l_em"] = sum(q["c2l_em"] for q in qs) / len(qs) if qs else None
+            c["time_travel"] = any(q.get("prev_em") == 0.0 for q in qs)
+        payload = {"repo_id": d["repo_id"], "commits": commits, "qnas": per_commit,
+                   "max_input_tokens": d["max_input_tokens"],
+                   "max_new_tokens": d["max_new_tokens"]}
+        with gzip.open(out / f"{safe}.json.gz", "wt", encoding="utf-8") as f:
+            json.dump(payload, f)
+        np.save(out / f"{safe}.ctx.npy", d["ctx"].astype(np.float32))
+        dates = [c.get("timestamp") for c in commits if c.get("timestamp")]
+        index.append({
+            "repo_id": d["repo_id"], "safe": safe, "n_commits": len(commits),
+            "n_qnas": len(d["qnas"]), "base_em": base, "c2l_em": c2l,
+            "n_time_travel": tt, "first_date": min(dates) if dates else "",
+            "last_date": max(dates) if dates else "", **github_meta(d["repo_id"]),
+        })
+        print(f"{d['repo_id']:45s} commits={len(commits):4d} base={base:.3f} "
+              f"c2l={c2l:.3f} gain={gain:+.3f} time_travel={tt}")
+
+    index.sort(key=lambda r: r["repo_id"].lower())
+    (out / "index.json").write_text(json.dumps(index, indent=1))
+    print(f"wrote {len(index)} repos to {out}")
+
+
+if __name__ == "__main__":
+    main()

tools/precompute.py

@@ -0,0 +1,209 @@
+#!/usr/bin/env python3
+"""Precompute everything the Code2LoRA timeline Space serves from cache.
+
+For every cross-repo test repository (held out from training) this walks the
+Code2LoRA-Evo GRU over the precomputed per-commit embeddings and records:
+
+* ``ctx``  -- output_norm(h_t) after every commit (what the head consumes),
+* for every QnA attached to a commit (first <=8, the paper's eval cap):
+  - the base-model prediction (no adapter),
+  - the prediction with the adapter generated at that commit (paper protocol),
+  - for QnAs the adapter solves and the base model does not, the prediction
+    with the adapter from the *previous* commit ("time travel").
+
+One pickle per repository is written to ``--out-dir``; existing files are
+skipped, so the job is resumable. Run from the RepoPeftData root:
+
+    python hf_space_demo/tools/precompute.py --out-dir $SCRATCH/C2L_SPACE_CACHE \
+        --shard-i 0 --num-shards 8
+"""
+
+from __future__ import annotations
+
+import argparse
+import pickle
+import sys
+import time
+from collections import defaultdict
+from pathlib import Path
+from typing import Dict, List, Optional
+
+import numpy as np
+import pyarrow.parquet as pq
+import torch
+
+_ROOT = Path(__file__).resolve().parents[2]
+for p in (_ROOT, _ROOT / "hypernetwork"):
+    if str(p) not in sys.path:
+        sys.path.insert(0, str(p))
+
+from code2lora_core import (  # noqa: E402
+    LoRA,
+    inject_lora_weights,
+    load_commit_rows_for_gru,
+    load_qna_rows,
+)
+from evaluation.metrics import compute_metrics  # noqa: E402
+from evaluation.run_code2lora_gru_v2_eval import (  # noqa: E402
+    DEFAULT_BASE_MODEL,
+    DEFAULT_TARGET_MODULES,
+    _generate_batch,
+    _get_bos_id,
+    _load_v2_gru_ckpt,
+    _prepare_prefix_ids,
+)
+
+COMMITS = "/scratch/lhotsko/REPO_DATASET/commit_parquet_hf_v2/commits/cr_test.parquet"
+QNAS = "/scratch/lhotsko/REPO_DATASET/code2lora_snapshots_hf/qna/cr_test.parquet"
+CKPT = "/scratch/lhotsko/hf_staging/code2lora_gru/gru_head.best.pt"
+
+
+def clear_all_lora_weights(model: torch.nn.Module, specs) -> None:
+    named = dict(model.named_modules())
+    for sp in specs:
+        mod = named[sp.full_name]
+        if isinstance(mod, LoRA):
+            mod.A, mod.B = None, None
+
+
+def safe_name(repo_id: str) -> str:
+    return repo_id.replace("/", "__")
+
+
+def commit_metadata(repo_ids: List[str]) -> Dict[str, Dict[str, dict]]:
+    t = pq.read_table(COMMITS, columns=[
+        "repo_id", "commit_sha", "commit_index", "commit_timestamp",
+        "n_new_assertions", "production_code_diff"],
+        filters=[("repo_id", "in", repo_ids)])
+    out: Dict[str, Dict[str, dict]] = defaultdict(dict)
+    for r in t.to_pylist():
+        diff = r["production_code_diff"] or ""
+        out[r["repo_id"]][r["commit_sha"]] = {
+            "timestamp": r["commit_timestamp"],
+            "n_new_assertions": int(r["n_new_assertions"] or 0),
+            "diff_preview": diff[:6000],
+            "diff_chars": len(diff),
+        }
+    return out
+
+
+def score(pred: str, target: str) -> Dict[str, float]:
+    m = compute_metrics(pred, target)
+    return {"em": float(bool(m["exact_match"])), "es": float(m["edit_similarity"])}
+
+
+@torch.no_grad()
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--out-dir", required=True)
+    ap.add_argument("--checkpoint", default=CKPT)
+    ap.add_argument("--shard-i", type=int, default=0)
+    ap.add_argument("--num-shards", type=int, default=1)
+    ap.add_argument("--max-input-tokens", type=int, default=4096)
+    ap.add_argument("--max-new-tokens", type=int, default=64)
+    ap.add_argument("--batch-size", type=int, default=8)
+    ap.add_argument("--qnas-per-commit-limit", type=int, default=8)
+    ap.add_argument("--repos", nargs="*", default=None)
+    args = ap.parse_args()
+
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    rows_by_repo = load_commit_rows_for_gru(Path(COMMITS))
+    repos = sorted(rows_by_repo)
+    if args.repos:
+        repos = [r for r in repos if r in set(args.repos)]
+    repos = [r for i, r in enumerate(repos) if i % args.num_shards == args.shard_i]
+    repos = [r for r in repos if not (out_dir / f"{safe_name(r)}.pkl").exists()]
+    print(f"[shard {args.shard_i}/{args.num_shards}] {len(repos)} repos to do: {repos}", flush=True)
+    if not repos:
+        return
+
+    base_model, gru, head, specs, tok = _load_v2_gru_ckpt(
+        Path(args.checkpoint), DEFAULT_BASE_MODEL, DEFAULT_TARGET_MODULES, device)
+    bos = _get_bos_id(tok)
+    meta = commit_metadata(repos)
+
+    qnas_by_key: Dict[tuple, List[dict]] = defaultdict(list)
+    for q in load_qna_rows(Path(QNAS), repo_ids=repos):
+        qnas_by_key[(q.repo_id, q.commit_sha)].append(q)
+
+    def gen(ids_list: List[List[int]]) -> List[str]:
+        out: List[str] = []
+        for i in range(0, len(ids_list), args.batch_size):
+            out += _generate_batch(base_model, tok, device, ids_list[i:i + args.batch_size],
+                                   max_new_tokens=args.max_new_tokens)
+        return out
+
+    for repo in repos:
+        t0 = time.time()
+        rows = rows_by_repo[repo]
+        h = gru.init_hidden(torch.from_numpy(rows[0].repo_state_embedding).to(device).unsqueeze(0))
+        ctxs = []
+        for row in rows:
+            h = gru.step(torch.from_numpy(row.diff_embedding).to(device).unsqueeze(0), h)
+            ctxs.append(gru.output_norm(h[-1]).float())
+        ctx = torch.cat(ctxs)                                   # [T, 2048]
+
+        commits, qnas = [], []
+        for k, row in enumerate(rows):
+            m = meta[repo].get(row.commit_sha, {})
+            commits.append({"sha": row.commit_sha, "commit_index": int(row.commit_index),
+                            "in_repo_split": row.in_repo_split, **m})
+            for q in qnas_by_key.get((repo, row.commit_sha), [])[: args.qnas_per_commit_limit]:
+                ids = _prepare_prefix_ids(tok, q.prefix, args.max_input_tokens, bos)
+                had_bos = bos is not None and ids[0] == bos
+                qnas.append({
+                    "k": k, "test_file": q.test_file, "lineno": int(q.lineno or 0),
+                    "test_function": q.test_function or "",
+                    "prompt": tok.decode(ids[1:] if had_bos else ids),
+                    "had_bos": had_bos, "n_prompt_tokens": len(ids),
+                    "target": q.target, "_ids": ids,
+                })
+
+        clear_all_lora_weights(base_model, specs)
+        for q, p in zip(qnas, gen([q["_ids"] for q in qnas])):
+            q["base_pred"] = p
+            q.update({f"base_{k}": v for k, v in score(p, q["target"]).items()})
+
+        by_k: Dict[int, List[dict]] = defaultdict(list)
+        for q in qnas:
+            by_k[q["k"]].append(q)
+        for k, group in by_k.items():
+            inject_lora_weights(base_model, specs, head(ctx[k:k + 1]), batch_index=0)
+            for q, p in zip(group, gen([q["_ids"] for q in group])):
+                q["c2l_pred"] = p
+                q.update({f"c2l_{kk}": v for kk, v in score(p, q["target"]).items()})
+
+        wins = [q for q in qnas if q["c2l_em"] == 1.0 and q["base_em"] == 0.0 and q["k"] > 0]
+        by_prev: Dict[int, List[dict]] = defaultdict(list)
+        for q in wins:
+            by_prev[q["k"] - 1].append(q)
+        for k_prev, group in by_prev.items():
+            inject_lora_weights(base_model, specs, head(ctx[k_prev:k_prev + 1]), batch_index=0)
+            for q, p in zip(group, gen([q["_ids"] for q in group])):
+                q["prev_pred"] = p
+                q.update({f"prev_{kk}": v for kk, v in score(p, q["target"]).items()})
+        clear_all_lora_weights(base_model, specs)
+
+        for q in qnas:
+            q.pop("_ids")
+        payload = {"repo_id": repo, "ctx": ctx.cpu().numpy().astype(np.float32),
+                   "commits": commits, "qnas": qnas,
+                   "max_input_tokens": args.max_input_tokens,
+                   "max_new_tokens": args.max_new_tokens}
+        tmp = out_dir / f"{safe_name(repo)}.pkl.tmp"
+        with open(tmp, "wb") as f:
+            pickle.dump(payload, f)
+        tmp.replace(out_dir / f"{safe_name(repo)}.pkl")
+        n = max(len(qnas), 1)
+        print(f"[{repo}] {len(rows)} commits, {len(qnas)} qnas, "
+              f"base EM={sum(q['base_em'] for q in qnas)/n:.3f} "
+              f"c2l EM={sum(q['c2l_em'] for q in qnas)/n:.3f} "
+              f"time-travel={sum(1 for q in wins if q.get('prev_em') == 0.0)} "
+              f"({(time.time()-t0)/60:.1f} min)", flush=True)
+
+
+if __name__ == "__main__":
+    main()

tools/precompute.sbatch

@@ -0,0 +1,27 @@
+#!/bin/bash
+# Precompute the Code2LoRA timeline Space cache (cr_test, all 51 repos).
+#   sbatch --array=0-7 hf_space_demo/tools/precompute.sbatch
+#SBATCH --job-name=c2l_space_cache
+#SBATCH --output=slurm_logs/c2l_space_cache_%A_%a.out
+#SBATCH --error=slurm_logs/c2l_space_cache_%A_%a.err
+#SBATCH --time=03:00:00
+#SBATCH --gres=gpu:h100:1
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=80G
+#SBATCH --account=def-yuntian_gpu
+#SBATCH --export=ALL
+
+set -euo pipefail
+module purge 2>/dev/null || true
+module load StdEnv/2023 python/3.12 gcc/12.3 cuda/12.6 arrow
+source /scratch/lhotsko/venvs/qwen-cu126-py312/bin/activate
+cd /home/lhotsko/RepoPeftData
+mkdir -p slurm_logs
+
+export PYTHONUNBUFFERED=1
+export HF_HOME="${HF_HOME:-$HOME/.cache/huggingface}"
+export HF_HUB_OFFLINE=1
+
+python hf_space_demo/tools/precompute.py \
+  --out-dir /home/lhotsko/projects/def-yuntian/lhotsko/c2l_space/cache \
+  --shard-i "${SLURM_ARRAY_TASK_ID:-0}" --num-shards 8

tools/smoke_test.py

@@ -0,0 +1,140 @@
+#!/usr/bin/env python3
+"""CPU smoke test for the Space.
+
+Builds a one-repo data directory (real GRU contexts and prompts, placeholder
+cached predictions), then drives every app handler and serves the UI once.
+
+    C2L_DATA_DIR=/tmp/c2l_smoke python tools/smoke_test.py --repo MiniMax-AI/Mini-Agent
+"""
+
+from __future__ import annotations
+
+import argparse
+import gzip
+import json
+import os
+import sys
+import time
+import urllib.request
+from collections import defaultdict
+from pathlib import Path
+
+import numpy as np
+import pyarrow.parquet as pq
+import torch
+
+HERE = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(HERE))
+
+from c2l.core import CommitGRU  # noqa: E402
+
+COMMITS = "/scratch/lhotsko/REPO_DATASET/commit_parquet_hf_v2/commits/cr_test.parquet"
+QNAS = "/scratch/lhotsko/REPO_DATASET/code2lora_snapshots_hf/qna/cr_test.parquet"
+
+
+def build(repo: str, out: Path) -> None:
+    from huggingface_hub import hf_hub_download
+    from transformers import AutoTokenizer
+
+    out.mkdir(parents=True, exist_ok=True)
+    safe = repo.replace("/", "__")
+    ct = pq.read_table(COMMITS, filters=[("repo_id", "=", repo)]).to_pylist()
+    ct.sort(key=lambda r: r["commit_index"])
+    qt = pq.read_table(QNAS, columns=["repo_id", "commit_sha", "test_file", "lineno", "test_function",
+                                      "prefix", "target"], filters=[("repo_id", "=", repo)]).to_pylist()
+    by_sha = defaultdict(list)
+    for q in qt:
+        by_sha[q["commit_sha"]].append(q)
+
+    state = torch.load(hf_hub_download("code2lora/code2lora-gru", "code2lora_gru.pt"),
+                       map_location="cpu", weights_only=False, mmap=True)
+    gc = state.get("gru_config") or {}
+    gru = CommitGRU(int(gc.get("diff_input_dim", 2048)), int(gc.get("repo_state_dim", 2048)),
+                    int(gc.get("hidden_dim", 2048)))
+    gru.load_state_dict(state["gru_state"])
+    gru.eval()
+    with torch.no_grad():
+        h = gru.init_hidden(torch.tensor(ct[0]["repo_state_embedding"]).unsqueeze(0))
+        ctx = []
+        for r in ct:
+            h = gru.step(torch.tensor(r["diff_embedding"]).unsqueeze(0), h)
+            ctx.append(gru.context(h).float())
+    np.save(out / f"{safe}.ctx.npy", torch.cat(ctx).numpy().astype(np.float32))
+
+    tok = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-1.5B")
+    bos = tok.bos_token_id if tok.bos_token_id is not None else tok.eos_token_id
+    commits, qnas = [], []
+    for k, r in enumerate(ct):
+        group = []
+        for j, q in enumerate(by_sha[r["commit_sha"]][:8]):
+            ids = ([bos] + tok.encode(q["prefix"], add_special_tokens=False))[-4096:]
+            had_bos = ids[0] == bos
+            group.append({"k": k, "test_file": q["test_file"], "lineno": int(q["lineno"] or 0),
+                          "test_function": q["test_function"] or "",
+                          "prompt": tok.decode(ids[1:] if had_bos else ids), "had_bos": had_bos,
+                          "n_prompt_tokens": len(ids), "target": q["target"],
+                          "base_pred": "placeholder", "base_em": 0.0, "base_es": 0.0,
+                          "c2l_pred": q["target"], "c2l_em": 1.0, "c2l_es": 1.0,
+                          **({"prev_pred": "placeholder", "prev_em": 0.0, "prev_es": 0.0} if j == 0 and k else {})})
+        qnas.append(group)
+        commits.append({"sha": r["commit_sha"], "commit_index": r["commit_index"],
+                        "timestamp": r["commit_timestamp"], "n_new_assertions": r["n_new_assertions"],
+                        "diff_preview": (r["production_code_diff"] or "")[:6000],
+                        "n_qnas": len(group), "base_em": 0.0 if group else None,
+                        "c2l_em": 1.0 if group else None, "time_travel": bool(group and k)})
+    with gzip.open(out / f"{safe}.json.gz", "wt", encoding="utf-8") as f:
+        json.dump({"repo_id": repo, "commits": commits, "qnas": qnas,
+                   "max_input_tokens": 4096, "max_new_tokens": 64}, f)
+    (out / "index.json").write_text(json.dumps([{
+        "repo_id": repo, "safe": safe, "n_commits": len(commits), "n_qnas": sum(map(len, qnas)),
+        "base_em": 0.0, "c2l_em": 1.0, "n_time_travel": sum(c["time_travel"] for c in commits),
+        "first_date": commits[0]["timestamp"], "last_date": commits[-1]["timestamp"],
+        "description": "smoke test", "stars": 1234}]))
+    print(f"[build] {repo}: {len(commits)} commits, {sum(map(len, qnas))} qnas", flush=True)
+
+
+def drive() -> None:
+    t0 = time.time()
+    import app
+    print(f"[app] import + engine {time.time() - t0:.0f}s", flush=True)
+    safe = app.INDEX[0]["safe"]
+    r = app.on_repo(safe)
+    k1 = r[1].value if hasattr(r[1], "value") else r[1]["value"]
+    print("[on_repo] ok, start commit", k1, flush=True)
+    r = app.on_commit(safe, k1, None)
+    i = r[3].value if hasattr(r[3], "value") else r[3]["value"]
+    print("[on_commit] ok, pick", i, flush=True)
+    t = time.time()
+    status, st, *_ = app.on_adapt(safe, k1, i)
+    print(f"[on_adapt] {time.time() - t:.2f}s  state={st}", flush=True)
+    print(status[:400])
+    _, preds, _ = app.on_pick(safe, k1, i, st)
+    assert "Code2LoRA @ commit" in preds, preds
+    t = time.time()
+    live = app.on_live(safe, k1, i, st, progress=lambda *a, **k: None)
+    print(f"[on_live] {time.time() - t:.1f}s\n{live}", flush=True)
+    t = time.time()
+    prompt = app.on_prefill(safe, k1, i)
+    custom = app.on_custom(safe, k1, prompt, 32, st, progress=lambda *a, **k: None)
+    print(f"[on_custom] {time.time() - t:.1f}s\n{custom[:1500]}", flush=True)
+    for k in range(1, len(app.load_repo(safe)["commits"]) + 1):
+        app.on_commit(safe, k, st)
+        app.on_adapt(safe, k, 0)
+    print("[all commits] ok", flush=True)
+
+    app.demo.queue(max_size=32).launch(server_port=7861, prevent_thread_lock=True)
+    body = urllib.request.urlopen("http://127.0.0.1:7861/", timeout=30).read()
+    print(f"[serve] GET / -> {len(body)} bytes", flush=True)
+    app.demo.close()
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--repo", default="MiniMax-AI/Mini-Agent")
+    ap.add_argument("--build-only", action="store_true")
+    args = ap.parse_args()
+    data = Path(os.environ["C2L_DATA_DIR"])
+    if not (data / "index.json").exists():
+        build(args.repo, data)
+    if not args.build_only:
+        drive()

tools/smoke_test.sbatch

@@ -0,0 +1,20 @@
+#!/bin/bash
+# CPU smoke test mirroring the free Space tier (2 vCPU).
+#   sbatch hf_space_demo/tools/smoke_test.sbatch
+#SBATCH --job-name=c2l_space_smoke
+#SBATCH --output=slurm_logs/c2l_space_smoke_%j.out
+#SBATCH --time=01:00:00
+#SBATCH --cpus-per-task=2
+#SBATCH --mem=24G
+#SBATCH --account=def-yuntian
+
+set -euo pipefail
+module purge 2>/dev/null || true
+module load StdEnv/2023 python/3.12 arrow
+P=/home/lhotsko/projects/def-yuntian/lhotsko/c2l_space
+source $P/venv/bin/activate
+cd /home/lhotsko/RepoPeftData/hf_space_demo
+
+export PYTHONUNBUFFERED=1 HF_HOME=$P/hf HF_HUB_OFFLINE=1 C2L_DATA_DIR=$P/smoke_data
+export OMP_NUM_THREADS=2 GRADIO_ANALYTICS_ENABLED=False
+python tools/smoke_test.py

upload_checkpoint.py

@@ -1,47 +0,0 @@
-#!/usr/bin/env python3
-"""Upload a trained Code2LoRA-GRU checkpoint to a Hugging Face model repo.
-
-Usage:
-    huggingface-cli login          # once, to set your token
-    python upload_checkpoint.py \
-        --ckpt /path/to/gru_head.best.pt \
-        --repo your-name/code2lora-gru
-"""
-
-from __future__ import annotations
-
-import argparse
-from pathlib import Path
-
-
-def main() -> None:
-    ap = argparse.ArgumentParser(description=__doc__)
-    ap.add_argument("--ckpt", required=True, help="Path to gru_head.*.pt")
-    ap.add_argument("--repo", required=True,
-                    help="Target HF model repo id, e.g. your-name/code2lora-gru")
-    ap.add_argument("--filename", default="gru_head.best.pt",
-                    help="Name to store it under in the repo.")
-    ap.add_argument("--private", action="store_true")
-    args = ap.parse_args()
-
-    ckpt = Path(args.ckpt)
-    if not ckpt.exists():
-        raise SystemExit(f"Checkpoint not found: {ckpt}")
-
-    from huggingface_hub import HfApi
-    api = HfApi()
-    api.create_repo(args.repo, repo_type="model", private=args.private,
-                    exist_ok=True)
-    print(f"Uploading {ckpt} ({ckpt.stat().st_size / 1e9:.2f} GB) -> "
-          f"{args.repo}/{args.filename} ...")
-    api.upload_file(
-        path_or_fileobj=str(ckpt),
-        path_in_repo=args.filename,
-        repo_id=args.repo,
-        repo_type="model",
-    )
-    print("Done. Set CODE2LORA_CKPT_REPO=%s in the Space settings." % args.repo)
-
-
-if __name__ == "__main__":
-    main()