File size: 1,342 Bytes
edccc2c | 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 | """© KAND CA 2026 - evaluate a router checkpoint without retraining.
Reuses train_router_head.py's own loaders and report() so a published model and
a new one are scored by identical code in one session.
A checkpoint saved by train_router_head.py is a directory with the backbone plus
router_model.pt (the span scorer). The published Nawah-Router-v3 has the same
layout, so both load through the same path.
"""
import os, sys
os.environ.setdefault("DATA_DIR", "ds/data")
import torch
from transformers import AutoTokenizer
import train_router_head as T
def load_model(mid):
m = T.RouterModel(mid)
w = os.path.join(mid, "router_model.pt") if os.path.isdir(mid) else None
if w is None:
from huggingface_hub import hf_hub_download
w = hf_hub_download(mid, "router_model.pt")
if os.path.exists(w):
m.load_state_dict(torch.load(w, map_location="cpu", weights_only=True))
return m.cuda().eval()
for mid in sys.argv[1:]:
tok = AutoTokenizer.from_pretrained(mid)
m = load_model(mid)
n = sum(p.numel() for p in m.parameters())
print(f"\n{'='*66}\n{mid} params={n:,}")
for name in ("eval_unseen_lanes", "eval_unseen_domain", "eval_unseen_axis", "eval_hard"):
rows = T.load(name)
T.report(m, tok, rows, name.replace("eval_", ""))
del m; torch.cuda.empty_cache()
|