mirror of
https://github.com/zenlm/zen-router.git
synced 2026-10-30 14:57:44 +00:00
zeekay
f888212afb
Reads a prompt, picks the best model. 0.6B encoder-pooled router with task/route/feature heads; SFT from eval profiles + usage ledger, GRPO on quality - lambda*cost - mu*latency; serves via hanzo-engine and plugs the hanzo-router Classifier and enso Featurizer seams. Apache-2.0.
25 lines
890 B
Python
25 lines
890 B
Python
"""Route a prompt with a trained zen-router checkpoint."""
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from transformers import AutoTokenizer
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
from training.sft import ZenRouter
|
|
|
|
model_dir = Path("out/zen-router")
|
|
cfg = json.loads((model_dir / "router_config.json").read_text())
|
|
tok = AutoTokenizer.from_pretrained(model_dir)
|
|
model = ZenRouter(str(model_dir), len(cfg["tasks"]), len(cfg["catalog"]), 256)
|
|
model.load_state_dict(torch.load(model_dir / "zen-router.pt", map_location="cpu"))
|
|
model.eval()
|
|
|
|
prompt = sys.argv[1] if len(sys.argv) > 1 else "refactor this rust function to be async"
|
|
enc = tok(prompt, return_tensors="pt")
|
|
with torch.no_grad():
|
|
t_logits, r_logits, feats = model(**enc)
|
|
print("task :", cfg["tasks"][t_logits.argmax().item()])
|
|
print("route:", cfg["catalog"][r_logits.argmax().item()])
|