Add the 6M BERT encoder as a switchable backbone alongside the 52M decoder

app.py

@@ -16,14 +16,30 @@ from transformers import AutoTokenizer
 
 from routing_model import MAX_ROUTES, RouterModel, build_text, route
 
-MODEL_ID = os.environ.get("MODEL_ID", "oddadmix/Nawah-Router-v3")
 HF_TOKEN = os.environ.get("MODEL_HF_TOKEN") or os.environ.get("HF_TOKEN")
 
-print(f"[*] loading {MODEL_ID}", flush=True)
-TOK = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
-MODEL = RouterModel.from_pretrained(MODEL_ID, token=HF_TOKEN)
+# Two backbones on the same head and the same task. The 6M is a bidirectional
+# BERT encoder; the 52M is a Llama decoder. Measured on the same eval sets in one
+# session, the encoder wins three of four - see the model cards.
+REGISTRY = {
+    "52M":     {"repo": "oddadmix/Nawah-Router-v3",         "label": "Nawah-Router-v3 (Llama decoder)"},
+    "6M-BERT": {"repo": "oddadmix/Nawah-Router-BERT-6M-v2", "label": "Nawah-Router-BERT-6M-v2 (BERT encoder)"},
+}
+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)))
-print(f"[+] ready — {sum(p.numel() for p in MODEL.parameters())/1e6:.2f}M params", flush=True)
+
+MODEL_ID = REGISTRY[DEFAULT]["repo"]
+TOK, MODEL = LOADED[DEFAULT]["tok"], LOADED[DEFAULT]["model"]
 
 app = FastAPI()
 
@@ -31,6 +47,7 @@ app = FastAPI()
 class RouteReq(BaseModel):
     text: str = ""
     cats: list[str] = []
+    model: str = DEFAULT
 
 
 @app.get("/")
@@ -40,8 +57,10 @@ def index():
 
 @app.get("/api/ready")
 def ready():
-    return {"ready": True, "model": MODEL_ID,
-            "params": sum(p.numel() for p in MODEL.parameters()), "max_routes": MAX_ROUTES}
+    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()]}
 
 
 @app.post("/api/route")
@@ -50,16 +69,18 @@ def api_route(req: RouteReq):
     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(MODEL, TOK, text, cats)
+    res = route(sel["model"], sel["tok"], text, cats)
     full, _ = build_text(text, cats)
-    ntok = len(TOK.encode(full, add_special_tokens=False))
+    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})
+        "ms": round((time.perf_counter() - t0) * 1000), "tokens": ntok,
+        "model": sel["label"], "params": sel["params"]})
 
 
 app.mount("/static", StaticFiles(directory="static"), name="static")

static/app.js

@@ -1,4 +1,8 @@
 /* Nawah-Router demo — editable categories, live routing, connectors weighted by probability. */
+/* Two backbones on the same head: a 6M bidirectional BERT encoder and the 52M
+   Llama decoder. Switching re-routes immediately so the comparison lands on the
+   same text and the same categories. Declared here because run() references it. */
+var currentModel = "6M-BERT";
 const $ = id => document.getElementById(id);
 
 const PRESETS = [
@@ -132,7 +136,7 @@ async function run() {
   try {
     const res = await fetch("/api/route", {
       method: "POST", headers: { "Content-Type": "application/json" },
-      body: JSON.stringify({ text, cats: cats() })
+      body: JSON.stringify({ text, cats: cats(), model: currentModel })
     });
     const data = await res.json();
     if (my !== seq) return;                      // a newer keystroke already superseded this
@@ -169,7 +173,25 @@ PRESETS.forEach((p, i) => {
 
 fetch("/api/ready").then(r => r.json()).then(d => {
   $("dot").classList.add("on");
-  $("modelname").textContent = d.model.split("/").pop();
-  $("modelsize").textContent = (d.params / 1e6).toFixed(1) + "M";
   $("nodemeta").textContent = `CPU · مسار واحد · حتى ${d.max_routes} فئات`;
-}).catch(() => { $("modelname").textContent = "غير متصل"; });
+  currentModel = d.default || currentModel;
+  const pick = $("modelpick");
+  (d.models || []).forEach(m => {
+    const o = document.createElement("option");
+    o.value = m.key;
+    o.textContent = m.key === "6M-BERT" ? "🔵 6M encoder" : "🧱 52M decoder";
+    o.title = m.label;
+    if (m.key === currentModel) o.selected = true;
+    pick.appendChild(o);
+  });
+  const sizeOf = k => {
+    const m = (d.models || []).find(x => x.key === k);
+    return m ? (m.params / 1e6).toFixed(1) + "M" : "";
+  };
+  $("modelsize").textContent = sizeOf(currentModel);
+  pick.addEventListener("change", () => {
+    currentModel = pick.value;
+    $("modelsize").textContent = sizeOf(currentModel);
+    run();                       // same text, same categories, other backbone
+  });
+}).catch(() => { $("modelsize").textContent = "غير متصل"; });

static/index.html

@@ -22,7 +22,8 @@
     </div>
     <div class="modelchip" id="modelchip">
       <span class="dot" id="dot"></span>
-      <span id="modelname">…</span>
+      <select id="modelpick" style="background:transparent;border:0;color:inherit;
+              font:inherit;cursor:pointer;outline:none"></select>
       <span class="mono" id="modelsize"></span>
     </div>
   </header>