| """Check subgroup patterns using forward out-of-fold chronological TRAIN only.""" |
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import pandas as pd |
| from sklearn.metrics import roc_curve |
| from sklearn.feature_extraction.text import TfidfVectorizer |
| from sklearn.linear_model import SGDClassifier |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| CORPUS = ROOT / "data/processed/corpus.parquet" |
| TRAIN = ROOT / "data/splits/chronological/train.json" |
| DATE_CACHE = ROOT / "data/processed/secvuleval_commit_dates.jsonl" |
| OUT_JSON = ROOT / "results/phase7/char_forward_train_subgroups.json" |
| OUT_MD = ROOT / "results/phase7/char_forward_train_subgroups.md" |
| SEEDS = tuple(range(10)) |
| FOLDS = (2019, 2020, 2021) |
|
|
|
|
| def wilson(k: int, n: int) -> list[float] | None: |
| if not n: |
| return None |
| z, p = 1.96, k / n |
| d = 1 + z*z/n |
| c = p + z*z/(2*n) |
| h = z*np.sqrt(p*(1-p)/n + z*z/(4*n*n)) |
| return [float((c-h)/d), float((c+h)/d)] |
|
|
|
|
| def metric(mask: np.ndarray, y: np.ndarray, pred: np.ndarray) -> dict: |
| safe = mask & (y == 0) |
| vuln = mask & (y == 1) |
| good = pred == y |
| n, k = int(mask.sum()), int(np.sum(mask & good)) |
| ks, kv = int(np.sum(safe & good)), int(np.sum(vuln & good)) |
| return { |
| "n_rows": n, |
| "n_pairs": n // 2, |
| "accuracy": k/n if n else None, |
| "wilson_95": wilson(k, n), |
| "safe_n": int(safe.sum()), |
| "safe_accuracy": ks/int(safe.sum()) if safe.any() else None, |
| "safe_wilson_95": wilson(ks, int(safe.sum())), |
| "vulnerable_n": int(vuln.sum()), |
| "vulnerable_accuracy": kv/int(vuln.sum()) if vuln.any() else None, |
| "vulnerable_wilson_95": wilson(kv, int(vuln.sum())), |
| } |
|
|
|
|
| def fit_score(fit: pd.DataFrame, held: pd.DataFrame, |
| sample_weight: np.ndarray | None = None) -> np.ndarray: |
| vectorizer = TfidfVectorizer(analyzer="char", ngram_range=(3, 5), min_df=2, |
| max_features=300_000, sublinear_tf=True, |
| dtype=np.float32) |
| xfit = vectorizer.fit_transform(fit.text.tolist()) |
| xheld = vectorizer.transform(held.text.tolist()) |
| yfit = fit.target.to_numpy(dtype=np.int8) |
| all_scores = [] |
| for seed in SEEDS: |
| model = SGDClassifier(loss="log_loss", alpha=1e-5, max_iter=200, |
| tol=1e-3, random_state=seed, average=True) |
| model.fit(xfit, yfit, sample_weight=sample_weight) |
| all_scores.append(model.decision_function(xheld)) |
| return np.mean(np.stack(all_scores), axis=0) |
|
|
|
|
| def length_label_weights(fit: pd.DataFrame) -> np.ndarray: |
| """Equalize the six observed label × fixed-length cells using TRAIN only.""" |
| labels = fit.target.to_numpy(dtype=np.int8) |
| lines = fit.vuln_lines.to_numpy() |
| bins = np.where(lines <= 30, 0, np.where(lines <= 100, 1, 2)) |
| cell = labels * 3 + bins |
| counts = np.bincount(cell, minlength=6) |
| present = counts > 0 |
| n_present = int(present.sum()) |
| weights = np.asarray([len(fit) / (n_present * counts[c]) for c in cell], |
| dtype=np.float64) |
| if not np.isclose(weights.mean(), 1.0): |
| raise AssertionError("sample weights must have mean one") |
| return weights |
|
|
|
|
| def select_threshold(scores: np.ndarray, labels: np.ndarray) -> float: |
| """Select a balanced-accuracy threshold on calibration-only predictions.""" |
| fpr, tpr, thresholds = roc_curve(labels, scores, pos_label=1) |
| finite = np.isfinite(thresholds) |
| thresholds, fpr, tpr = thresholds[finite], fpr[finite], tpr[finite] |
| prevalence = float(np.mean(labels == 1)) |
| ba = (tpr + 1 - fpr) / 2 |
| accuracy = prevalence * tpr + (1 - prevalence) * (1 - fpr) |
| order = np.lexsort((- |
| np.abs(thresholds), accuracy, ba)) |
| return float(thresholds[order[-1]]) |
|
|
|
|
| def nested_length_thresholds(oof: pd.DataFrame) -> list[dict]: |
| """Calibrate per fixed length band on earlier OOF years, score next year.""" |
| reports = [] |
| rng = np.random.default_rng(42) |
| bands = (("short_le_30", 1, 30), ("medium_31_100", 31, 100), |
| ("long_gt_100", 101, None)) |
| for year in FOLDS: |
| calibration = oof[oof.year < year] |
| held = oof[oof.year == year] |
| if calibration.empty: |
| continue |
| cy = calibration.target.to_numpy(dtype=np.int8) |
| cs = calibration.score.to_numpy() |
| global_threshold = select_threshold(cs, cy) |
| thresholds = {} |
| for name, low, high in bands: |
| cmask = calibration.vuln_lines.to_numpy() >= low |
| if high is not None: |
| cmask &= calibration.vuln_lines.to_numpy() <= high |
| counts = [int(np.sum(cmask & (cy == label))) for label in (0, 1)] |
| if min(counts) >= 50: |
| threshold = select_threshold(cs[cmask], cy[cmask]) |
| fallback = False |
| else: |
| threshold = global_threshold |
| fallback = True |
| thresholds[name] = {"threshold": threshold, |
| "calibration_safe_n": counts[0], |
| "calibration_vulnerable_n": counts[1], |
| "pooled_threshold_fallback": fallback} |
|
|
| ey = held.target.to_numpy(dtype=np.int8) |
| es = held.score.to_numpy() |
| lens = held.vuln_lines.to_numpy() |
| pred = np.empty(len(held), dtype=np.int8) |
| for name, low, high in bands: |
| mask = lens >= low |
| if high is not None: |
| mask &= lens <= high |
| pred[mask] = (es[mask] >= thresholds[name]["threshold"]).astype(np.int8) |
| fixed = (es >= 0).astype(np.int8) |
| pair_effects = [] |
| rows = pd.DataFrame({"pair_id": held.pair_id.to_numpy(), "y": ey, |
| "fixed": fixed, "calibrated": pred}) |
| for _, group in rows.groupby("pair_id", sort=True): |
| if len(group) != 2 or set(group.y.astype(int)) != {0, 1}: |
| raise ValueError("pair integrity failed in threshold comparison") |
| pair_effects.append(float(np.mean(group.calibrated.to_numpy() == group.y.to_numpy()) - |
| np.mean(group.fixed.to_numpy() == group.y.to_numpy()))) |
| effects = np.asarray(pair_effects) |
| draws = effects[rng.integers(0, len(effects), size=(10_000, len(effects)))].mean(axis=1) |
| reports.append({ |
| "held_year": int(year), |
| "calibration_years": sorted(int(y) for y in calibration.year.unique()), |
| "calibration_pairs": int(calibration.pair_id.nunique()), |
| "thresholds_by_length": thresholds, |
| "fixed_zero_threshold": metric(np.ones(len(held), dtype=bool), ey, fixed), |
| "length_calibrated": metric(np.ones(len(held), dtype=bool), ey, pred), |
| "balanced_accuracy_delta_pp": float(100 * effects.mean()), |
| "pair_bootstrap_95_ci_delta_pp": [ |
| float(100 * x) for x in np.quantile(draws, [0.025, 0.975])], |
| }) |
| return reports |
|
|
|
|
| def main() -> None: |
| frame = pd.read_parquet(CORPUS) |
| date_map = {} |
| for line in DATE_CACHE.read_text().splitlines(): |
| obj = json.loads(line) |
| if obj.get("date"): |
| date_map[str(obj["commit_id"])] = str(obj["date"]) |
| missing = frame.commit_date.isna() |
| frame.loc[missing, "commit_date"] = frame.loc[missing, "commit_id"].astype(str).map(date_map) |
| frame["year"] = pd.to_numeric(frame.commit_date.astype(str).str[:4], errors="coerce") |
| train_ids = set(map(str, json.loads(TRAIN.read_text()))) |
| sub = frame[frame.pair_id.astype(str).isin(train_ids)] |
| records = [] |
| for pid, grp in sub.groupby("pair_id", sort=True): |
| if len(grp) != 2 or set(grp.target.astype(int)) != {0, 1}: |
| raise ValueError(f"pair integrity failure: {pid}") |
| if grp.func_vuln.nunique() != 1 or grp.func_patch.nunique() != 1: |
| raise ValueError(f"paired code mismatch: {pid}") |
| if grp.cwe.nunique(dropna=False) != 1 or grp.year.nunique(dropna=False) != 1: |
| raise ValueError(f"paired metadata mismatch: {pid}") |
| vrow = grp[grp.target.astype(int) == 1].iloc[0] |
| for target in (0, 1): |
| row = grp[grp.target.astype(int) == target].iloc[0] |
| records.append({ |
| "pair_id": str(pid), "target": target, "year": row.year, |
| "text": str(row.func_vuln if target else row.func_patch), |
| "cwe": str(vrow.cwe) if pd.notna(vrow.cwe) else "UNKNOWN", |
| "vuln_lines": max(1, len(str(vrow.func_vuln).splitlines())), |
| }) |
| train = pd.DataFrame(records) |
| oof_parts = [] |
| weighted_parts = [] |
| fold_info = [] |
| for year in FOLDS: |
| fit = train[train.year < year].reset_index(drop=True) |
| held = train[train.year == year].reset_index(drop=True) |
| if fit.empty or held.empty or held.target.nunique() != 2: |
| raise ValueError(f"invalid forward fold {year}") |
| scores = fit_score(fit, held) |
| weighted_scores = fit_score(fit, held, length_label_weights(fit)) |
| held = held.copy() |
| held["score"] = scores |
| held["prediction"] = (scores >= 0).astype(np.int8) |
| oof_parts.append(held) |
| weighted_held = held.drop(columns=["score", "prediction"]).copy() |
| weighted_held["score"] = weighted_scores |
| weighted_held["prediction"] = (weighted_scores >= 0).astype(np.int8) |
| weighted_parts.append(weighted_held) |
| fold_info.append({"held_year": year, |
| "fit_pairs": int(fit.pair_id.nunique()), |
| "held_pairs": int(held.pair_id.nunique()), |
| "metrics": metric(np.ones(len(held), dtype=bool), |
| held.target.to_numpy(), |
| held.prediction.to_numpy())}) |
| print(json.dumps({"held_year": year, |
| "fit_pairs": int(fit.pair_id.nunique()), |
| "held_pairs": int(held.pair_id.nunique())}), flush=True) |
| oof = pd.concat(oof_parts, ignore_index=True) |
| weighted_oof = pd.concat(weighted_parts, ignore_index=True) |
| y = oof.target.to_numpy(dtype=np.int8) |
| pred = oof.prediction.to_numpy(dtype=np.int8) |
| all_mask = np.ones(len(oof), dtype=bool) |
| nested_calibration = nested_length_thresholds(oof) |
|
|
| |
| wy = weighted_oof.target.to_numpy(dtype=np.int8) |
| wp = weighted_oof.prediction.to_numpy(dtype=np.int8) |
| base_rows = oof[["pair_id", "year", "target", "vuln_lines", "prediction"]].copy() |
| base_rows = base_rows.rename(columns={"prediction": "base_prediction"}) |
| weighted_rows = weighted_oof[["pair_id", "target", "prediction"]].rename( |
| columns={"prediction": "weighted_prediction"}) |
| paired_compare = base_rows.merge(weighted_rows, on=["pair_id", "target"], |
| validate="one_to_one") |
| if len(paired_compare) != len(oof): |
| raise AssertionError("weighted and baseline OOF row alignment failed") |
| pair_effects = [] |
| for _, group in paired_compare.groupby("pair_id", sort=True): |
| if len(group) != 2 or set(group.target.astype(int)) != {0, 1}: |
| raise ValueError("pair integrity failure in weighted comparison") |
| base_ok = group.base_prediction.to_numpy() == group.target.to_numpy() |
| weighted_ok = group.weighted_prediction.to_numpy() == group.target.to_numpy() |
| pair_effects.append(float(weighted_ok.mean() - base_ok.mean())) |
| pair_effects = np.asarray(pair_effects) |
| rng = np.random.default_rng(77) |
| pair_draws = pair_effects[rng.integers(0, len(pair_effects), |
| size=(10_000, len(pair_effects)))].mean(axis=1) |
| weighted_comparison = { |
| "weighting": "equal total sample weight for each observed target × fixed length band cell", |
| "baseline": metric(np.ones(len(oof), dtype=bool), y, pred), |
| "length_label_weighted": metric(np.ones(len(weighted_oof), dtype=bool), wy, wp), |
| "paired_accuracy_delta_pp": float(100 * pair_effects.mean()), |
| "pair_bootstrap_95_ci_delta_pp": [ |
| float(100 * v) for v in np.quantile(pair_draws, [0.025, 0.975])], |
| "by_held_year": {}, |
| "by_fixed_length_band": {}, |
| } |
| for year in FOLDS: |
| bm = oof.year.to_numpy() == year |
| wm = weighted_oof.year.to_numpy() == year |
| weighted_comparison["by_held_year"][str(year)] = { |
| "baseline": metric(bm, y, pred), |
| "length_label_weighted": metric(wm, wy, wp), |
| } |
| for name, low, high in (("short_le_30", 1, 30), |
| ("medium_31_100", 31, 100), |
| ("long_gt_100", 101, None)): |
| lm = oof.vuln_lines.to_numpy() >= low |
| if high is not None: |
| lm &= oof.vuln_lines.to_numpy() <= high |
| wlm = weighted_oof.vuln_lines.to_numpy() >= low |
| if high is not None: |
| wlm &= weighted_oof.vuln_lines.to_numpy() <= high |
| weighted_comparison["by_fixed_length_band"][name] = { |
| "baseline": metric(lm, y, pred), |
| "length_label_weighted": metric(wlm, wy, wp), |
| } |
|
|
| cwe_results = [] |
| for cwe, inds in oof.groupby("cwe", sort=True).groups.items(): |
| mask = np.zeros(len(oof), dtype=bool) |
| mask[np.asarray(list(inds), dtype=int)] = True |
| by_year = {} |
| for year in FOLDS: |
| ym = mask & (oof.year.to_numpy() == year) |
| by_year[str(year)] = metric(ym, y, pred) |
| cwe_results.append({"cwe": cwe, "pooled": metric(mask, y, pred), |
| "by_held_year": by_year}) |
|
|
| length_results = [] |
| bands = [("short_le_30", 1, 30), ("medium_31_100", 31, 100), |
| ("long_gt_100", 101, None)] |
| for name, low, high in bands: |
| lens = oof.vuln_lines.to_numpy() |
| mask = (lens >= low) & ((lens <= high) if high is not None else True) |
| by_year = {str(year): metric(mask & (oof.year.to_numpy() == year), y, pred) |
| for year in FOLDS} |
| length_results.append({"band": name, "min_lines": low, |
| "max_lines": high, "pooled": metric(mask, y, pred), |
| "by_held_year": by_year}) |
|
|
| result = { |
| "method": "fixed ten-seed char TF-IDF/SGD, forward TRAIN OOF, zero threshold", |
| "folds": fold_info, |
| "pooled_oof": metric(all_mask, y, pred), |
| "cwe_slices": cwe_results, |
| "vulnerable_function_length_bands": length_results, |
| "nested_forward_length_calibration": nested_calibration, |
| "length_label_weighting_comparison": weighted_comparison, |
| "coverage": { |
| "train_pairs": len(train_ids), |
| "oof_pairs": int(oof.pair_id.nunique()), |
| "oof_rows": len(oof), |
| "undated_or_other_year_pairs_not_in_folds": len(train_ids - set(oof.pair_id)), |
| }, |
| "protocol": { |
| "held_years": list(FOLDS), "seeds": list(SEEDS), |
| "threshold": 0.0, |
| "subgroup_analysis_is_descriptive": True, |
| "length_bands_prespecified_lines": [30, 100], |
| "valid_used_for_selection": False, "test_read": False, |
| }, |
| "provenance": { |
| "corpus_sha256": hashlib.sha256(CORPUS.read_bytes()).hexdigest(), |
| "train_split_sha256": hashlib.sha256(TRAIN.read_bytes()).hexdigest(), |
| "date_cache_sha256": hashlib.sha256(DATE_CACHE.read_bytes()).hexdigest(), |
| }, |
| "limitations": [ |
| "Only TRAIN pairs dated 2019–2021 with earlier dated TRAIN history are scored.", |
| "CWE labels may be incomplete/noisy; small subgroup estimates are unstable.", |
| "Forward folds share overlapping fit histories and are not independent replications.", |
| "Per-row Wilson intervals do not account for repository clustering.", |
| ], |
| } |
| OUT_JSON.parent.mkdir(parents=True, exist_ok=True) |
| OUT_JSON.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n") |
| lines = ["# Forward TRAIN subgroup replication", "", |
| "Fixed model and zero threshold; all rows are forward out-of-fold TRAIN predictions. " |
| "No VALID selection or TEST access.", "", |
| f"Pooled: {result['pooled_oof']['accuracy']:.1%} accuracy " |
| f"(Wilson 95% {result['pooled_oof']['wilson_95'][0]:.1%}–" |
| f"{result['pooled_oof']['wilson_95'][1]:.1%}); SAFE " |
| f"{result['pooled_oof']['safe_accuracy']:.1%}, VULNERABLE " |
| f"{result['pooled_oof']['vulnerable_accuracy']:.1%}; " |
| f"{result['coverage']['oof_pairs']} pairs.", "", |
| "## Prespecified function-length bands", "", |
| "| Band | pairs | SAFE accuracy | VULNERABLE accuracy | overall |", |
| "|---|---:|---:|---:|---:|"] |
| for item in length_results: |
| m = item["pooled"] |
| lines.append(f"| {item['band']} | {m['n_pairs']} | {m['safe_accuracy']:.1%} | " |
| f"{m['vulnerable_accuracy']:.1%} | {m['accuracy']:.1%} |") |
| lines += ["", "## CWE slices with at least 30 pooled OOF pairs", "", |
| "| CWE | pairs | SAFE accuracy | VULNERABLE accuracy | held-year vulnerable recalls |", |
| "|---|---:|---:|---:|---|"] |
| for item in cwe_results: |
| m = item["pooled"] |
| if m["n_pairs"] < 30: |
| continue |
| yearly = ", ".join(f"{yr}: {v['vulnerable_accuracy']:.1%}" |
| for yr, v in item["by_held_year"].items() |
| if v["vulnerable_n"]) |
| lines.append(f"| {item['cwe']} | {m['n_pairs']} | {m['safe_accuracy']:.1%} | " |
| f"{m['vulnerable_accuracy']:.1%} | {yearly} |") |
| lines += ["", "## Nested prior-year length calibration", "", |
| "For each held year, thresholds use only earlier-year OOF scores; year 2019 has no prior OOF calibration and is omitted.", "", |
| "| Held year | fixed balanced acc. | calibrated balanced acc. | delta pp (pair-bootstrap 95% CI) |", |
| "|---:|---:|---:|---:|"] |
| for item in nested_calibration: |
| fixed, calibrated = item["fixed_zero_threshold"], item["length_calibrated"] |
| fixed_ba = (fixed["safe_accuracy"] + fixed["vulnerable_accuracy"]) / 2 |
| cal_ba = (calibrated["safe_accuracy"] + calibrated["vulnerable_accuracy"]) / 2 |
| ci = item["pair_bootstrap_95_ci_delta_pp"] |
| lines.append(f"| {item['held_year']} | {fixed_ba:.2%} | {cal_ba:.2%} | " |
| f"{item['balanced_accuracy_delta_pp']:+.2f} ({ci[0]:+.2f}, {ci[1]:+.2f}) |") |
| lines += ["", "## TRAIN-only length × label sample weighting", "", |
| f"Pooled paired accuracy delta: {weighted_comparison['paired_accuracy_delta_pp']:+.2f} pp " |
| f"(pair-bootstrap 95% CI {weighted_comparison['pair_bootstrap_95_ci_delta_pp'][0]:+.2f}, " |
| f"{weighted_comparison['pair_bootstrap_95_ci_delta_pp'][1]:+.2f}).", "", |
| "| Held year | baseline accuracy | weighted accuracy |", |
| "|---:|---:|---:|"] |
| for year, values in weighted_comparison["by_held_year"].items(): |
| lines.append(f"| {year} | {values['baseline']['accuracy']:.2%} | " |
| f"{values['length_label_weighted']['accuracy']:.2%} |") |
| lines += ["", "This is one forward TRAIN diagnostic; length × label weights use only the corresponding fit fold. " |
| "CWE and length slices are descriptive; no VALID selection or TEST access.", ""] |
| OUT_MD.write_text("\n".join(lines)) |
| print(json.dumps({"pooled_oof": result["pooled_oof"], |
| "coverage": result["coverage"], |
| "artifacts": [str(OUT_JSON), str(OUT_MD)]}, indent=2), flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
