| """Gate a trained router on a held-out test set. |
| |
| Usage: |
| python evaluate_router.py --model router-model --test my-sealed-test.jsonl \ |
| --baseline-correct 172 --repeats 3 |
| |
| The test file has the same format as the training data ({"text", "label"}). |
| Write it yourself, BEFORE you train, and never train on it. |
| |
| It prints overall accuracy, a confusion matrix, the gate lines that matter for |
| a router, and latency, then exits non-zero if any predeclared line fails. |
| |
| Latency is local PyTorch inference on whatever device runs this script. Set |
| --max-p95-ms for YOUR target hardware; it is not comparable to a TensorRT |
| engine measured client-side on a Jetson. |
| """ |
| import argparse, json, statistics, sys, time |
| from collections import Counter |
|
|
| import torch |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer |
|
|
| MEDIA = {"image", "audio", "video"} |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--model", required=True) |
| ap.add_argument("--test", required=True) |
| ap.add_argument("--repeats", type=int, default=1) |
| ap.add_argument("--baseline-correct", type=int, help="correct calls by the router you are replacing, same cases") |
| ap.add_argument("--max-chat-to-media", type=int, default=0) |
| ap.add_argument("--max-media-to-chat", type=int, default=0) |
| ap.add_argument("--max-p95-ms", type=float, default=100.0) |
| args = ap.parse_args() |
|
|
| device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" |
| tok = AutoTokenizer.from_pretrained(args.model) |
| tok.truncation_side = "left" |
| model = AutoModelForSequenceClassification.from_pretrained(args.model).to(device).eval() |
| id2label = model.config.id2label |
|
|
| rows = [json.loads(line) for line in open(args.test) if line.strip()] |
| confusion = Counter() |
| latencies = [] |
| with torch.no_grad(): |
| model(**tok("warm up", return_tensors="pt").to(device)) |
| for _ in range(args.repeats): |
| for row in rows: |
| t0 = time.perf_counter() |
| enc = tok(row["text"], truncation=True, max_length=512, return_tensors="pt").to(device) |
| pred = id2label[int(model(**enc).logits.argmax(-1))] |
| if device == "cuda": |
| torch.cuda.synchronize() |
| latencies.append((time.perf_counter() - t0) * 1000) |
| confusion[(row["label"], pred)] += 1 |
|
|
| calls = sum(confusion.values()) |
| correct = sum(n for (gold, pred), n in confusion.items() if gold == pred) |
| chat_to_media = sum(n for (gold, pred), n in confusion.items() if gold == "chat" and pred in MEDIA) |
| media_to_chat = sum(n for (gold, pred), n in confusion.items() if gold in MEDIA | {"search"} and pred == "chat") |
| lat = sorted(latencies) |
| p50 = statistics.median(lat) |
| p95 = lat[min(len(lat) - 1, int(0.95 * len(lat)))] |
|
|
| labels = sorted({g for g, _ in confusion} | {p for _, p in confusion}) |
| print(f"accuracy: {correct}/{calls}") |
| print("confusion (rows = gold, cols = predicted):") |
| print(" " + " ".join(f"{l[:6]:>6}" for l in labels)) |
| for g in labels: |
| print(f"{g[:6]:>6} " + " ".join(f"{confusion[(g, p)]:>6}" for p in labels)) |
| print(f"latency ms: p50 {p50:.1f} p95 {p95:.1f} ({device})") |
|
|
| gates = [ |
| ("chat sent to a media lane", chat_to_media <= args.max_chat_to_media, f"{chat_to_media} (max {args.max_chat_to_media})"), |
| ("specialist request sent to chat", media_to_chat <= args.max_media_to_chat, f"{media_to_chat} (max {args.max_media_to_chat})"), |
| ("p95 latency", p95 <= args.max_p95_ms, f"{p95:.1f} ms (max {args.max_p95_ms})"), |
| ] |
| if args.baseline_correct is not None: |
| gates.append(("at least as accurate as the baseline", correct >= args.baseline_correct, |
| f"{correct} vs {args.baseline_correct}")) |
| failed = False |
| for name, ok, detail in gates: |
| print(f"{'PASS' if ok else 'FAIL'} {name}: {detail}") |
| failed |= not ok |
| sys.exit(1 if failed else 0) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
