File size: 2,754 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 | #!/usr/bin/env python3
"""Merge the real Code2LoRA/RepoPeftBench corpus (73,849 real repo-commit
docs, 443,798 real assertion-completion QnAs) with our synthetic corpus
(211 docs: the Code2LoRA paper + coding-agent-harness + agile/Jira +
general fact-sheets) into one combined training set for the hypernetwork.
Output:
data/embeddings/combined_embeddings.parquet
data/qna/combined_qna.jsonl
Usage:
python scripts/merge_corpora.py
"""
from __future__ import annotations
import json
import sys
from pathlib import Path
import pyarrow as pa
import 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, ensure_dirs # noqa: E402
def merge_embeddings() -> int:
synthetic_path = EMBEDDINGS_DIR / "doc_embeddings.parquet"
real_path = EMBEDDINGS_DIR / "real_code2lora_embeddings.parquet"
out_path = EMBEDDINGS_DIR / "combined_embeddings.parquet"
tables = []
for p, label in [(synthetic_path, "synthetic"), (real_path, "real")]:
if not p.exists():
print(f" [skip] {p} not found", flush=True)
continue
t = pq.read_table(p)
print(f" {label}: {t.num_rows} rows, columns={t.column_names}", flush=True)
tables.append(t.select(["doc_id", "doc_version", "split", "category", "doc_embedding"]))
combined = pa.concat_tables(tables)
pq.write_table(combined, out_path)
print(f"Wrote {combined.num_rows} combined embeddings -> {out_path}", flush=True)
return combined.num_rows
def merge_qna() -> int:
synthetic_path = QNA_DIR / "qna.jsonl"
real_path = QNA_DIR / "real_code2lora_qna.jsonl"
out_path = QNA_DIR / "combined_qna.jsonl"
n = 0
with out_path.open("w") as out:
for p, label in [(synthetic_path, "synthetic"), (real_path, "real")]:
if not p.exists():
print(f" [skip] {p} not found", flush=True)
continue
count = 0
with p.open() as f:
for line in f:
line = line.strip()
if not line:
continue
out.write(line + "\n")
count += 1
n += 1
print(f" {label}: {count} QnA rows", flush=True)
print(f"Wrote {n} combined QnA pairs -> {out_path}", flush=True)
return n
def main() -> None:
ensure_dirs()
print("Merging embeddings...", flush=True)
n_docs = merge_embeddings()
print("\nMerging QnA...", flush=True)
n_qna = merge_qna()
print(f"\nCombined corpus: {n_docs} documents, {n_qna} QnA pairs.", flush=True)
if __name__ == "__main__":
main()
|