#!/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()