LiquidAI/prompt-routing

1"""Prompt routing — server-side CPU inference.
2
3The LFM2.5 Encoder router runs on the Space's CPU. Weights are pulled from the Hub
4model repo at boot (no local copy). The browser sends {text, cats} to /api/route
5and gets back the routing distribution; nothing is downloaded to the visitor and no
6WebGPU is required.
7"""
8
9import os
10import threading
11import time
12
13import torch
14from flask import Flask, request, send_from_directory
15
16import router as R
17
18APP_DIR = os.path.dirname(os.path.abspath(__file__))
19_CPUS = R.effective_cpus()
20torch.set_num_threads(_CPUS)
21try:
22    torch.set_num_interop_threads(1)
23except RuntimeError:
24    pass
25
26app = Flask(__name__)
27
28HF_TOKEN = os.environ.get("HF_TOKEN")
29print(f"[boot] pulling {R.MODEL_ID} from the Hub, loading on {_CPUS} CPU thread(s), f32…", flush=True)
30from tokenizers import Tokenizer
31from huggingface_hub import hf_hub_download
32TOK = Tokenizer.from_file(hf_hub_download(R.MODEL_ID, "tokenizer.json", token=HF_TOKEN))
33MODEL = R.Lfm2Router.from_hub(token=HF_TOKEN, quantize=False)
34_LOCK = threading.Lock()
35# warm up so the first real request isn't the slow one
36R.route(MODEL, TOK, "Set a timer.", ["Simple tool use", "Creative writing"])
37print("[boot] router ready", flush=True)
38
39
40@app.route("/")
41def index():
42    return send_from_directory(APP_DIR, "index.html")
43
44
45@app.route("/<path:name>")
46def static_file(name):
47    if name in ("style.css", "shared-ui.js", "lliquid.gif"):
48        return send_from_directory(APP_DIR, name)
49    return ("not found", 404)
50
51
52@app.route("/api/ready")
53def ready():
54    return {"ready": MODEL is not None}
55
56
57@app.route("/api/route", methods=["POST"])
58def api_route():
59    data = request.get_json(silent=True) or {}
60    text = (data.get("text") or "").strip()
61    cats = [c.strip() for c in (data.get("cats") or []) if c and c.strip()]
62    if not text or not cats:
63        return {"probs": [], "top": -1, "tokens": 0, "ms": 0}
64    t0 = time.perf_counter()
65    with _LOCK:
66        probs, top, tokens = R.route(MODEL, TOK, text, cats)
67    return {"probs": probs, "top": top, "tokens": tokens,
68            "ms": round((time.perf_counter() - t0) * 1000)}
69
70
71if __name__ == "__main__":
72    app.run(host="0.0.0.0", port=int(os.environ.get("PORT", 7860)), threaded=True)
73