File size: 2,111 Bytes
bdb0955 | 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 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | """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) # context: short summary, first 600 characters
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)}) # probability the flag is true
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))
|