File size: 3,779 Bytes
481fbb6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | #!/usr/bin/env python3
"""Assemble the final ALIGNED 6-view training dataset by joining the
multi-view embeddings (inputs, doc_id = 'owner/repo') to ALL tech-lead QA
(targets) matched by repo name (commitpack/swe QA carry 'owner/repo@commit'
-> stripped to 'owner/repo'). This turns 132 repo-scoped-only aligned repos
into 515 aligned repos (~4000 QA) -- the best use of what we built.
Splits (per repo, deterministic): 80% train / 10% cr_val / 10% cr_test, so
the hypernetwork is evaluated on held-out repos it never trained on. Within
train repos, ~15% of QA -> qna_split=held_out (feeds ir_test).
Per-repo QA cap keeps any single repo (e.g. django) from dominating.
Output: data/embeddings/aligned6_embeddings.parquet
data/qna/aligned6_qna.jsonl
"""
from __future__ import annotations
import json, hashlib, random, sys
from collections import defaultdict
from pathlib import Path
import pyarrow as pa, pyarrow.parquet as pq
HERE = Path(__file__).resolve().parent; REPO_ROOT = HERE.parent
sys.path.insert(0, str(REPO_ROOT))
from memory_lora.data_paths import EMBEDDINGS_DIR, QNA_DIR
PER_REPO_CAP = 20
def split_of(repo: str) -> str:
h = int(hashlib.md5(repo.encode()).hexdigest(), 16) % 100
if h < 10: return "cr_test"
if h < 20: return "cr_val"
return "train"
def main():
rng = random.Random(3407)
mv = pq.read_table(EMBEDDINGS_DIR / "multiview_embeddings.parquet").to_pylist()
mv_by_repo = {}
for r in mv: # dedupe by repo (keep first)
mv_by_repo.setdefault(r["doc_id"], r)
mv_repos = set(mv_by_repo)
qa_by_repo = defaultdict(list)
for f in ["repo_scoped_qa.jsonl", "techlead_qa_commitpack.jsonl", "techlead_qa.jsonl"]:
p = QNA_DIR / f
if not p.exists(): continue
for l in open(p):
d = json.loads(l); repo = d["doc_id"].split("@")[0]
if repo in mv_repos:
qa_by_repo[repo].append(d)
# embeddings with splits (only repos that have >=1 QA are trainable/evaluable)
emb_rows = []
for repo, r in mv_by_repo.items():
if repo not in qa_by_repo: continue
emb_rows.append({"doc_id": repo, "doc_version": "head", "split": split_of(repo),
"category": "aligned6", "doc_embedding": r["doc_embedding"]})
t = pa.table({k: [e[k] for e in emb_rows] for k in
["doc_id", "doc_version", "split", "category", "doc_embedding"]})
pq.write_table(t, EMBEDDINGS_DIR / "aligned6_embeddings.parquet")
# qna with qna_split
n_qa = 0
from collections import Counter
split_ct = Counter()
with (QNA_DIR / "aligned6_qna.jsonl").open("w") as out:
for repo, rows in qa_by_repo.items():
sp = split_of(repo)
if len(rows) > PER_REPO_CAP:
rows = rng.sample(rows, PER_REPO_CAP)
for i, d in enumerate(rows):
if sp == "train":
qsplit = "held_out" if rng.random() < 0.15 else "train"
else:
qsplit = "held_out" # cr_val/cr_test: all held out
out.write(json.dumps({"doc_id": repo, "doc_version": "head",
"split": sp, "qna_split": qsplit, "aspect": d.get("aspect", "?"),
"question": d.get("question", ""), "prefix": d["prefix"],
"target": d["target"], "lang": d.get("lang", "unknown")}) + "\n")
n_qa += 1
split_ct[sp] += 1
print(f"aligned6 dataset: {len(emb_rows)} repos (embeddings) | {n_qa} QA")
print(f" repo splits: {Counter(split_of(r) for r in qa_by_repo)}")
print(f" QA by split: {dict(split_ct)}")
print(f" -> aligned6_embeddings.parquet + aligned6_qna.jsonl")
if __name__ == "__main__":
main()
|