File size: 5,514 Bytes
994182c | 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 | #!/usr/bin/env python3
"""In-memory held-out evaluation used by run_sft.py during training.
On a single GPU you cannot serve vLLM while training holds VRAM, so the per-epoch
benchmark runs on the *training* model directly (model.generate) on a small sample
of the same held-out sets eval_endpoint.py uses. Scoring/parse logic is imported
from eval_endpoint so in-training numbers are comparable to the base baseline.
Robustness: callers should wrap this in try/except — a benchmark hiccup must never
kill a training run.
"""
from __future__ import annotations
import json
import sys
from pathlib import Path
from typing import Any
sys.path.insert(0, str(Path(__file__).resolve().parent))
from eval_endpoint import MCQ_INSTRUCTION, parse_letter, parse_verdict # noqa: E402
def _read_jsonl(path: Path, limit: int | None) -> list[dict[str, Any]]:
rows: list[dict[str, Any]] = []
with path.open("r", encoding="utf-8") as fh:
for line in fh:
line = line.strip()
if line:
rows.append(json.loads(line))
if limit and len(rows) >= limit:
break
return rows
def _generate(model, tokenizer, messages: list[dict[str, str]], max_new_tokens: int,
enable_thinking: bool = False) -> str:
import torch
try:
prompt = tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True, enable_thinking=enable_thinking
)
except TypeError:
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
out = model.generate(
**inputs,
max_new_tokens=max_new_tokens,
do_sample=False,
pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
)
return tokenizer.decode(out[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=False)
def _score_vuln(model, tokenizer, rows, max_new_tokens, enable_thinking=False) -> dict[str, Any]:
tp = fp = tn = fn = unparsed = 0
for r in rows:
pred = parse_verdict(_generate(model, tokenizer, r["messages"], max_new_tokens, enable_thinking))
gold = r["gold_label"]
if pred is None:
unparsed += 1
continue
if gold == "vulnerable":
tp += pred == "vulnerable"; fn += pred != "vulnerable"
else:
tn += pred == "not_vulnerable"; fp += pred != "not_vulnerable"
n = len(rows)
prec = tp / (tp + fp) if (tp + fp) else 0.0
rec = tp / (tp + fn) if (tp + fn) else 0.0
return {
"kind": "vuln_detection", "n": n,
"accuracy": (tp + tn) / n if n else 0.0,
"precision_vuln": prec, "recall_vuln": rec,
"f1_vuln": 2 * prec * rec / (prec + rec) if (prec + rec) else 0.0,
"unparsed": unparsed,
}
def _score_mcq(model, tokenizer, rows, max_new_tokens, enable_thinking=False) -> dict[str, Any]:
correct = unparsed = 0
for r in rows:
choices = r["choices"]
user = r["question"] + "\n\n" + "\n".join(
f"{chr(ord('A') + j)}. {c}" for j, c in enumerate(choices)
) + "\n\n" + MCQ_INSTRUCTION
messages = [
{"role": "system", "content": "You are a cybersecurity expert. Authorized security research context."},
{"role": "user", "content": user},
]
pred = parse_letter(_generate(model, tokenizer, messages, max_new_tokens, enable_thinking), len(choices))
if pred is None:
unparsed += 1
continue
correct += int(pred == r["gold_index"])
n = len(rows)
return {"kind": "mcq", "n": n, "accuracy": correct / n if n else 0.0, "unparsed": unparsed}
def run_eval_sets(model, tokenizer, eval_files: list[str], sample_per_set: int = 80,
max_new_tokens: int = 256, enable_thinking: bool = False) -> dict[str, Any]:
"""Return {set_name: metrics} for each eval file (sampled for speed).
enable_thinking defaults to False: this tracking eval measures direct-answer
accuracy so generations stay short, fast and parseable. The base baseline and
every per-epoch run use the same setting, so the numbers stay comparable. (The
primary CyberGym metric exercises the full agentic thinking flow separately.)
"""
was_training = model.training
model.eval()
# generate() needs the kv-cache, which gradient checkpointing disables
prev_use_cache = getattr(model.config, "use_cache", None)
try:
model.config.use_cache = True
except Exception:
pass
results: dict[str, Any] = {}
try:
for f in eval_files:
path = Path(f)
if not path.is_file():
results[path.stem] = {"error": "missing"}
continue
rows = _read_jsonl(path, sample_per_set)
if not rows:
results[path.stem] = {"error": "empty"}
continue
kind = rows[0].get("kind", "vuln_detection")
scorer = _score_mcq if kind == "mcq" else _score_vuln
results[path.stem] = scorer(model, tokenizer, rows, max_new_tokens, enable_thinking)
finally:
if prev_use_cache is not None:
try:
model.config.use_cache = prev_use_cache
except Exception:
pass
if was_training:
model.train()
return results
|