| """Score a Vela 2.0 model (or the regex scanner) on the eval set. |
| |
| python scripts/run_baseline.py --model vllm-sr/Vela-2.0-0.8B |
| python scripts/run_baseline.py --model regex # secrets task only, no GPU |
| python scripts/run_baseline.py --model vllm-sr/Vela-2.0-0.8B --heads runs/vela-0.8b-guard |
| |
| Writes results/<slug>/predictions.jsonl and results/<slug>/metrics.json. |
| """ |
| import argparse |
| import json |
| import logging |
| import time |
| from pathlib import Path |
|
|
| from vela_guard import evaluate, regex_baseline |
|
|
| TASKS = ("secrets", "injection", "toolcall") |
|
|
| log = logging.getLogger("run_baseline") |
|
|
|
|
| def load_rows(data_dir: Path, tasks: tuple[str, ...]) -> dict[str, list[dict]]: |
| rows = {} |
| for task in tasks: |
| path = data_dir / f"{task}.jsonl" |
| if not path.exists(): |
| raise FileNotFoundError(f"{path} missing; run scripts/build_eval.py first") |
| rows[task] = [json.loads(line) for line in path.read_text().splitlines() if line.strip()] |
| return rows |
|
|
|
|
| def load_vela(name: str, device: str, heads_dir: Path | None): |
| import torch |
| from transformers import AutoModel |
|
|
| if device == "cuda" and not torch.cuda.is_available(): |
| raise RuntimeError("CUDA not available; this script is meant to run on a RunPod GPU") |
| if heads_dir is not None: |
| from trainer.guard_model import GuardModel |
|
|
| return GuardModel.load(name, heads_dir, device) |
| model = AutoModel.from_pretrained(name, trust_remote_code=True) |
| return model.to(device).eval() |
|
|
|
|
| def predict_vela(model, row: dict) -> tuple[dict, float]: |
| state, questions = evaluate.build_request(row) |
| start = time.perf_counter() |
| response = model.system_one(state, questions) |
| return evaluate.parse_response(row["task"], response), time.perf_counter() - start |
|
|
|
|
| def predict_regex(row: dict) -> tuple[dict, float]: |
| start = time.perf_counter() |
| spans = [{k: s[k] for k in ("label", "start", "end")} for s in regex_baseline.find_secrets(row["state"]["request"])] |
| return {"spans": spans, "score": 1.0 if spans else 0.0}, time.perf_counter() - start |
|
|
|
|
| def run(args: argparse.Namespace) -> None: |
| is_regex = args.model == "regex" |
| tasks = ("secrets",) if is_regex else tuple(args.tasks) |
| rows_by_task = load_rows(args.data, tasks) |
| model = None if is_regex else load_vela(args.model, args.device, args.heads) |
|
|
| out_dir = args.out / (args.model.replace("/", "__") + ("__guard" if args.heads else "")) |
| out_dir.mkdir(parents=True, exist_ok=True) |
| report, failures = {"model": args.model, "tasks": {}}, 0 |
| with (out_dir / "predictions.jsonl").open("w") as fh: |
| for task, rows in rows_by_task.items(): |
| kept_rows, preds, latencies = [], [], [] |
| for row in rows: |
| try: |
| pred, seconds = predict_regex(row) if is_regex else predict_vela(model, row) |
| except Exception: |
| log.exception("prediction failed for %s", row["id"]) |
| failures += 1 |
| continue |
| kept_rows.append(row) |
| preds.append(pred) |
| latencies.append(seconds) |
| fh.write(json.dumps({"id": row["id"], "pred": pred, "seconds": seconds}) + "\n") |
| task_report = evaluate.score_task(task, kept_rows, preds) |
| task_report["median_ms"] = 1000 * sorted(latencies)[len(latencies) // 2] if latencies else None |
| report["tasks"][task] = task_report |
| log.info("%s done: %d/%d rows", task, len(kept_rows), len(rows)) |
| report["failures"] = failures |
| (out_dir / "metrics.json").write_text(json.dumps(report, indent=2)) |
| log.info("%s", headline(report)) |
|
|
|
|
| def headline(report: dict) -> str: |
| t = report["tasks"] |
| parts = [report["model"]] |
| if "secrets" in t: |
| parts.append(f"secrets exactF1={t['secrets']['all']['exact']['f1']:.3f} overlapF1={t['secrets']['all']['overlap']['f1']:.3f}") |
| if "injection" in t: |
| parts.append(f"injection acc={t['injection']['all']['accuracy']:.3f} auroc={t['injection']['all']['auroc']:.3f}") |
| if "toolcall" in t: |
| parts.append(f"toolcall acc={t['toolcall']['all']['accuracy']:.3f} block->allow={t['toolcall']['all']['block_passed_as_allow']:.3f}") |
| return " | ".join(parts) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| parser.add_argument("--model", required=True, help="HF repo id, or 'regex'") |
| parser.add_argument("--data", type=Path, default=Path("data/eval")) |
| parser.add_argument("--out", type=Path, default=Path("results")) |
| parser.add_argument("--device", default="cuda") |
| parser.add_argument("--tasks", nargs="+", choices=TASKS, default=list(TASKS)) |
| parser.add_argument("--heads", type=Path, default=None, help="trained guard heads dir (trainer.train_heads --out)") |
| run(parser.parse_args()) |
|
|
|
|
| if __name__ == "__main__": |
| logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s") |
| main() |
|
|
