| """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() |
|
|
