eval_testset.py

#!/usr/bin/env python3
"""Evaluate the tuned model on the held-out 200-utterance test set → metrics for §7.

Runs the merged model locally (works on CUDA or Apple mps — the Mac Studio is fine for inference).
Prints a confusion matrix, per-class precision/recall/F1, accuracy, and the false-positive rate
(benign flagged as malevolent) — the key indicator that legitimate members are not blocked.

  python eval_testset.py --model ./bcbsma-injection-gemma2-9b-merged
"""
import argparse, json, torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from data_utils import to_messages   # shared with train_lora.py (single source of truth)

TEST_FILE = "bcbsma_injection_ft_test.jsonl"   # same file train_lora.py holds out; either format works

def parse_label(text):
    t = text.strip().lower()
    return "malevolent" if "malevolent" in t else ("benign" if "benign" in t else "unknown")

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="./bcbsma-injection-gemma2-9b-merged")
    ap.add_argument("--test", default=TEST_FILE)
    args = ap.parse_args()

    tok = AutoTokenizer.from_pretrained(args.model)
    model = AutoModelForCausalLM.from_pretrained(
        args.model, torch_dtype=torch.bfloat16, device_map="auto", attn_implementation="eager",
    ).eval()

    rows = [json.loads(l) for l in open(args.test, encoding="utf-8")]
    # positive class = malevolent
    tp = fp = tn = fn = unknown = 0
    for r in rows:
        msgs = to_messages(r)                       # dual-format (chat or Vertex) -> folded prompt
        gold = msgs[-1]["content"].strip().lower()  # the assistant label
        prompt = tok.apply_chat_template(msgs[:-1], # user turn only; the model generates the label
                                         tokenize=False, add_generation_prompt=True)
        ids = tok(prompt, return_tensors="pt").to(model.device)
        with torch.no_grad():
            out = model.generate(**ids, max_new_tokens=4, do_sample=False)
        pred = parse_label(tok.decode(out[0][ids["input_ids"].shape[1]:], skip_special_tokens=True))
        if pred == "unknown":
            unknown += 1; pred = "benign"  # count a non-answer as a miss on the safe side
        if gold == "malevolent" and pred == "malevolent": tp += 1
        elif gold == "benign" and pred == "malevolent": fp += 1
        elif gold == "benign" and pred == "benign": tn += 1
        elif gold == "malevolent" and pred == "benign": fn += 1

    n = len(rows)
    acc = (tp + tn) / n
    def prf(tp, fp, fn):
        p = tp / (tp + fp) if (tp + fp) else 0.0
        r = tp / (tp + fn) if (tp + fn) else 0.0
        f = 2 * p * r / (p + r) if (p + r) else 0.0
        return p, r, f
    mp, mr, mf = prf(tp, fp, fn)      # malevolent
    bp, br, bf = prf(tn, fn, fp)      # benign (swap roles)
    fpr = fp / (fp + tn) if (fp + tn) else 0.0

    print(f"\nModel: {args.model}   |   test examples: {n}   |   unanswered: {unknown}\n")
    print("Confusion matrix (rows = actual, cols = predicted)")
    print(f"                 pred malevolent   pred benign")
    print(f"  malevolent          {tp:>4}            {fn:>4}")
    print(f"  benign              {fp:>4}            {tn:>4}\n")
    print(f"  Accuracy:                 {acc:6.3f}")
    print(f"  Malevolent  P/R/F1:       {mp:.3f} / {mr:.3f} / {mf:.3f}")
    print(f"  Benign      P/R/F1:       {bp:.3f} / {br:.3f} / {bf:.3f}")
    print(f"  False-positive rate:      {fpr:6.3f}   (benign flagged as malevolent — blocks legit members)")

if __name__ == "__main__":
    main()