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