#!/usr/bin/env python3 """Build held-out scoring sets for eval_endpoint.py. Two modes: vuln_detection -- from the reserved `eval_split` (test) of the training sources. Each row -> {messages: [system, user], gold_label, cwe}. The model is asked the same detection question used in training; the gold label comes straight from the dataset's ground truth. mcq -- from an eval-only MCQ source (disjoint from training). Each row -> {question, choices, gold_index, gold_letter}. These sets are never used for training and must be passed through decontaminate.py's eval side, not its train side. Examples: # 1) download the held-out test splits first: python training/scripts/hf_download.py --key primevul --eval python training/scripts/hf_download.py --key megavul --eval python training/scripts/build_eval_sets.py --mode vuln_detection \ --out data/eval/vuln_detection_test.jsonl # 2) MCQ eval (download the disjoint config, then build): python training/scripts/hf_download.py --hf-id theelderemo/pentesting-explanations \ --config mitre_attack --split train --out data/download/pentest_mcq_eval/raw.jsonl python training/scripts/build_eval_sets.py --mode mcq \ --mcq-input data/download/pentest_mcq_eval/raw.jsonl \ --out data/eval/knowledge_mcq.jsonl """ from __future__ import annotations import argparse import json import random import sys from pathlib import Path from typing import Any import yaml sys.path.insert(0, str(Path(__file__).resolve().parent)) from sft_adapters import apply_adapter, coerce_list, get_first # noqa: E402 def read_yaml(path: str | Path) -> dict[str, Any]: with Path(path).open("r", encoding="utf-8") as fh: return yaml.safe_load(fh) or {} def read_jsonl(path: Path): with path.open("r", encoding="utf-8") as fh: for line in fh: line = line.strip() if line: yield json.loads(line) def drop_final_assistant(messages: list[dict[str, str]]) -> list[dict[str, str]]: for i in range(len(messages) - 1, -1, -1): if messages[i].get("role") == "assistant": return messages[:i] return messages def build_vuln_detection(manifest: dict[str, Any], raw_dir: Path, scan_cap: int, target_per_class: int, seed: int) -> tuple[list[dict], list[str]]: """Class-balanced vuln-detection set. These test splits are heavily majority `not_vulnerable`, so an unbalanced sample lets a model score high by always answering "not vulnerable" (accuracy high, F1=0). We collect both classes across sources, then take an equal number of each so accuracy and F1 are meaningful. """ pos: list[dict] = [] # vulnerable neg: list[dict] = [] # not_vulnerable notes: list[str] = [] for key, source in manifest.get("sources", {}).items(): if not source.get("eval_split"): continue raw = raw_dir / key / "eval.jsonl" if not raw.is_file(): notes.append(f"{key}: missing {raw} (run hf_download.py --key {key} --eval)") continue scanned = 0 for i, raw_row in enumerate(read_jsonl(raw)): if scanned >= scan_cap: break scanned += 1 for ex in apply_adapter(source["adapter"], raw_row, source.get("params", {}) or {}): if not ex.verify or ex.verify.get("mode") != "label": continue row = { "id": f"{source['hf_id']}:eval:{i}", "source": source["hf_id"], "kind": "vuln_detection", "messages": drop_final_assistant(ex.messages), "gold_label": ex.verify["expected"], "cwe": ex.verify.get("cwe", []), } (pos if ex.verify["expected"] == "vulnerable" else neg).append(row) break rng = random.Random(seed) rng.shuffle(pos) rng.shuffle(neg) n = min(target_per_class, len(pos), len(neg)) notes.append(f"collected vulnerable={len(pos)} not_vulnerable={len(neg)} -> balanced {n}+{n}") if n == 0: notes.append("WARNING: one class is empty; download more rows (hf_download --eval --max-rows N)") rows = pos[:n] + neg[:n] rng.shuffle(rows) return rows, notes def build_mcq(mcq_input: Path, source_label: str, cap: int, seed: int) -> tuple[list[dict], list[str]]: """MCQ set with choices shuffled so the correct answer isn't positionally biased. The raw source lists the correct choice first for most rows (gold ~76% 'A'), which a model can game by always answering 'A'. We permute choices per question and remap the gold index. """ rng = random.Random(seed) rows: list[dict] = [] notes: list[str] = [] for i, raw in enumerate(read_jsonl(mcq_input)): if len(rows) >= cap: break question = get_first(raw, ["question", "Question", "prompt"]) choices = raw.get("choices") or raw.get("options") if not question or not isinstance(choices, list) or not choices: continue gold_index = raw.get("answer_idx", raw.get("correct_idx")) # MMLU stores the answer as an int index in `answer`. if gold_index is None and isinstance(raw.get("answer"), int): gold_index = raw.get("answer") gold_letter = raw.get("correct_letter") or raw.get("answer") if gold_index is None and isinstance(gold_letter, str) and len(gold_letter) == 1: gold_index = ord(gold_letter.upper()) - ord("A") if gold_index is None: correct_choice = get_first(raw, ["correct_choice"]) if correct_choice and correct_choice in choices: gold_index = choices.index(correct_choice) if gold_index is None: continue gold_index = int(gold_index) if not 0 <= gold_index < len(choices): continue # shuffle choices, track where the correct one lands correct_text = str(choices[gold_index]) shuffled = [str(c) for c in choices] rng.shuffle(shuffled) new_gold = shuffled.index(correct_text) rows.append( { "id": f"{source_label}:eval:{i}", "source": source_label, "kind": "mcq", "question": str(question), "choices": shuffled, "gold_index": new_gold, "gold_letter": chr(ord("A") + new_gold), } ) from collections import Counter notes.append("gold_letter dist: " + str(dict(Counter(r["gold_letter"] for r in rows)))) return rows, notes def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument("--mode", choices=["vuln_detection", "mcq"], required=True) parser.add_argument("--manifest", default="training/configs/datasets.yaml") parser.add_argument("--raw-dir", default="data/download") parser.add_argument("--out", required=True) parser.add_argument("--scan-cap", type=int, default=8000, help="Rows to read per source (vuln mode).") parser.add_argument("--target-per-class", type=int, default=60, help="Balanced count per class (vuln mode).") parser.add_argument("--mcq-cap", type=int, default=200, help="Max MCQ rows.") parser.add_argument("--seed", type=int, default=1337) parser.add_argument("--mcq-input", help="Raw MCQ JSONL (for --mode mcq).") parser.add_argument("--mcq-source", default="pentest_mcq_eval") return parser.parse_args() def main() -> int: args = parse_args() out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) notes: list[str] = [] if args.mode == "vuln_detection": manifest = read_yaml(args.manifest) rows, notes = build_vuln_detection(manifest, Path(args.raw_dir), args.scan_cap, args.target_per_class, args.seed) else: if not args.mcq_input: print("--mcq-input is required for --mode mcq", file=sys.stderr) return 2 rows, notes = build_mcq(Path(args.mcq_input), args.mcq_source, args.mcq_cap, args.seed) with out_path.open("w", encoding="utf-8") as out: for row in rows: out.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n") summary = {"mode": args.mode, "rows": len(rows), "out": str(out_path)} if notes: summary["notes"] = notes print(json.dumps(summary, indent=2)) return 0 if __name__ == "__main__": raise SystemExit(main())