| """Run s1-llm-auto-router locally on CPU. Needs: pip install onnxruntime tokenizers numpy huggingface_hub""" |
| import json |
| import sys |
|
|
| import numpy as np |
| import onnxruntime as ort |
| from huggingface_hub import snapshot_download |
| from tokenizers import Tokenizer |
|
|
| CATEGORIES = ["coding", "agentic", "math", "knowledge", "long_context", "tool_use", "design", "summarisation", "general"] |
| DIFFICULTY = ["trivial", "easy", "moderate", "hard", "frontier"] |
| STAKES = ["negligible", "low", "medium", "high"] |
| FLAGS = ["needs_tools", "needs_vision", "needs_long_context", "follow_up"] |
|
|
| root = snapshot_download("decisionmodels/s1-llm-auto-router") |
| temps = json.load(open(f"{root}/meta.json"))["temperatures"] |
| tok = Tokenizer.from_file(f"{root}/tokenizer/tokenizer.json") |
| tok.enable_truncation(512, strategy="only_second") |
| sess = ort.InferenceSession(f"{root}/model.gq.onnx", providers=["CPUExecutionProvider"]) |
|
|
|
|
| def softmax(x): |
| e = np.exp(x - x.max()) |
| return e / e.sum() |
|
|
|
|
| def classify(request: str, context: str = "(new conversation)") -> dict: |
| enc = tok.encode(context[:600], request) |
| ids = np.asarray([enc.ids], dtype=np.int64) |
| cat, dif, stk, flg = sess.run(None, {"input_ids": ids, "attention_mask": np.ones_like(ids)}) |
| pc = softmax(cat[0] / temps["category"]) |
| pd = softmax(dif[0] / temps["difficulty"]) |
| ps = softmax(stk[0] / temps["stakes"]) |
| pf = 1 / (1 + np.exp(-flg[0] / temps["flags"])) |
| out = { |
| "category": CATEGORIES[int(pc.argmax())], |
| "category_confidence": round(float(pc.max()), 3), |
| "difficulty": DIFFICULTY[int(pd.argmax())], |
| "difficulty_expected": round(float((np.arange(5) * pd).sum()), 2), |
| "stakes": STAKES[int(ps.argmax())], |
| "stakes_expected": round(float((np.arange(4) * ps).sum()), 2), |
| } |
| out.update({name: round(float(p), 3) for name, p in zip(FLAGS, pf)}) |
| return out |
|
|
|
|
| if __name__ == "__main__": |
| text = " ".join(sys.argv[1:]) or "Refactor this ABAP report to use CDS views" |
| print(json.dumps(classify(text), indent=2)) |
|
|
