Upload app.py

app.py

@@ -76,7 +76,11 @@ if os.getenv("LAYA_HIDE_CUDA", "0") in ("1", "true", "True"):
 SERVING_ELSEWHERE = os.environ.get("LAYA_SERVING") == "1"
 BACKEND = os.getenv("LAYA_BACKEND", "").strip()
 MODEL = os.getenv("LAYA_MODEL", "").strip()
-DTYPE = os.getenv("LAYA_DTYPE", "float16").strip()
+# Empty by default: the right dtype depends on the backend — see agent().
+# float16 on a CPU produced NaNs here, which surface as a UNIFORM probability
+# distribution and confidence 0.0 rather than as an error (seen 22 Sep 2026 on
+# the Space: every choice 0.3333).
+DTYPE = os.getenv("LAYA_DTYPE", "").strip()
 TOKEN = os.getenv("LAYA_TOKEN", "").strip()
 ORIGINS = [o.strip() for o in os.getenv("LAYA_CORS_ORIGINS", "").split(",") if o.strip()]
 EAGER = os.getenv("LAYA_EAGER", "1") not in ("0", "false", "False") and not SERVING_ELSEWHERE
@@ -111,13 +115,15 @@ def agent():
         name = _pick_backend()
         mod = importlib.import_module(name)
         model = MODEL or DEFAULT_MODELS[name]
+        # MLX runs fp16 on the GPU happily; torch on a CPU needs fp32.
+        dtype = DTYPE or ("float16" if name == "laya_mlx" else "float32")
         # Both packages expose load(model_id, dtype=...); backends that do not
         # take dtype raise TypeError, so fall back to the plain call.
         try:
-            a = mod.load(model, dtype=DTYPE)
+            a = mod.load(model, dtype=dtype)
         except TypeError:
             a = mod.load(model)
-        _state.update(agent=a, backend=name, model=model, load_ms=round((time.time() - t0) * 1000))
+        _state.update(agent=a, backend=name, model=model, dtype=dtype, load_ms=round((time.time() - t0) * 1000))
         return a
     except Exception as e:  # noqa: BLE001
         _state["error"] = f"{type(e).__name__}: {e}"
@@ -144,7 +150,11 @@ def _auth(token: Optional[str]) -> None:
 class Question(BaseModel):
     type: Literal["choice", "score", "noul"]
     instructions: str
-    criteria: Optional[List[str]] = None       # required for choice and score
+    # Laya scores each option at its own [MASK] token, so `criteria` is a MAP
+    # of option -> what that option means. A bare list is accepted and each
+    # entry becomes its own description; passing a list to the model itself
+    # yields a uniform distribution with confidence 0.0 rather than an error.
+    criteria: Optional[Union[Dict[str, str], List[str]]] = None
 
 
 class DecideIn(BaseModel):
@@ -182,20 +192,35 @@ def _decide(state: Any, questions: Dict[str, Dict[str, Any]]) -> Dict[str, Any]:
 
 
 def _route(question: str, candidates: List[Candidate], instructions: str, top_k: int) -> Dict[str, Any]:
-    labels = [c.label if not c.hint else f"{c.label} — {c.hint}" for c in candidates]
-    out = _decide(question, {"pick": {"type": "choice", "instructions": instructions, "criteria": labels}})
+    """One choice question whose options are the candidates.
+
+    `criteria` maps each option to what it means — that description is what
+    Laya scores against the state, so a candidate's `hint` is worth giving.
+    Duplicate labels would collide as dict keys, so they are suffixed.
+    """
+    options: Dict[str, Candidate] = {}
+    for c in candidates:
+        key = c.label
+        if key in options:
+            key = f"{c.label} ({c.id})"
+        options[key] = c
+    criteria = {k: (c.hint or c.label) for k, c in options.items()}
+
+    out = _decide(question, {"pick": {"type": "choice", "instructions": instructions, "criteria": criteria}})
     ans = out["answers"].get("pick", {}) if isinstance(out["answers"], dict) else {}
     probs = ans.get("probabilities") or ans.get("probs") or {}
     if isinstance(probs, list):                      # some builds return a bare vector
-        probs = {labels[i]: p for i, p in enumerate(probs) if i < len(labels)}
+        probs = {k: probs[i] for i, k in enumerate(options) if i < len(probs)}
     ranked = sorted(
-        ({"id": c.id, "label": c.label, "p": float(probs.get(labels[i], 0.0))} for i, c in enumerate(candidates)),
+        ({"id": c.id, "label": c.label, "p": float(probs.get(k, 0.0))} for k, c in options.items()),
         key=lambda r: r["p"], reverse=True,
     )[: max(1, top_k)]
-    chosen = ans.get("value") or ans.get("label")
-    top = next((r for r in ranked if str(r["label"]) == str(chosen)), ranked[0] if ranked else None)
+    chosen = ans.get("choice") or ans.get("value") or ans.get("label")
+    picked = options.get(str(chosen))
+    top = (next((r for r in ranked if r["id"] == picked.id), None) if picked else None) or (ranked[0] if ranked else None)
     return {"id": top["id"] if top else None, "label": top["label"] if top else None,
-            "p": top["p"] if top else None, "ranked": ranked, "ms": out["ms"]}
+            "p": top["p"] if top else None, "confidence": ans.get("confidence"),
+            "ranked": ranked, "ms": out["ms"]}
 
 
 # ── REST routes (attached to whichever server ends up running) ──────────────
@@ -208,6 +233,7 @@ def health() -> Dict[str, Any]:
         "ok": _state["agent"] is not None,
         "backend": _state["backend"],
         "model": _state["model"] or MODEL or "(default)",
+        "dtype": _state.get("dtype"),
         "load_ms": _state["load_ms"],
         "error": _state["error"],
         "auth_required": bool(TOKEN),
@@ -217,7 +243,12 @@ def health() -> Dict[str, Any]:
 @api.post("/v1/decide")
 def decide(body: DecideIn, x_laya_token: Optional[str] = Header(None)) -> Dict[str, Any]:
     _auth(x_laya_token)
-    qs = {k: {kk: vv for kk, vv in v.model_dump().items() if vv is not None} for k, v in body.questions.items()}
+    qs = {}
+    for k, v in body.questions.items():
+        q = {kk: vv for kk, vv in v.model_dump().items() if vv is not None}
+        if isinstance(q.get("criteria"), list):
+            q["criteria"] = {c: c for c in q["criteria"]}
+        qs[k] = q
     return _decide(body.state, qs)