code/training/scripts/build_eval_sets.py

#!/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())