| """Phase 1 §3.3: builds the three split regimes (random, chronological, |
| repo-disjoint) plus the OOD manifest placeholder, at the pair level so the |
| pair-integrity constraint (§3.3.4: a pair_id never crosses partitions) holds |
| by construction. |
| |
| Each split is written to data/splits/<regime>/{train,valid,test}.json as a |
| list of pair_id (not sample_id — both halves of a pair always travel |
| together, so pair_id is the right unit of assignment and the smaller file). |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import random |
| from collections import defaultdict |
| from pathlib import Path |
|
|
| import pandas as pd |
|
|
| SPLITS_DIR = Path("data/splits") |
|
|
|
|
| def _pair_table(df: pd.DataFrame) -> pd.DataFrame: |
| """One row per pair_id, carrying the fields splits are built from.""" |
| vuln_rows = df[df["target"] == 1].drop_duplicates("pair_id") |
| return vuln_rows[["pair_id", "project", "source", "commit_date"]].reset_index(drop=True) |
|
|
|
|
| def build_random_split( |
| pairs: pd.DataFrame, seed: int = 42, ratios: tuple[float, float, float] = (0.8, 0.1, 0.1) |
| ) -> dict[str, list[str]]: |
| """80/10/10 by pair, stratified by source so each split keeps a |
| comparable primevul/secvuleval mix (spec §3.3.1).""" |
| rng = random.Random(seed) |
| splits: dict[str, list[str]] = {"train": [], "valid": [], "test": []} |
| for _source, group in pairs.groupby("source"): |
| ids = group["pair_id"].tolist() |
| rng.shuffle(ids) |
| n = len(ids) |
| n_train = int(n * ratios[0]) |
| n_valid = int(n * ratios[1]) |
| splits["train"].extend(ids[:n_train]) |
| splits["valid"].extend(ids[n_train : n_train + n_valid]) |
| splits["test"].extend(ids[n_train + n_valid :]) |
| return splits |
|
|
|
|
| def build_chronological_split( |
| pairs: pd.DataFrame, |
| train_before: str, |
| valid_before: str, |
| ) -> tuple[dict[str, list[str]], dict]: |
| """Train on commit_date < train_before, valid in [train_before, valid_before), |
| test >= valid_before (spec §3.3.2). Only pairs with a known commit_date are |
| eligible; pairs without one are returned separately as `undated` so callers |
| can decide whether to fold them into train or exclude them (we exclude by |
| default — undated data must not leak into a label that claims to be |
| time-ordered).""" |
| dated = pairs[pairs["commit_date"].notna()].copy() |
| |
| |
| |
| |
| |
| |
| |
| dated["commit_date"] = pd.to_datetime(dated["commit_date"], utc=True, format="mixed", errors="coerce") |
| dated = dated[dated["commit_date"].notna()] |
|
|
| train_cut = pd.Timestamp(train_before, tz="UTC") |
| valid_cut = pd.Timestamp(valid_before, tz="UTC") |
|
|
| train_ids = dated[dated["commit_date"] < train_cut]["pair_id"].tolist() |
| valid_ids = dated[(dated["commit_date"] >= train_cut) & (dated["commit_date"] < valid_cut)]["pair_id"].tolist() |
| test_ids = dated[dated["commit_date"] >= valid_cut]["pair_id"].tolist() |
|
|
| stats = { |
| "n_total_pairs": len(pairs), |
| "n_dated_pairs": len(dated), |
| "n_undated_pairs_excluded": len(pairs) - len(dated), |
| "n_train": len(train_ids), |
| "n_valid": len(valid_ids), |
| "n_test": len(test_ids), |
| "train_before": train_before, |
| "valid_before": valid_before, |
| } |
| return {"train": train_ids, "valid": valid_ids, "test": test_ids}, stats |
|
|
|
|
| def build_repo_disjoint_folds( |
| pairs: pd.DataFrame, n_folds: int = 5, seed: int = 42 |
| ) -> list[dict[str, list[str]]]: |
| """Assigns every pair of a project to exactly one fold (greedy balanced by |
| vulnerable-sample count per project; `pairs` here is already the |
| one-row-per-pair, target=1 table, so "vulnerable-sample count" == pair |
| count for that project). Returns `n_folds` {train, valid, test} splits, |
| each using one fold as test, the next as valid, the rest as train |
| (spec §3.3.3: "5 folds stored for cross-validation").""" |
| by_project: dict[str, list[str]] = defaultdict(list) |
| for row in pairs.itertuples(): |
| by_project[row.project].append(row.pair_id) |
|
|
| rng = random.Random(seed) |
| projects = list(by_project.keys()) |
| rng.shuffle(projects) |
| projects.sort(key=lambda p: -len(by_project[p])) |
|
|
| fold_sizes = [0] * n_folds |
| fold_projects: list[list[str]] = [[] for _ in range(n_folds)] |
| for project in projects: |
| i = min(range(n_folds), key=lambda k: fold_sizes[k]) |
| fold_projects[i].append(project) |
| fold_sizes[i] += len(by_project[project]) |
|
|
| fold_pair_ids = [ |
| [pid for project in fp for pid in by_project[project]] for fp in fold_projects |
| ] |
|
|
| folds = [] |
| for k in range(n_folds): |
| test_ids = fold_pair_ids[k] |
| valid_ids = fold_pair_ids[(k + 1) % n_folds] |
| train_ids = [ |
| pid for j in range(n_folds) if j not in (k, (k + 1) % n_folds) for pid in fold_pair_ids[j] |
| ] |
| folds.append({"train": train_ids, "valid": valid_ids, "test": test_ids}) |
| return folds |
|
|
|
|
| def assert_pair_integrity(split: dict[str, list[str]]) -> None: |
| train, valid, test = set(split["train"]), set(split["valid"]), set(split["test"]) |
| assert not (train & valid), "pair_id leaked between train/valid" |
| assert not (train & test), "pair_id leaked between train/test" |
| assert not (valid & test), "pair_id leaked between valid/test" |
|
|
|
|
| def write_split(name: str, split: dict[str, list[str]], out_dir: Path = SPLITS_DIR) -> None: |
| d = out_dir / name |
| d.mkdir(parents=True, exist_ok=True) |
| for part, ids in split.items(): |
| (d / f"{part}.json").write_text(json.dumps(sorted(ids))) |
|
|
