| """Phase 1 §3.3 CLI: reads data/processed/corpus.parquet (+ the SecVulEval |
| GitHub-resolved commit dates, if available) and writes all three split |
| regimes under data/splits/. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import pandas as pd |
|
|
| from src.data.build_splits import ( |
| assert_pair_integrity, |
| build_chronological_split, |
| build_random_split, |
| build_repo_disjoint_folds, |
| _pair_table, |
| write_split, |
| ) |
| from src.utils.seeding import set_seed |
|
|
| SECVULEVAL_DATES_CACHE = Path("data/processed/secvuleval_commit_dates.jsonl") |
|
|
|
|
| def _merge_secvuleval_dates(df: pd.DataFrame, cache_path: Path = SECVULEVAL_DATES_CACHE) -> pd.DataFrame: |
| """Fills in `commit_date` for secvuleval rows from the GitHub-resolved |
| cache (src/data/fetch_commit_dates.py), where PrimeVul's Big-Vul join |
| didn't already provide one. Does not overwrite an existing date.""" |
| if not cache_path.exists(): |
| print(f"no commit-date cache at {cache_path}; secvuleval rows stay undated") |
| return df |
| resolved: dict[str, str] = {} |
| with open(cache_path) as f: |
| for line in f: |
| rec = json.loads(line) |
| if rec.get("date"): |
| resolved[rec["commit_id"]] = rec["date"] |
| print(f"loaded {len(resolved)} resolved commit dates from {cache_path}") |
|
|
| df = df.copy() |
| needs_date = df["commit_date"].isna() |
| df.loc[needs_date, "commit_date"] = df.loc[needs_date, "commit_id"].map(resolved) |
| return df |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--seed", type=int, default=42) |
| parser.add_argument("--corpus", default="data/processed/corpus.parquet") |
| parser.add_argument("--chrono-train-before", default="2022-01-01") |
| parser.add_argument("--chrono-valid-before", default="2023-01-01") |
| parser.add_argument("--n-folds", type=int, default=5) |
| parser.add_argument("--results-dir", default="results/phase1") |
| args = parser.parse_args() |
| set_seed(args.seed) |
|
|
| results_dir = Path(args.results_dir) |
| results_dir.mkdir(parents=True, exist_ok=True) |
|
|
| df = pd.read_parquet(args.corpus) |
| df = _merge_secvuleval_dates(df) |
| pairs = _pair_table(df) |
| print(f"{len(pairs)} pairs total; date coverage now {pairs['commit_date'].notna().mean():.1%}") |
|
|
| random_split = build_random_split(pairs, seed=args.seed) |
| assert_pair_integrity(random_split) |
| write_split("random", random_split) |
| print(f"random split: {[(k, len(v)) for k, v in random_split.items()]}") |
|
|
| chrono_split, chrono_stats = build_chronological_split( |
| pairs, args.chrono_train_before, args.chrono_valid_before |
| ) |
| assert_pair_integrity(chrono_split) |
| write_split("chronological", chrono_split) |
| print(f"chronological split: {chrono_stats}") |
|
|
| folds = build_repo_disjoint_folds(pairs, n_folds=args.n_folds, seed=args.seed) |
| for i, fold in enumerate(folds): |
| assert_pair_integrity(fold) |
| write_split(f"repo_disjoint_fold{i}", fold) |
| print(f"repo-disjoint: {args.n_folds} folds written") |
|
|
| summary = { |
| "n_pairs_total": len(pairs), |
| "date_coverage_after_secvuleval_merge": float(pairs["commit_date"].notna().mean()), |
| "random_split_sizes": {k: len(v) for k, v in random_split.items()}, |
| "chronological_split_stats": chrono_stats, |
| "repo_disjoint_fold_sizes": [ |
| {k: len(v) for k, v in fold.items()} for fold in folds |
| ], |
| } |
| (results_dir / "split_summary.json").write_text(json.dumps(summary, indent=2)) |
| print(json.dumps(summary, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
