"""Model backends for DecisionLab: FalconDec and Arthur models found in the models folder, and Laya. The model list and its settings live in app/registry.py. Every backend takes the same Jev/Laya-style request: state: str or dict questions: {name: {"type": "choice"|"score"|"noul", "instructions": str, "criteria": {key: description} (choice) | [levels] (score)}} and return the same normalised answer per question, so the UI can compare them directly. """ from __future__ import annotations import importlib.util import os import threading import time import traceback from pathlib import Path import torch from huggingface_hub import snapshot_download from .registry import load_order, model_specs, resolve_local_dir, trusted_modeling, warmup_enabled from .security import Gate from .scoring import _lookup, normalise SPECS = model_specs(os.environ) DEVICE_PREF = os.getenv("DEVICE", "auto") # auto | cuda | cpu def gpu_label() -> str | None: """The GPU's name for the status panel, or None on CPU. On ZeroGPU there is no real GPU outside @spaces.GPU, so asking for its name can fail; the lab then reports how the GPU is provided instead of crashing /api/status.""" try: return torch.cuda.get_device_name(0) if torch.cuda.is_available() else None except Exception: return "attached per decision (ZeroGPU)" def pick_device() -> str: if DEVICE_PREF == "cpu": return "cpu" if torch.cuda.is_available(): return "cuda" if DEVICE_PREF == "cuda": print("[DecisionLab] DEVICE=cuda requested but CUDA is unavailable; using CPU") return "cpu" def _sync(device: str) -> None: if device == "cuda": torch.cuda.synchronize() # ----------------------------------------------------------------------------- base class class Backend: def __init__(self, spec: dict): self.spec = spec self.key, self.name, self.side = spec["key"], spec["name"], spec["side"] self.status = "idle" # idle | loading | ready | error self.error = "" self.info: dict = {} self.device = pick_device() self._lock = threading.Lock() self._loading = Gate(1) # DL-SA-005: one load at a time per model def describe(self) -> dict: # Where the model comes from is known from discovery, before (and whether or not) it loads. where = {k: self.spec[k] for k in ("path", "repo") if self.spec.get(k)} return dict(key=self.key, name=self.name, side=self.side, status=self.status, error=self.error, device=self.device, **{**where, **self.info}) def start_load(self) -> bool: """Start loading in a background thread unless a load is already running. Returns False if one is.""" if not self._loading.try_enter(): return False self.status, self.error = "loading", "" threading.Thread(target=self.load, kwargs={"gated": True}, daemon=True).start() return True def load(self, gated: bool = False) -> None: if not gated and not self._loading.try_enter(): return try: self._load_once() finally: self._loading.leave() def _load_once(self) -> None: self.status, self.error = "loading", "" t0 = time.perf_counter() try: self._load() self.info["load_seconds"] = round(time.perf_counter() - t0, 1) if warmup_enabled(os.environ): self._warmup() self.status = "ready" except Exception as exc: # surfaced in the UI traceback.print_exc() self.status, self.error = "error", f"{type(exc).__name__}: {exc}"[:600] def _warmup(self) -> None: q = {"w": {"type": "choice", "instructions": "Which team?", "criteria": {"a": "billing", "b": "shipping"}}} for _ in range(2): self._run("warm-up message", q) def decide(self, state, questions: dict) -> dict: if self.status != "ready": raise RuntimeError(f"{self.name} is not ready ({self.status})") with self._lock, torch.inference_mode(): _sync(self.device) t0 = time.perf_counter() answers = self._run(state, questions) _sync(self.device) ms = (time.perf_counter() - t0) * 1000 return dict(answers=answers, ms=round(ms, 1)) def _load(self): raise NotImplementedError def _run(self, state, questions): raise NotImplementedError # ----------------------------------------------------------------------------- LightDec class LightDecBackend(Backend): """A FalconDec checkpoint from the Hugging Face Hub (source "hub") or a local folder (source "local").""" def _load(self): sp, variant = self.spec, self.spec["variant"] if sp["source"] == "local": path = resolve_local_dir(sp["path"], variant, sp["name"], sp["path_env"]) where = dict(path=sp["path"]) else: path = Path(snapshot_download(sp["repo"], revision=sp["revision"])) if variant == "int8": path = path / "compact-int8" where = dict(repo=sp["repo"]) modeling = trusted_modeling(path, os.environ) # DL-SA-002: only known modeling code is executed spec = importlib.util.spec_from_file_location("falcondec_modeling", str(modeling)) fdm = importlib.util.module_from_spec(spec) spec.loader.exec_module(fdm) self.fdm = fdm self.model, self.tok = fdm.load_falcondec(str(path), device=self.device) wfile = next(iter(sorted(path.glob("*.safetensors"))), None) fc = self.model.fcfg self.info.update( **where, variant=variant, params_m=round(self.model.num_parameters() / 1e6, 1), weights_mb=round(wfile.stat().st_size / 1e6) if wfile else None, version=f"{fc.get('name', 'FalconDec')} v{fc.get('version', '?')}", backbone=fc.get("backbone", "jhu-clsp/ettin-encoder-150m"), confidence_native="top-option probability", ) def _run(self, state, questions): out = self.fdm.decide(self.model, self.tok, state, questions, defer_threshold=0.0)["answers"] res = {} for name, q in questions.items(): r = out.get(name) if r is None: continue probs = r.get("probs") or {} if q.get("type") == "noul": res[name] = normalise(q, None, p_true=r.get("p_true", _lookup(probs, "true"))) else: res[name] = normalise(q, probs, choice=r.get("choice"), level=r.get("expected_level")) return res # ----------------------------------------------------------------------------- Arthur class ArthurBackend(Backend): """An Arthur model (config.json + model.safetensors) from the Hub or the models folder. Its code ships with DecisionLab (app/arthur.py); only data is downloaded or read, never code.""" def _load(self): if self.spec.get("source") == "hub": folder = Path(snapshot_download(self.spec["repo"], allow_patterns=["config.json", "model.safetensors"])) where = dict(repo=self.spec["repo"]) else: folder = Path(self.spec["path"]) if not folder.is_dir(): raise FileNotFoundError(f"{self.name} folder not found at {folder}. It was in the models folder at " "start-up; put it back or restart DecisionLab.") where = dict(path=str(folder)) from . import arthur # imported here: torch-heavy, and only needed if Arthur is present self.arthur = arthur self.net, self.temps, cfg = arthur.load_arthur(folder, self.device) wfile = folder / "model.safetensors" self.info.update( **where, params_m=round(sum(p.numel() for p in self.net.parameters()) / 1e6, 1), weights_mb=round(wfile.stat().st_size / 1e6), version=f"Arthur {cfg.get('tier', '?')} (notebook v0.8.0 layout)", confidence_native="top-option probability (temperature-calibrated)", ) def _run(self, state, questions): out = self.arthur.decide(self.net, self.temps, state, questions) res = {} for name, q in questions.items(): r = out.get(name) if r is None: continue if q.get("type") == "noul": res[name] = normalise(q, None, p_true=r["p_true"]) else: res[name] = normalise(q, r["probs"], choice=r.get("choice"), level=r.get("expected_level")) return res # ----------------------------------------------------------------------------- Laya class LayaBackend(Backend): def _load(self): os.environ.setdefault("USE_TF", "0") import laya # noqa: WPS433 (heavy import, done lazily) errors = [] for repo in dict.fromkeys([self.spec["repo"], self.spec["fallback_repo"]]): if not repo: continue try: try: self.agent = laya.load(repo, device=self.device) except TypeError: self.agent = laya.load(repo) self.info["repo"] = repo break except Exception as exc: errors.append(f"{repo}: {type(exc).__name__}: {exc}") else: raise RuntimeError(" | ".join(errors)) if errors: self.info["note"] = f"Primary repo failed, loaded {self.info['repo']} instead" n = None for attr in ("model", "net", "module"): m = getattr(self.agent, attr, None) if isinstance(m, torch.nn.Module): n = sum(p.numel() for p in m.parameters()) break self.info.update( params_m=round(n / 1e6, 1) if n else 421.0, version=f"laya {getattr(laya, '__version__', '?')}", backbone="answerdotai/ModernBERT-large", confidence_native="1 − normalised entropy", ) def _run(self, state, questions): raw = self.agent.predict(state, questions) answers = raw.get("answers", raw) if isinstance(raw, dict) else {} res = {} for name, q in questions.items(): a = answers.get(name) if a is None: continue if not isinstance(a, dict): a = {"value": a} probs = next((a[k] for k in ("probabilities", "probs", "distribution", "scores") if isinstance(a.get(k), dict)), None) qtype = q.get("type", "choice") if qtype == "noul": p = a.get("noul", a.get("p_true", a.get("probability"))) if p is None and probs: p = _lookup(probs, "true") if p is None and isinstance(a.get("value"), (int, float)): p = a["value"] res[name] = normalise(q, probs, p_true=p) else: res[name] = normalise(q, probs, choice=a.get("choice"), level=a.get("score")) return res _KINDS = {"lightdec": LightDecBackend, "arthur": ArthurBackend, "laya": LayaBackend} BACKENDS = {s["key"]: _KINDS[s["kind"]](s) for s in SPECS} # insertion order = comparison order _LOADER = Gate(1) def load_all() -> None: """Load the models one after another in LOAD_ORDER. Runs in a background thread at startup. Only one loader runs at a time: a second call while one is running returns at once.""" if not _LOADER.try_enter(): return try: _load_each() finally: _LOADER.leave() def _load_each() -> None: for key in load_order(os.environ): b = BACKENDS.get(key) if b and b.status in ("idle", "error"): b.load()