app/models.py

"""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()