#!/usr/bin/env python3 """Merge RepoPeftBench (Code2LoRA's own benchmark) into the aligned6 corpus. RepoPeftBench contributes ~500K *assertion-completion* items over 512 repos -- short, exact code targets. aligned6 contributes ~27K *prose* QA over 2066 repos -- conventions and architecture. They are complementary: the benchmark run showed the head learns a repo's stack and conventions but not what it actually does, and exact-recall data is what addresses that. Split integrity is the thing to get right. A repo must never appear in both a training split and an eval split, or cross-repo evaluation becomes meaningless. RepoPeftBench already partitions BY REPO (cr_val / cr_test hold out whole repositories), so we carry its partition through unchanged and only ever add repos to `train` when the benchmark itself calls them training repos. doc split : which repos the hypernetwork trains on (train/cr_val/cr_test) qna_split : within a train repo, held-out QA for in-repo eval (ir_*) Usage: python scripts/merge_repopeft_corpus.py \ --repopeft-emb data/embeddings/repopeft_6view.parquet \ --out-emb data/embeddings/all_lora_embeddings.parquet \ --out-qna data/qna/all_lora_qna.jsonl """ from __future__ import annotations import argparse import json import sys from collections import Counter, defaultdict from pathlib import Path import pyarrow.parquet as pq HERE = Path(__file__).resolve().parent REPO_ROOT = HERE.parent sys.path.insert(0, str(REPO_ROOT)) # file stem -> (doc split contributed, qna_split) # cr_* hold out whole repos; ir_* are held-out QA of repos that stay in train. SPLIT_MAP = { "train": ("train", "train"), "cr_val": ("cr_val", "train"), "cr_test": ("cr_test", "train"), "ir_val": ("train", "held_out"), "ir_test": ("train", "held_out"), } def main() -> None: ap = argparse.ArgumentParser() ap.add_argument("--repopeft-emb", default="data/embeddings/repopeft_6view.parquet") ap.add_argument("--aligned-emb", default="data/embeddings/aligned6_embeddings.parquet") ap.add_argument("--aligned-qna", default="data/qna/aligned6_qna.jsonl") ap.add_argument("--repopeft-glob", default="data/real_code2lora/*/qna/*.parquet") ap.add_argument("--out-emb", default="data/embeddings/all_lora_embeddings.parquet") ap.add_argument("--out-qna", default="data/qna/all_lora_qna.jsonl") ap.add_argument("--max-qna-per-repo", type=int, default=400, help="cap per repo: evo alone has ~1000/repo, which would " "swamp the prose QA and bias the head toward one task") ap.add_argument("--max-target-chars", type=int, default=400) args = ap.parse_args() import glob as _glob import pyarrow as pa # ---------------- embeddings ---------------- ali = pq.read_table(args.aligned_emb) ali_dim = len(ali.column("doc_embedding")[0].as_py()) print(f"aligned6: {ali.num_rows} repos, dim {ali_dim}") rp_path = Path(args.repopeft_emb) if not rp_path.exists(): print(f"!! missing {rp_path} -- run build_repo_multiview.py first", file=sys.stderr) sys.exit(1) rp = pq.read_table(rp_path) rp_dim = len(rp.column("doc_embedding")[0].as_py()) print(f"repopeft: {rp.num_rows} repos, dim {rp_dim}") if rp_dim != ali_dim: print(f"!! dim mismatch {rp_dim} != {ali_dim}; the head cannot consume both", file=sys.stderr) sys.exit(1) # Which split does each RepoPeftBench repo belong to? Derived from the QA # files it appears in, so we inherit the benchmark's own repo partition. repo_split: dict[str, str] = {} for f in sorted(_glob.glob(args.repopeft_glob)): stem = Path(f).stem if stem not in SPLIT_MAP: continue doc_split, _ = SPLIT_MAP[stem] ids = set(pq.read_table(f, columns=["repo_id"]).column("repo_id").to_pylist()) for r in ids: # An eval assignment always wins: if a repo is used to hold out # cross-repo performance anywhere, it must never be trained on. if repo_split.get(r) in ("cr_val", "cr_test"): continue repo_split[r] = doc_split rp_ids = rp.column("doc_id").to_pylist() rp_embs = rp.column("doc_embedding").to_pylist() ali_ids = set(ali.column("doc_id").to_pylist()) out_ids, out_ver, out_split, out_cat, out_emb = [], [], [], [], [] for c, col in (("doc_id", out_ids), ("doc_version", out_ver), ("split", out_split), ("category", out_cat)): if c in ali.column_names: col.extend(ali.column(c).to_pylist()) else: col.extend(["v1"] * ali.num_rows if c == "doc_version" else ["aligned6"] * ali.num_rows) out_emb.extend(ali.column("doc_embedding").to_pylist()) added = 0 for rid, emb in zip(rp_ids, rp_embs): if rid in ali_ids: # already covered by aligned6 continue sp = repo_split.get(rid) if sp is None: # embedded but no QA -> useless continue out_ids.append(rid) out_ver.append("head") out_split.append(sp) out_cat.append("repopeftbench") out_emb.append(emb) added += 1 print(f"merged embeddings: {len(out_ids)} repos (+{added} from RepoPeftBench)") print(" split counts:", dict(Counter(out_split))) pq.write_table(pa.table({ "doc_id": out_ids, "doc_version": out_ver, "split": out_split, "category": out_cat, "doc_embedding": out_emb, }), args.out_emb) # ---------------- QA ---------------- have_emb = set(out_ids) split_of = dict(zip(out_ids, out_split)) per_repo: dict[str, int] = defaultdict(int) n_written = 0 src_counts: Counter = Counter() with open(args.out_qna, "w") as out: # aligned6 first, verbatim with open(args.aligned_qna) as f: for line in f: line = line.strip() if not line: continue out.write(line + "\n") n_written += 1 src_counts["aligned6"] += 1 for f in sorted(_glob.glob(args.repopeft_glob)): stem = Path(f).stem if stem not in SPLIT_MAP: continue _, qna_split = SPLIT_MAP[stem] t = pq.read_table(f, columns=["repo_id", "prefix", "target"]) rid_c = t.column("repo_id").to_pylist() pre_c = t.column("prefix").to_pylist() tgt_c = t.column("target").to_pylist() kept = 0 for rid, pre, tgt in zip(rid_c, pre_c, tgt_c): if rid not in have_emb: continue if not pre or not tgt: continue if len(tgt) > args.max_target_chars: continue if per_repo[rid] >= args.max_qna_per_repo: continue per_repo[rid] += 1 out.write(json.dumps({ "doc_id": rid, "doc_version": "head", "split": split_of[rid], "qna_split": qna_split, "question": "", "prefix": pre, "target": tgt, }) + "\n") kept += 1 n_written += 1 src_counts[Path(f).parent.parent.name + "/" + stem] += kept print(f"\nmerged QA: {n_written} rows -> {args.out_qna}") for k, v in src_counts.most_common(): print(f" {k:<34} {v}") print(f"\nrepos with QA: {len(per_repo)}") if __name__ == "__main__": main()