Code2LoRA-GRU interactive demo

README.md

@@ -1,13 +1,130 @@
 ---
-title: Code2lora Gru Demo
-emoji: 🏆
-colorFrom: purple
-colorTo: indigo
+title: Code2LoRA-GRU
+emoji: 🧬
+colorFrom: blue
+colorTo: green
 sdk: gradio
-sdk_version: 6.17.3
-python_version: '3.13'
+sdk_version: 5.9.1
 app_file: app.py
 pinned: false
+license: mit
+short_description: Turn a repo's commit history into a LoRA adapter, live.
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Code2LoRA-GRU — interactive demo
+
+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.
+
+This Space implements the inference pipeline from **Code2LoRA: Hypernetworks
+for Repository-Conditioned and Commit-Streaming Adapters of Code Language
+Models**.
+
+---
+
+## 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
+```

app.py

@@ -0,0 +1,353 @@
+"""Code2LoRA-GRU — interactive demo (Hugging Face Space).
+
+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.
+"""
+
+from __future__ import annotations
+
+import html
+import tempfile
+import traceback
+from pathlib import Path
+from typing import List, Optional
+
+import gradio as gr
+
+from c2l_demo.engine import Engine, RepoSession
+
+# ---------------------------------------------------------------------------
+# Process-wide engine (lazy model loading happens on first use)
+# ---------------------------------------------------------------------------
+
+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.
+"""
+
+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.
+"""
+
+
+# ---------------------------------------------------------------------------
+# 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"),
+    ])
+
+    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
+
+
+# ---------------------------------------------------------------------------
+# Callbacks
+# ---------------------------------------------------------------------------
+
+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>")
+
+
+# ---------------------------------------------------------------------------
+# UI
+# ---------------------------------------------------------------------------
+
+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;}
+"""
+
+
+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)
+
+        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>""")
+
+        with gr.Accordion("How it works", open=False):
+            gr.Markdown(HOW_IT_WORKS_MD)
+
+        # 0) Instructions / requirements
+        with gr.Accordion("0 · Repository requirements (read me first)", open=True):
+            gr.Markdown(REQUIREMENTS_MD)
+
+        # 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.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
+
+
+if __name__ == "__main__":
+    build_demo().queue().launch(server_name="0.0.0.0", server_port=7860)

c2l_demo/__init__.py

@@ -0,0 +1,17 @@
+"""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

@@ -0,0 +1,360 @@
+"""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/core.py

@@ -0,0 +1,295 @@
+"""Model core for Code2LoRA-GRU inference.
+
+Faithful port of ``hypernetwork/code2lora_core.py`` from the RepoPeftData
+project, trimmed to the pieces needed for inference (the parquet loaders that
+the trainer uses are intentionally omitted so this Space has no pyarrow
+dependency).
+
+The two networks that get loaded from a ``gru_head.*.pt`` checkpoint are:
+
+* :class:`CommitGRU`     -- a 1-layer GRU that ingests per-commit diff
+  embeddings; ``h_0`` comes from a learnable projection of the initial
+  repo-state embedding.
+* :class:`Code2LoRAHead`  -- maps the GRU hidden state to one LoRA (A, B)
+  pair per module *type* (q_proj, v_proj, ...), shared across all
+  transformer layers.
+
+The generated weights are pushed into :class:`LoRA` wrappers that replace the
+target ``nn.Linear`` modules of the frozen base LLM.
+"""
+
+from __future__ import annotations
+
+import math
+import re
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ---------------------------------------------------------------------------
+# LoRA module + injection
+# ---------------------------------------------------------------------------
+
+class LoRA(nn.Module):
+    """Wraps an ``nn.Linear`` with an additive low-rank update.
+
+    Forward: ``y = base(x) + scaling * (x @ A^T) @ B^T`` where the per-call
+    A: ``[rank, in_features]`` and B: ``[out_features, rank]`` come from the
+    hypernet via :meth:`set_lora_weights`.
+    """
+
+    def __init__(self, base: nn.Linear, in_features: int, out_features: int,
+                 rank: int, alpha: float):
+        super().__init__()
+        self.base = base
+        for p in self.base.parameters():
+            p.requires_grad = False
+        self.in_features = in_features
+        self.out_features = out_features
+        self.rank = rank
+        self.scaling = float(alpha) / float(max(1, rank))
+        self.A: Optional[torch.Tensor] = None  # [rank, in_features]
+        self.B: Optional[torch.Tensor] = None  # [out_features, rank]
+
+    def set_lora_weights(self, A: torch.Tensor, B: torch.Tensor) -> None:
+        self.A = A
+        self.B = B
+
+    def clear_lora_weights(self) -> None:
+        self.A = None
+        self.B = None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = self.base(x)
+        if self.A is None or self.B is None:
+            return y
+        x_f32 = x.detach().to(torch.float32)
+        A = self.A.to(torch.float32)
+        B = self.B.to(torch.float32)
+        delta = F.linear(F.linear(x_f32, A), B) * self.scaling
+        return y + delta.to(dtype=y.dtype)
+
+
+@dataclass
+class ModuleSpec:
+    full_name: str   # e.g. 'model.layers.5.self_attn.q_proj'
+    layer_idx: int
+    type: str        # e.g. 'q_proj'
+    in_features: int
+    out_features: int
+
+
+def get_module_specs(model: nn.Module, target_module_types: List[str]
+                     ) -> List[ModuleSpec]:
+    """Discover every nn.Linear whose name contains a target type."""
+    specs: List[ModuleSpec] = []
+    pat = re.compile(r"\bmodel\.layers\.(\d+)\.")
+    for name, m in model.named_modules():
+        match_type = next((t for t in target_module_types if t in name), None)
+        if match_type is None:
+            continue
+        if not isinstance(m, nn.Linear):
+            continue
+        m_layer = pat.search(name)
+        layer_idx = int(m_layer.group(1)) if m_layer else -1
+        specs.append(ModuleSpec(
+            full_name=name,
+            layer_idx=layer_idx,
+            type=match_type,
+            in_features=int(m.in_features),
+            out_features=int(m.out_features),
+        ))
+    specs.sort(key=lambda s: (s.layer_idx, s.full_name))
+    return specs
+
+
+def replace_with_lora(model: nn.Module, specs: List[ModuleSpec],
+                      rank: int, alpha: float) -> None:
+    """Replace each target ``nn.Linear`` with a :class:`LoRA` wrapper. Idempotent."""
+    named = dict(model.named_modules())
+    device = next(model.parameters()).device
+    dtype = next(model.parameters()).dtype
+    for sp in specs:
+        parent_name, attr = sp.full_name.rsplit(".", 1)
+        orig = getattr(named[parent_name], attr)
+        if isinstance(orig, LoRA):
+            continue
+        assert isinstance(orig, nn.Linear), \
+            f"{sp.full_name} is not nn.Linear (got {type(orig)})"
+        wrapped = LoRA(orig, sp.in_features, sp.out_features,
+                       rank, alpha).to(device=device, dtype=dtype)
+        setattr(named[parent_name], attr, wrapped)
+
+
+def inject_lora_weights(model: nn.Module, specs: List[ModuleSpec],
+                        head_out: Dict[str, Dict[str, torch.Tensor]],
+                        batch_index: int = 0) -> None:
+    """Push ``head_out['A'][type]`` / ``head_out['B'][type]`` into each LoRA."""
+    A_by_type = head_out["A"]
+    B_by_type = head_out["B"]
+    named = dict(model.named_modules())
+    for sp in specs:
+        named[sp.full_name].set_lora_weights(
+            A_by_type[sp.type][batch_index],
+            B_by_type[sp.type][batch_index],
+        )
+
+
+def clear_all_lora_weights(model: nn.Module, specs: List[ModuleSpec]) -> None:
+    """Reset every wrapper LoRA to identity (base model only)."""
+    named = dict(model.named_modules())
+    for sp in specs:
+        mod = named.get(sp.full_name)
+        if isinstance(mod, LoRA):
+            mod.clear_lora_weights()
+
+
+def discover_module_types_and_dims(specs: List[ModuleSpec]
+                                   ) -> Dict[str, Tuple[int, int]]:
+    type_dims: Dict[str, Tuple[int, int]] = {}
+    for sp in specs:
+        if sp.type in type_dims:
+            assert type_dims[sp.type] == (sp.in_features, sp.out_features), \
+                f"type {sp.type} appears with inconsistent dims"
+            continue
+        type_dims[sp.type] = (sp.in_features, sp.out_features)
+    return type_dims
+
+
+# ---------------------------------------------------------------------------
+# Shared LoRA generation head
+# ---------------------------------------------------------------------------
+
+class Code2LoRAHead(nn.Module):
+    """Maps a context vector to one LoRA (A, B) pair per module type."""
+
+    def __init__(
+        self,
+        input_dim: int,
+        type_dims: Dict[str, Tuple[int, int]],
+        hidden_dim: int = 1024,
+        rank: int = 16,
+        init_log_scale: float = -3.5,
+    ):
+        super().__init__()
+        self.input_dim = input_dim
+        self.hidden_dim = hidden_dim
+        self.rank = rank
+        self.type_dims = dict(type_dims)
+        self.types = sorted(type_dims.keys())
+
+        self.trunk = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.GELU(),
+        )
+        self.heads_A = nn.ModuleDict({
+            t: nn.Linear(hidden_dim, rank * type_dims[t][0]) for t in self.types
+        })
+        self.heads_B = nn.ModuleDict({
+            t: nn.Linear(hidden_dim, type_dims[t][1] * rank) for t in self.types
+        })
+        self.log_scale_A = nn.ParameterDict({
+            t: nn.Parameter(torch.tensor(init_log_scale)) for t in self.types
+        })
+        self.log_scale_B = nn.ParameterDict({
+            t: nn.Parameter(torch.tensor(init_log_scale)) for t in self.types
+        })
+
+    def forward(self, ctx: torch.Tensor) -> Dict[str, Dict[str, torch.Tensor]]:
+        if ctx.dim() == 3:
+            ctx = torch.max(ctx, dim=1).values
+        h = self.trunk(ctx.float())
+        h = F.normalize(h, p=2, dim=-1) * math.sqrt(self.hidden_dim)
+
+        A_out: Dict[str, torch.Tensor] = {}
+        B_out: Dict[str, torch.Tensor] = {}
+        for t in self.types:
+            in_f, out_f = self.type_dims[t]
+            A_raw = self.heads_A[t](h).view(-1, self.rank, in_f)
+            B_raw = self.heads_B[t](h).view(-1, out_f, self.rank)
+            scale_A = torch.exp(self.log_scale_A[t]).clamp(1e-5, 0.3)
+            scale_B = torch.exp(self.log_scale_B[t]).clamp(1e-5, 0.3)
+            A_out[t] = torch.tanh(A_raw) * scale_A
+            B_out[t] = torch.tanh(B_raw) * scale_B
+        return {"A": A_out, "B": B_out}
+
+    def config_dict(self) -> Dict[str, Any]:
+        return {
+            "input_dim": self.input_dim,
+            "hidden_dim": self.hidden_dim,
+            "rank": self.rank,
+            "types": self.types,
+            "type_dims": {t: list(v) for t, v in self.type_dims.items()},
+        }
+
+
+# ---------------------------------------------------------------------------
+# Commit-streaming GRU
+# ---------------------------------------------------------------------------
+
+class CommitGRU(nn.Module):
+    """One-layer GRU that ingests per-commit diff embeddings.
+
+    ``h_0`` is a learnable projection of the initial repo-state embedding;
+    ``h_T`` is fed into :class:`Code2LoRAHead`.
+    """
+
+    def __init__(
+        self,
+        diff_input_dim: int = 2048,
+        repo_state_dim: int = 2048,
+        hidden_dim: int = 2048,
+    ):
+        super().__init__()
+        self.diff_input_dim = diff_input_dim
+        self.repo_state_dim = repo_state_dim
+        self.hidden_dim = hidden_dim
+
+        self.diff_proj = nn.Sequential(
+            nn.Linear(diff_input_dim, hidden_dim),
+            nn.LayerNorm(hidden_dim),
+        )
+        self.repo_init_proj = nn.Sequential(
+            nn.Linear(repo_state_dim, hidden_dim),
+            nn.GELU(),
+            nn.LayerNorm(hidden_dim),
+        )
+        self.gru = nn.GRU(
+            input_size=hidden_dim,
+            hidden_size=hidden_dim,
+            num_layers=1,
+            batch_first=True,
+        )
+        self.output_norm = nn.LayerNorm(hidden_dim)
+
+    def init_hidden(self, repo_state_emb_0: torch.Tensor) -> torch.Tensor:
+        h0 = self.repo_init_proj(repo_state_emb_0.float())
+        return h0.unsqueeze(0)
+
+    def step(self, diff_emb: torch.Tensor, h_prev: torch.Tensor) -> torch.Tensor:
+        x = self.diff_proj(diff_emb.float()).unsqueeze(1)  # [B, 1, H]
+        _, h_new = self.gru(x, h_prev)
+        return h_new
+
+    def context(self, h: torch.Tensor) -> torch.Tensor:
+        """Match the eval path: ctx = output_norm(h[-1])."""
+        return self.output_norm(h[-1])
+
+
+__all__ = [
+    "LoRA",
+    "ModuleSpec",
+    "get_module_specs",
+    "replace_with_lora",
+    "inject_lora_weights",
+    "clear_all_lora_weights",
+    "discover_module_types_and_dims",
+    "Code2LoRAHead",
+    "CommitGRU",
+]

c2l_demo/embedding.py

@@ -0,0 +1,165 @@
+"""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

@@ -0,0 +1,400 @@
+"""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

@@ -0,0 +1,498 @@
+"""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/metrics.py

@@ -0,0 +1,116 @@
+"""Evaluation metrics for the demo.
+
+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:
+        s = s.replace(tok, "")
+    return s.strip()
+
+
+def strip_comments(s: str) -> str:
+    return s.split("#")[0].strip()
+
+
+def normalize_for_match(s: str) -> str:
+    s = s.strip().rstrip(":, \t")
+    s = s.replace(", ", ",").replace(" ,", ",")
+    return " ".join(s.split())
+
+
+def exact_match(pred: str, ref: str) -> bool:
+    return normalize_for_match(pred) == normalize_for_match(ref)
+
+
+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:
+        pred = pred.split("\n")[0]
+    pred = strip_comments(pred)
+    ref = strip_comments(target).strip()
+    if ref:
+        if "," not in ref and "," in pred:
+            pred = pred.split(",")[0].strip()
+        if len(ref.split()) == 1 and len(pred.strip().split()) > 1:
+            pred = pred.strip().split()[0]
+    return pred.strip()
+
+
+def compute_metrics(pred: str, ref: str) -> dict:
+    pred_clean = postprocess_prediction(pred, ref)
+    ref_clean = strip_comments(ref)
+    return {
+        "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/viz.py

@@ -0,0 +1,192 @@
+"""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

@@ -0,0 +1,9 @@
+gradio>=4.44,<6
+torch>=2.3
+transformers>=4.45
+accelerate>=0.34
+huggingface_hub>=0.25
+numpy>=1.24
+matplotlib>=3.7
+# optional, improves CodeBLEU fidelity if present:
+# codebleu>=0.7

upload_checkpoint.py

@@ -0,0 +1,47 @@
+#!/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()