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