Download scripts/analyze_char_forward_train_subgroups.py from Dhruv1000/tracegrpo: direct link, hf CLI and curl.
- Browser
- Download file 19.9 kB
-
https://huggingface.co/Dhruv1000/tracegrpo/resolve/main/scripts/analyze_char_forward_train_subgroups.py
- Command line
-
hf download hf://Dhruv1000/tracegrpo/scripts/analyze_char_forward_train_subgroups.py
-
curl -L -o analyze_char_forward_train_subgroups.py https://huggingface.co/Dhruv1000/tracegrpo/resolve/main/scripts/analyze_char_forward_train_subgroups.py
19.9 kB
| """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) | |
| # Paired out-of-fold comparison for TRAIN-only length × label weighting. | |
| 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() | |