Spaces:
Running
Running
| """ | |
| Nawah-Router — FastAPI backend for the interactive routing demo. | |
| One endpoint. The model scores every supplied category in a single forward pass, so the frontend | |
| can re-route on each keystroke without a per-category cost. | |
| """ | |
| import os | |
| import time | |
| import torch | |
| from fastapi import FastAPI | |
| from fastapi.responses import FileResponse, JSONResponse | |
| from fastapi.staticfiles import StaticFiles | |
| from pydantic import BaseModel | |
| from transformers import AutoTokenizer | |
| from routing_model import MAX_ROUTES, RouterModel, build_text, detect_lang, route | |
| HF_TOKEN = os.environ.get("MODEL_HF_TOKEN") or os.environ.get("HF_TOKEN") | |
| # Three backbones on the same head and the same task. The two 6M's are bidirectional BERT | |
| # encoders (one Arabic-only, one pretrained jointly on Arabic+English); the 52M is a Llama | |
| # decoder. Measured on the same eval sets in one session each - see the model cards. | |
| REGISTRY = { | |
| "52M": {"repo": "oddadmix/Nawah-Router-v3", | |
| "label": "Nawah-Router-v3 (Llama decoder, Arabic)"}, | |
| "6M-BERT": {"repo": "oddadmix/Nawah-Router-BERT-6M-v2", | |
| "label": "Nawah-Router-BERT-6M-v2 (BERT encoder, Arabic)"}, | |
| "6M-BILINGUAL": {"repo": "oddadmix/Nawah-Router-BERT-6M-bilingual-pretrained", | |
| "label": "Nawah-Router-BERT-6M-bilingual-pretrained (BERT encoder, English+Arabic)"}, | |
| } | |
| DEFAULT = os.environ.get("MODEL_KEY", "6M-BERT") | |
| LOADED = {} | |
| for key, spec in REGISTRY.items(): | |
| print(f"[*] loading {spec['repo']}", flush=True) | |
| tok = AutoTokenizer.from_pretrained(spec["repo"], token=HF_TOKEN) | |
| mdl = RouterModel.from_pretrained(spec["repo"], token=HF_TOKEN) | |
| LOADED[key] = {"tok": tok, "model": mdl, "label": spec["label"], | |
| "repo": spec["repo"], | |
| "params": sum(p.numel() for p in mdl.parameters())} | |
| print(f"[+] {key}: {LOADED[key]['params']/1e6:.2f}M params", flush=True) | |
| torch.set_num_threads(int(os.environ.get("OMP_NUM_THREADS", 4))) | |
| MODEL_ID = REGISTRY[DEFAULT]["repo"] | |
| TOK, MODEL = LOADED[DEFAULT]["tok"], LOADED[DEFAULT]["model"] | |
| app = FastAPI() | |
| class RouteReq(BaseModel): | |
| text: str = "" | |
| cats: list[str] = [] | |
| model: str = DEFAULT | |
| def index(): | |
| return FileResponse("static/index.html") | |
| def ready(): | |
| return {"ready": True, "model": MODEL_ID, "default": DEFAULT, | |
| "params": LOADED[DEFAULT]["params"], "max_routes": MAX_ROUTES, | |
| "models": [{"key": k, "label": v["label"], "repo": v["repo"], | |
| "params": v["params"]} for k, v in LOADED.items()]} | |
| def api_route(req: RouteReq): | |
| cats = [c.strip() for c in req.cats if c and c.strip()][:MAX_ROUTES] | |
| text = (req.text or "").strip() | |
| if not text or not cats: | |
| return JSONResponse({"results": [], "ms": 0, "tokens": 0}) | |
| sel = LOADED.get(req.model) or LOADED[DEFAULT] | |
| t0 = time.perf_counter() | |
| res = route(sel["model"], sel["tok"], text, cats) | |
| full, _ = build_text(text, cats) | |
| ntok = len(sel["tok"].encode(full, add_special_tokens=False)) | |
| order = {r["route"]: r["score"] for r in res} | |
| return JSONResponse({ | |
| # returned in the caller's order so the UI does not reshuffle rows under the cursor | |
| "results": [{"route": c, "score": order.get(c, 0.0)} for c in cats], | |
| "top": max(range(len(cats)), key=lambda i: order.get(cats[i], 0.0)), | |
| "ms": round((time.perf_counter() - t0) * 1000), "tokens": ntok, "lang": detect_lang(text), | |
| "model": sel["label"], "params": sel["params"]}) | |
| app.mount("/static", StaticFiles(directory="static"), name="static") | |