tools/precompute.py

#!/usr/bin/env python3
"""Precompute everything the Code2LoRA timeline Space serves from cache.

For every cross-repo test repository (held out from training) this walks the
Code2LoRA-Evo GRU over the precomputed per-commit embeddings and records:

* ``ctx``  -- output_norm(h_t) after every commit (what the head consumes),
* for every QnA attached to a commit (first <=8, the paper's eval cap):
  - the base-model prediction (no adapter),
  - the prediction with the adapter generated at that commit (paper protocol),
  - for QnAs the adapter solves and the base model does not, the prediction
    with the adapter from the *previous* commit ("time travel").

One pickle per repository is written to ``--out-dir``; existing files are
skipped, so the job is resumable. Run from the RepoPeftData root:

    python hf_space_demo/tools/precompute.py --out-dir $SCRATCH/C2L_SPACE_CACHE \
        --shard-i 0 --num-shards 8
"""

from __future__ import annotations

import argparse
import pickle
import sys
import time
from collections import defaultdict
from pathlib import Path
from typing import Dict, List, Optional

import numpy as np
import pyarrow.parquet as pq
import torch

_ROOT = Path(__file__).resolve().parents[2]
for p in (_ROOT, _ROOT / "hypernetwork"):
    if str(p) not in sys.path:
        sys.path.insert(0, str(p))

from code2lora_core import (  # noqa: E402
    LoRA,
    inject_lora_weights,
    load_commit_rows_for_gru,
    load_qna_rows,
)
from evaluation.metrics import compute_metrics  # noqa: E402
from evaluation.run_code2lora_gru_v2_eval import (  # noqa: E402
    DEFAULT_BASE_MODEL,
    DEFAULT_TARGET_MODULES,
    _generate_batch,
    _get_bos_id,
    _load_v2_gru_ckpt,
    _prepare_prefix_ids,
)

COMMITS = "/scratch/lhotsko/REPO_DATASET/commit_parquet_hf_v2/commits/cr_test.parquet"
QNAS = "/scratch/lhotsko/REPO_DATASET/code2lora_snapshots_hf/qna/cr_test.parquet"
CKPT = "/scratch/lhotsko/hf_staging/code2lora_gru/gru_head.best.pt"


def clear_all_lora_weights(model: torch.nn.Module, specs) -> None:
    named = dict(model.named_modules())
    for sp in specs:
        mod = named[sp.full_name]
        if isinstance(mod, LoRA):
            mod.A, mod.B = None, None


def safe_name(repo_id: str) -> str:
    return repo_id.replace("/", "__")


def commit_metadata(repo_ids: List[str]) -> Dict[str, Dict[str, dict]]:
    t = pq.read_table(COMMITS, columns=[
        "repo_id", "commit_sha", "commit_index", "commit_timestamp",
        "n_new_assertions", "production_code_diff"],
        filters=[("repo_id", "in", repo_ids)])
    out: Dict[str, Dict[str, dict]] = defaultdict(dict)
    for r in t.to_pylist():
        diff = r["production_code_diff"] or ""
        out[r["repo_id"]][r["commit_sha"]] = {
            "timestamp": r["commit_timestamp"],
            "n_new_assertions": int(r["n_new_assertions"] or 0),
            "diff_preview": diff[:6000],
            "diff_chars": len(diff),
        }
    return out


def score(pred: str, target: str) -> Dict[str, float]:
    m = compute_metrics(pred, target)
    return {"em": float(bool(m["exact_match"])), "es": float(m["edit_similarity"])}


@torch.no_grad()
def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out-dir", required=True)
    ap.add_argument("--checkpoint", default=CKPT)
    ap.add_argument("--shard-i", type=int, default=0)
    ap.add_argument("--num-shards", type=int, default=1)
    ap.add_argument("--max-input-tokens", type=int, default=4096)
    ap.add_argument("--max-new-tokens", type=int, default=64)
    ap.add_argument("--batch-size", type=int, default=8)
    ap.add_argument("--qnas-per-commit-limit", type=int, default=8)
    ap.add_argument("--repos", nargs="*", default=None)
    ap.add_argument("--min-commits", type=int, default=10)
    ap.add_argument("--max-commits", type=int, default=400)
    args = ap.parse_args()

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    rows_by_repo = load_commit_rows_for_gru(Path(COMMITS))
    repos = sorted(r for r, rows in rows_by_repo.items()
                   if args.min_commits <= len(rows) <= args.max_commits)
    if args.repos:
        repos = [r for r in repos if r in set(args.repos)]
    repos = [r for i, r in enumerate(repos) if i % args.num_shards == args.shard_i]
    repos = [r for r in repos if not (out_dir / f"{safe_name(r)}.pkl").exists()]
    print(f"[shard {args.shard_i}/{args.num_shards}] {len(repos)} repos to do: {repos}", flush=True)
    if not repos:
        return

    base_model, gru, head, specs, tok = _load_v2_gru_ckpt(
        Path(args.checkpoint), DEFAULT_BASE_MODEL, DEFAULT_TARGET_MODULES, device)
    bos = _get_bos_id(tok)
    meta = commit_metadata(repos)

    qnas_by_key: Dict[tuple, List[dict]] = defaultdict(list)
    for q in load_qna_rows(Path(QNAS), repo_ids=repos):
        qnas_by_key[(q.repo_id, q.commit_sha)].append(q)

    def gen(ids_list: List[List[int]]) -> List[str]:
        out: List[str] = []
        for i in range(0, len(ids_list), args.batch_size):
            out += _generate_batch(base_model, tok, device, ids_list[i:i + args.batch_size],
                                   max_new_tokens=args.max_new_tokens)
        return out

    for repo in repos:
        t0 = time.time()
        rows = rows_by_repo[repo]
        h = gru.init_hidden(torch.from_numpy(rows[0].repo_state_embedding).to(device).unsqueeze(0))
        ctxs = []
        for row in rows:
            h = gru.step(torch.from_numpy(row.diff_embedding).to(device).unsqueeze(0), h)
            ctxs.append(gru.output_norm(h[-1]).float())
        ctx = torch.cat(ctxs)                                   # [T, 2048]

        commits, qnas = [], []
        for k, row in enumerate(rows):
            m = meta[repo].get(row.commit_sha, {})
            commits.append({"sha": row.commit_sha, "commit_index": int(row.commit_index),
                            "in_repo_split": row.in_repo_split, **m})
            for q in qnas_by_key.get((repo, row.commit_sha), [])[: args.qnas_per_commit_limit]:
                ids = _prepare_prefix_ids(tok, q.prefix, args.max_input_tokens, bos)
                had_bos = bos is not None and ids[0] == bos
                qnas.append({
                    "k": k, "test_file": q.test_file, "lineno": int(q.lineno or 0),
                    "test_function": q.test_function or "",
                    "prompt": tok.decode(ids[1:] if had_bos else ids),
                    "had_bos": had_bos, "n_prompt_tokens": len(ids),
                    "target": q.target, "_ids": ids,
                })

        clear_all_lora_weights(base_model, specs)
        for q, p in zip(qnas, gen([q["_ids"] for q in qnas])):
            q["base_pred"] = p
            q.update({f"base_{k}": v for k, v in score(p, q["target"]).items()})

        by_k: Dict[int, List[dict]] = defaultdict(list)
        for q in qnas:
            by_k[q["k"]].append(q)
        for k, group in by_k.items():
            inject_lora_weights(base_model, specs, head(ctx[k:k + 1]), batch_index=0)
            for q, p in zip(group, gen([q["_ids"] for q in group])):
                q["c2l_pred"] = p
                q.update({f"c2l_{kk}": v for kk, v in score(p, q["target"]).items()})

        wins = [q for q in qnas if q["c2l_em"] == 1.0 and q["base_em"] == 0.0 and q["k"] > 0]
        by_prev: Dict[int, List[dict]] = defaultdict(list)
        for q in wins:
            by_prev[q["k"] - 1].append(q)
        for k_prev, group in by_prev.items():
            inject_lora_weights(base_model, specs, head(ctx[k_prev:k_prev + 1]), batch_index=0)
            for q, p in zip(group, gen([q["_ids"] for q in group])):
                q["prev_pred"] = p
                q.update({f"prev_{kk}": v for kk, v in score(p, q["target"]).items()})
        clear_all_lora_weights(base_model, specs)

        for q in qnas:
            q.pop("_ids")
        payload = {"repo_id": repo, "ctx": ctx.cpu().numpy().astype(np.float32),
                   "commits": commits, "qnas": qnas,
                   "max_input_tokens": args.max_input_tokens,
                   "max_new_tokens": args.max_new_tokens}
        tmp = out_dir / f"{safe_name(repo)}.pkl.tmp"
        with open(tmp, "wb") as f:
            pickle.dump(payload, f)
        tmp.replace(out_dir / f"{safe_name(repo)}.pkl")
        n = max(len(qnas), 1)
        print(f"[{repo}] {len(rows)} commits, {len(qnas)} qnas, "
              f"base EM={sum(q['base_em'] for q in qnas)/n:.3f} "
              f"c2l EM={sum(q['c2l_em'] for q in qnas)/n:.3f} "
              f"time-travel={sum(1 for q in wins if q.get('prev_em') == 0.0)} "
              f"({(time.time()-t0)/60:.1f} min)", flush=True)


if __name__ == "__main__":
    main()