the-jashthakkar/CodeMode

1'''2 3Aggregate synthetic datasets from multiple runs into a single combined dataset generated using triplets_synthesis.py.4 5'''6 7import json8from pathlib import Path9from datetime import datetime10from typing import List, Dict11 12BASE_SYNTHETIC_DIR = Path("data/synthetic")13OUTPUT_DIR = BASE_SYNTHETIC_DIR / "combined"14 15 16def load_jsonl(path: Path) -> List[Dict]:17    with path.open("r", encoding="utf-8") as f:18        return [json.loads(line) for line in f]19 20 21def save_jsonl(path: Path, records: List[Dict]):22    path.parent.mkdir(parents=True, exist_ok=True)23    with path.open("w", encoding="utf-8") as f:24        for r in records:25            f.write(json.dumps(r, ensure_ascii=False) + "\n")26 27 28def save_json(path: Path, records: List[Dict]):29    path.parent.mkdir(parents=True, exist_ok=True)30    with path.open("w", encoding="utf-8") as f:31        json.dump(records, f, indent=2)32 33 34def aggregate():35    positive_pairs_all = []36    triplets_all = []37    included_runs = []38 39    for run_dir in BASE_SYNTHETIC_DIR.iterdir():40        if not run_dir.is_dir():41            continue42        if run_dir.name == "combined":43            continue44 45        pos_path = run_dir / "positive_pairs.jsonl"46        tri_path = run_dir / "triplets.jsonl"47 48        if pos_path.exists() and tri_path.exists():49            positive_pairs_all.extend(load_jsonl(pos_path))50            triplets_all.extend(load_jsonl(tri_path))51            included_runs.append(run_dir.name)52 53    # Save JSONL (training)54    save_jsonl(OUTPUT_DIR / "positive_pairs.jsonl", positive_pairs_all)55    save_jsonl(OUTPUT_DIR / "triplets.jsonl", triplets_all)56 57    # Save JSON (inspection / upload)58    save_json(OUTPUT_DIR / "positive_pairs.json", positive_pairs_all)59    save_json(OUTPUT_DIR / "triplets.json", triplets_all)60 61    # Metadata62    metadata = {63        "type": "combined_dataset",64        "included_runs": included_runs,65        "total_positive_pairs": len(positive_pairs_all),66        "total_triplets": len(triplets_all),67        "created_at": datetime.utcnow().isoformat(),68    }69 70    with (OUTPUT_DIR / "metadata.json").open("w", encoding="utf-8") as f:71        json.dump(metadata, f, indent=2)72 73    print("✅ Combined dataset created at:", OUTPUT_DIR)74 75 76if __name__ == "__main__":77    aggregate()78