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.
43 lines
1.1 KiB
Python
43 lines
1.1 KiB
Python
"""Stage-2 reward tuning (GRPO) on the routing objective.
|
|
|
|
Reward per sampled route: quality - lambda*cost - mu*latency, settled from
|
|
the usage ledger for that (prompt, model) pair. Uses the zenlm grpo trainer;
|
|
this file only adapts the routing task to it.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
import yaml
|
|
|
|
|
|
def main() -> None:
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--config", type=Path, required=True)
|
|
args = ap.parse_args()
|
|
cfg = yaml.safe_load(args.config.read_text())
|
|
g = cfg["grpo"]
|
|
|
|
try:
|
|
from grpo import RoutingRewardTrainer # zenlm/grpo
|
|
except ImportError as e:
|
|
raise SystemExit(
|
|
"zenlm/grpo not installed — `uv pip install -e ../grpo` from the zen workspace"
|
|
) from e
|
|
|
|
trainer = RoutingRewardTrainer(
|
|
model_path=cfg["output_dir"],
|
|
lambda_cost=g["lambda_cost"],
|
|
mu_latency=g["mu_latency"],
|
|
group_size=g["group_size"],
|
|
lr=float(g["lr"]),
|
|
)
|
|
trainer.train("data/routing-sft.jsonl")
|
|
trainer.save(cfg["output_dir"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|