Spaces:
Sleeping
Sleeping
| #!/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"])} | |
| 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() | |