File size: 9,983 Bytes
646ccfc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 | """Fit a code-only label-score threshold on TRAIN and evaluate once on VALID."""
from __future__ import annotations
import argparse
import hashlib
import json
from pathlib import Path
import numpy as np
import pandas as pd
import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
from scripts.audit_sft_prefix_sensitivity import score_batch
from scripts.audit_code_only_label_signal import NEUTRAL_PREFIX
from scripts.eval_graph_feature_classifier import metrics
from src.contracts.prompt import code_for_sample
from src.rl.bounded_prompt import bounded_prompt
BASE_REVISION = "1cfa9a7208912126459214e8b04321603b3df60c"
CHECKPOINT = Path("checkpoints/sft-contract-v2/checkpoint-500")
def select_balanced_accuracy_threshold(y_true, scores):
"""Select threshold by TRAIN balanced accuracy; ties choose threshold nearest zero."""
y = np.asarray(y_true, dtype=int)
s = np.asarray(scores, dtype=float)
if len(y) != len(s) or not len(y) or set(np.unique(y)) != {0, 1}:
raise ValueError("threshold fitting requires aligned SAFE/VULNERABLE examples")
values = np.unique(s)
candidates = np.concatenate(([values[0] - 1e-9],
(values[:-1] + values[1:]) / 2,
[values[-1] + 1e-9]))
records = []
for threshold in candidates:
pred = (s >= threshold).astype(int)
safe_recall = float(np.mean(pred[y == 0] == 0))
vulnerable_recall = float(np.mean(pred[y == 1] == 1))
records.append((0.5 * (safe_recall + vulnerable_recall),
-abs(float(threshold)), float(threshold), safe_recall,
vulnerable_recall))
best = max(records)
return {"threshold": best[2], "balanced_accuracy": best[0],
"safe_accuracy": best[3], "vulnerable_accuracy": best[4],
"tie_break": "among equal balanced-accuracy thresholds, choose closest to zero"}
def _pair_rows(frame, pair_ids, split_name):
subset = frame[frame.pair_id.astype(str).isin(pair_ids)]
out = {}
for pair_id, group in subset.groupby("pair_id", sort=False):
if len(group) != 2 or set(group.target.astype(int)) != {0, 1}:
raise ValueError(f"{split_name} pair integrity failure: {pair_id}")
out[str(pair_id)] = {int(r.target): r for r in group.itertuples(index=False)}
if set(out) != set(pair_ids):
raise ValueError(f"some {split_name} pair IDs are missing from corpus")
return out
def run(*, train_pair_count=512, seed=42, batch_size=4,
output_path=Path("results/phase7/code_margin_threshold_train512_valid1104.json"),
corpus_path=Path("data/processed/corpus.parquet")):
train_all = json.loads(Path("data/splits/chronological/train.json").read_text())
valid_ids = json.loads(Path("data/splits/chronological/valid.json").read_text())
train_set, valid_set = set(train_all), set(valid_ids)
if train_set & valid_set:
raise ValueError("chronological TRAIN and VALID pair IDs overlap")
if train_pair_count > len(train_all):
raise ValueError("requested calibration pairs exceed chronological TRAIN")
rng = np.random.default_rng(seed)
selected_train = sorted(rng.choice(np.array(sorted(train_all), dtype=object),
size=train_pair_count, replace=False).tolist())
frame = pd.read_parquet(corpus_path, columns=["sample_id", "pair_id", "target", "func_vuln", "func_patch"])
train_rows = _pair_rows(frame, set(selected_train), "TRAIN")
valid_rows = _pair_rows(frame, valid_set, "VALID")
if set(train_rows) & set(valid_rows):
raise ValueError("sampled TRAIN and full VALID pairs overlap")
tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT, local_files_only=True)
tokenizer.pad_token = tokenizer.pad_token or tokenizer.eos_token
tokenizer.padding_side = "left"
entries, metadata = [], {}
for split, rows in (("train", train_rows), ("valid", valid_rows)):
for pair_id in sorted(rows):
for target in (0, 1):
row = rows[pair_id][target]
prompt, was_bounded = bounded_prompt(code_for_sample(row._asdict()), tokenizer)
if len(tokenizer.encode(prompt)) > 28672:
raise ValueError(f"bounded {split} prompt exceeds pinned token limit: {pair_id}")
key = f"{split}:{pair_id}:{target}"
entries.append({"key": key, "prompt": prompt, "prefix": NEUTRAL_PREFIX})
metadata[key] = {"split": split, "pair_id": pair_id, "target": target,
"was_bounded": was_bounded}
base = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen3-4B", revision=BASE_REVISION, torch_dtype=torch.bfloat16,
attn_implementation="sdpa", local_files_only=True).cuda()
model = PeftModel.from_pretrained(base, CHECKPOINT, local_files_only=True).eval()
scores = score_batch(model, tokenizer, entries, model.device, batch_size)
pair_margins = {"train": {}, "valid": {}}
for key, info in metadata.items():
vals = scores[key]
pair_margins[info["split"]].setdefault(info["pair_id"], {})[info["target"]] = (
vals["VULNERABLE"] - vals["SAFE"])
records = {}
for split, pair_map in pair_margins.items():
records[split] = []
for pair_id in sorted(pair_map):
for target in (0, 1):
records[split].append({"pair_id": pair_id, "target": target,
"margin": pair_map[pair_id][target]})
train_y = [r["target"] for r in records["train"]]
train_score = [r["margin"] for r in records["train"]]
chosen = select_balanced_accuracy_threshold(train_y, train_score)
threshold = chosen["threshold"]
def evaluation(split):
rows = records[split]
y = np.asarray([r["target"] for r in rows], dtype=int)
score = np.asarray([r["margin"] for r in rows], dtype=float)
pred = (score >= threshold).astype(int)
pair_ids = sorted({r["pair_id"] for r in rows})
by_pair = {}
for row, p in zip(rows, pred, strict=True):
by_pair.setdefault(row["pair_id"], {})[row["target"]] = int(p)
pair_acc = np.asarray([np.mean([by_pair[p][0] == 0, by_pair[p][1] == 1])
for p in pair_ids])
boot_rng = np.random.default_rng(seed + (0 if split == "train" else 1))
boot = pair_acc[boot_rng.integers(0, len(pair_acc),
size=(10000, len(pair_acc)))].mean(axis=1)
return {"metrics": metrics(y.tolist(), pred.tolist(), len(pair_ids)),
"balanced_accuracy": float(0.5 * (np.mean(pred[y == 0] == 0) +
np.mean(pred[y == 1] == 1))),
"pair_bootstrap_accuracy_95_ci": [float(x) for x in np.quantile(boot, [.025, .975])],
"predicted_vulnerable_fraction": float(pred.mean()),
"score_margin_quantiles": {str(q): float(np.quantile(score, q))
for q in [0, .1, .5, .9, 1]},
"per_row": [{**r, "prediction": int(p)} for r, p in zip(rows, pred, strict=True)]}
result = {"purpose": "threshold calibration on chronological TRAIN; one-shot evaluation on full chronological VALID",
"policy": "pinned Qwen3-4B plus frozen Contract v2 SFT checkpoint 500",
"base_revision": BASE_REVISION, "checkpoint": str(CHECKPOINT),
"neutral_scaffold": NEUTRAL_PREFIX,
"calibration": {"pair_count": len(train_rows), "row_count": len(records["train"]),
"seed": seed, "selected_threshold": chosen,
"objective": "maximum balanced accuracy on sampled TRAIN rows"},
"train": evaluation("train"), "valid": evaluation("valid"),
"bounded_prompt_counts": {split: sum(m["was_bounded"] for m in metadata.values()
if m["split"] == split)
for split in ("train", "valid")},
"data": {"corpus_sha256": hashlib.sha256(corpus_path.read_bytes()).hexdigest(),
"train_split_sha256": hashlib.sha256(Path("data/splits/chronological/train.json").read_bytes()).hexdigest(),
"valid_split_sha256": hashlib.sha256(Path("data/splits/chronological/valid.json").read_bytes()).hexdigest(),
"adapter_sha256": hashlib.sha256((CHECKPOINT / "adapter_model.safetensors").read_bytes()).hexdigest(),
"train_valid_pair_overlap": 0, "test_read": False},
"limitations": ["Neutral empty fields are counterfactual and may be out of distribution.",
"Threshold is calibrated on a deterministic 512-pair TRAIN sample.",
"This evaluates a frozen score readout only; no model weights or prompt were changed."]}
output_path.parent.mkdir(parents=True, exist_ok=True)
output_path.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
return result
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--train-pair-count", type=int, default=512)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--batch-size", type=int, default=4)
parser.add_argument("--output", type=Path,
default=Path("results/phase7/code_margin_threshold_train512_valid1104.json"))
args = parser.parse_args()
result = run(train_pair_count=args.train_pair_count, seed=args.seed,
batch_size=args.batch_size, output_path=args.output)
print(json.dumps({k: result[k] for k in ("calibration", "train", "valid") if k in result}, indent=2))
if __name__ == "__main__":
main()
|