| |
| """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 |
|
|
|
|
| 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] = [] |
| neg: list[dict] = [] |
| 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")) |
| |
| 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 |
| |
| 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()) |
|
|
