decisionnode-edge 0.9.2__py3-none-any.whl

decisionnode_edge/model.py

@@ -0,0 +1,273 @@
+"""Load a DecisionNode Edge model directory and answer requests.
+
+    from decisionnode_edge import EdgeModel
+
+    model = EdgeModel.load("path/to/decisionnode-edge")
+    response = model.decide({"state": "...", "questions": {"route": {"type": "choice", "criteria": {...}}}})
+
+The model reads the prompt once per question and never generates text. Its answer is read at the
+<|answer|> position: the last hidden state there is scored against a small readout table
+(readout.safetensors), one row per answer code, and only the rows of the question's codes are read.
+Calibration (calibration.json) turns those scores into probabilities.
+
+A point or box question is a few such prompts, pass by pass (places.py): whether the thing is in the picture,
+then where along each side, each later prompt given what the earlier ones read.
+
+The backbone is read by torch (transformers) or by a faster backend (backends.py: vLLM, OpenVINO, MLX,
+llama.cpp); the readout and calibration are this module's in every case.
+"""
+
+from __future__ import annotations
+
+import base64
+import io
+import json
+import logging
+import os
+from typing import Any
+
+import torch
+from safetensors.torch import load_file
+
+from . import answers, backends, places, request as requests
+from .template import Template
+
+CONFIG_FILE = "decisionnode_edge.json"
+
+log = logging.getLogger("decisionnode_edge")
+
+
+def load_backbone(path: str, device: str | None, dtype: str | torch.dtype | None) -> tuple[Any, bool]:
+    """The transformers backbone on `device` (default CUDA when present, else the CPU) in `dtype` (default bfloat16 on
+    CUDA, float32 elsewhere), and whether it runs under bf16 autocast."""
+    from transformers import AutoModelForImageTextToText
+
+    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+    if dtype is None:
+        dtype = torch.bfloat16 if str(device).startswith("cuda") else torch.float32
+    dtype = getattr(torch, dtype) if isinstance(dtype, str) else dtype
+    backbone = AutoModelForImageTextToText.from_pretrained(path, dtype=dtype).to(device).eval()
+    # bf16 weights run under bf16 autocast, the reference numerics: some operations then compute in float32
+    return backbone, dtype == torch.bfloat16
+
+
+class EdgeModel:
+    def __init__(self, backbone: Any, tokenizer: Any, readout: torch.Tensor, config: dict, calibration: dict, path: str, autocast: bool,
+                 runner: Any = None, torch_options: dict | None = None):  # fmt: skip
+        self.backbone = backbone  # the transformers model; None until torch is needed when a faster backend reads the prompts
+        self.tokenizer = tokenizer
+        self.readout = readout
+        self.config = config
+        self.calibration = calibration
+        self.path = path
+        self.autocast = autocast
+        self.runner = runner  # the backend object (backends.py), or None: torch reads every prompt
+        self.torch_options = torch_options or {}  # device and dtype for torch when it is loaded on first use
+        self.template = Template(tokenizer, bytes.fromhex(config["template"]["code_seed_hex"]).decode("utf-8"))
+        if self.template.first_id != config["readout"]["first_token_id"]:
+            raise ValueError("the tokenizer and the readout table disagree on the ids of the added tokens")
+        self.limits = config["limits"]
+        self._images = None
+
+    @classmethod
+    def load(cls, path: str, device: str | None = None, dtype: str | torch.dtype | None = None, backend: str | None = None,
+             artifacts: str | None = None, options: dict | None = None) -> "EdgeModel":  # fmt: skip
+        """Load the model directory.
+
+        `backend`: "auto" (the default, or DECISIONNODE_EDGE_BACKEND), "torch", "vllm", "openvino", "mlx" or "llamacpp".
+        Auto takes the fastest that is installed and can run here (backends.ORDER). `artifacts`: the directory the build
+        tool wrote for openvino or llamacpp (or DECISIONNODE_EDGE_ARTIFACTS); for either asked for by name it defaults to
+        the model directory when that holds a build (the GGUF that ships with the model). `options`: the backend's settings (the
+        README; DECISIONNODE_EDGE_OPTIONS as JSON is read under them).
+
+        `device` defaults to CUDA when present, else the CPU, and `dtype` to bfloat16 on CUDA (the reference
+        precision) and float32 elsewhere; both are torch's, also when torch only reads the prompts a
+        faster backend cannot take."""
+        from transformers import AutoTokenizer
+
+        with open(os.path.join(path, CONFIG_FILE), encoding="utf-8") as f:
+            config = json.load(f)
+        with open(os.path.join(path, config["calibration"]), encoding="utf-8") as f:
+            calibration = json.load(f)
+        tokenizer = AutoTokenizer.from_pretrained(path)
+        table = load_file(os.path.join(path, config["readout"]["file"]))[config["readout"]["tensor"]]
+        artifacts = artifacts or os.environ.get(backends.ENV_ARTIFACTS) or None
+        if artifacts is None and (backend or os.environ.get(backends.ENV_BACKEND) or "").strip().lower() in ("openvino", "llamacpp") and backends.read_build(path):
+            artifacts = path  # a build that ships in the model folder (the GGUF), for a backend asked for by name
+        options = {**backends.env_options(), **(options or {})}
+        name = backends.choose(backend, device, artifacts, options)
+        log.info("the %s backend reads the prompts", name)
+        if name == "torch":
+            backbone, autocast = load_backbone(path, device, dtype)
+            return cls(backbone, tokenizer, table.float().to(backbone.device), config, calibration, path, autocast)
+        runner = backends.load(name, path, config, artifacts, options, device)
+        # the readout and calibration stay here, in float32 on the CPU
+        return cls(None, tokenizer, table.float(), config, calibration, path, False, runner=runner, torch_options={"device": device, "dtype": dtype})
+
+    @property
+    def device(self) -> torch.device:
+        return self.readout.device
+
+    @property
+    def backend(self) -> str:
+        """The backend that reads the prompts ("torch" also reads, on any backend, the pictures and the prompts longer than
+        the backend takes)."""
+        return self.runner.name if self.runner is not None else "torch"
+
+    # -- public -----------------------------------------------------------
+
+    def decide(self, body: dict) -> dict:
+        """Answer a `/v1/decide` request body. Returns {"model", "answers", "usage"} as the hosted API does."""
+        return self._decide(requests.parse(body))
+
+    def _decide(self, parsed: requests.Request, exact_sum: bool = False) -> dict:
+        """The answers of a parsed request; `exact_sum`: choice and score lists that sum to exactly 1.00 (/v1/systemone)."""
+        probs, input_tokens, placed = self._predict(parsed)
+        out = {}
+        for q in parsed.questions:
+            if q.place is not None:
+                out[q.id] = placed[q.id]
+                continue
+            p = probs[q.id]
+            if q.type == "choice":
+                out[q.id] = answers.choice_answer(q.keys, p, exact_sum)
+            elif q.type == "score":
+                out[q.id] = answers.score_answer(q.levels, p, exact_sum)
+            elif q.type == "truth":
+                out[q.id] = answers.truth_answer(p)
+            else:
+                out[q.id] = answers.number_answer(q.grid, p)
+        return {"model": self.config["name"], "answers": out, "usage": {"input_tokens": input_tokens, "output_tokens": 0}}
+
+    def predict(self, body: dict) -> dict[str, list[float]]:
+        """The unrounded calibrated probabilities of every question: per option or level in request order, [yes, no]
+        for truth, per grid value for number. A point or box question is answered by `decide` only."""
+        parsed = requests.parse(body)
+        for q in parsed.questions:
+            if q.place is not None:
+                raise requests.RequestError(f"question {q.id!r}: a {q.type} question is answered by decide(), not predict()")
+        return self._predict(parsed)[0]
+
+    # -- internals --------------------------------------------------------
+
+    def _predict(self, parsed: requests.Request) -> tuple[dict[str, list[float]], int, dict[str, dict]]:
+        """(the probabilities of every question but the point and box ones, the input tokens, the point and box answers)."""
+        images = self._load_images(parsed.images) if parsed.images else None
+        image_ids = [image["id"] for image in parsed.images]
+        out: dict[str, list[float]] = {}
+        state_tokens, total = None, 0
+
+        def read(qid: str, row: dict) -> list[float]:
+            nonlocal state_tokens, total
+            encoded = self.template.encode(row, parsed.state, image_ids, images[2] if images else None)
+            if len(encoded.input_ids) > self.limits["max_input_tokens"]:
+                raise requests.RequestError(f"question {qid!r} needs {len(encoded.input_ids):,} tokens; the maximum context length is {self.limits['max_input_tokens']:,} tokens",
+                                            error_type="max_tokens_exceeded")  # fmt: skip
+            probs = self._probabilities(encoded, row, images, bool(image_ids))
+            # the state and its pictures count once, every question's own tokens on top (as the hosted usage counts)
+            total += len(encoded.input_ids) - (encoded.state_tokens if state_tokens is not None else 0)
+            state_tokens = encoded.state_tokens
+            return probs
+
+        specs = []
+        for q in parsed.questions:
+            if q.place is not None:
+                specs.append(q.place)
+                continue
+            answered = [read(q.id, row) for row in q.rows]
+            if q.type == "number":
+                out[q.id] = [1.0] if not q.rows else requests.compose(q.grid, answered)
+            else:
+                out[q.id] = answered[0] if answered else [1.0]
+        placed: dict[str, dict] = {}
+        if specs:  # every point and box question's readings, pass by pass, each pass given the readings before it
+            plan = places.Plan(specs)
+            while True:
+                rows = plan.rows()
+                if not rows:
+                    break
+                plan.take(rows, [read(plan.chains[n].spec.qid, row) for n, row in rows])
+            placed = plan.answers(self._image_sizes(parsed.images))
+        return out, total, placed
+
+    def _probabilities(self, encoded: Any, row: dict, images: tuple | None, has_images: bool) -> list[float]:
+        """One prompt's class probabilities: the readout's scores of its codes, calibrated (a coordinate row's softmax)."""
+        return answers.calibrated(self._logits(encoded, images), row["type"], self.calibration, has_images)
+
+    @torch.inference_mode()
+    def _logits(self, encoded: Any, images: tuple | None) -> list[float]:
+        if self.runner is not None and images is None and self.runner.accepts(len(encoded.input_ids)):
+            h = self.runner.hidden(encoded.input_ids).to(self.device, torch.float32)
+        else:
+            h = self._torch_hidden(encoded, images).to(self.device)
+        scores = h @ self.readout.t()  # one score per added token, never the vocabulary
+        rows = torch.tensor(encoded.class_ids, device=self.device) - self.template.first_id
+        return scores[rows].tolist()
+
+    def _torch_hidden(self, encoded: Any, images: tuple | None) -> torch.Tensor:
+        """The hidden state at the <|answer|> position (the last token) from the transformers backbone, float32."""
+        if self.backbone is None:
+            log.warning("loading torch for the prompts the %s backend does not take (pictures, longer inputs)", self.backend)
+            self.backbone, self.autocast = load_backbone(self.path, self.torch_options.get("device"), self.torch_options.get("dtype"))
+        device = self.backbone.device
+        input_ids = torch.tensor([encoded.input_ids], device=device)
+        extra = {}
+        if images is not None:
+            extra = {
+                "pixel_values": images[0].to(device),
+                "image_grid_thw": images[1].to(device),
+                "mm_token_type_ids": (input_ids == self.template.image_pad).long(),
+            }
+        with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=self.autocast):
+            hidden = self.backbone.model(input_ids=input_ids, attention_mask=torch.ones_like(input_ids), use_cache=False, **extra).last_hidden_state
+        return hidden[0, -1].float()
+
+    def _image_sizes(self, images: list[dict]) -> dict[str, tuple[int, int]]:
+        return image_sizes(images)
+
+    def _load_images(self, images: list[dict]) -> tuple[torch.Tensor, torch.Tensor, dict[str, int]]:
+        """(pixel values, grid, placeholder tokens per image id) of a request's pictures: EXIF orientation applied,
+        RGB, the longest side at most max_image_side, then the backbone's own image processor."""
+        from PIL import Image, ImageOps
+        from transformers import AutoImageProcessor
+
+        if self._images is None:
+            self._images = AutoImageProcessor.from_pretrained(self.path)
+        max_side = int(self.limits["max_image_side"])
+        fitted = []
+        for image in images:
+            source = image["path"] if image.get("path") is not None else io.BytesIO(base64.b64decode(image["data"]))
+            try:
+                with Image.open(source) as im:
+                    im = ImageOps.exif_transpose(im).convert("RGB")
+                    if max(im.size) > max_side:
+                        scale = max_side / max(im.size)
+                        im = im.resize((max(1, round(im.width * scale)), max(1, round(im.height * scale))), Image.BICUBIC)
+                    fitted.append(im.copy())
+            except (OSError, ValueError) as error:
+                raise requests.RequestError(f"Image {image['id']!r} could not be read: {error}") from None
+        out = self._images(images=fitted, return_tensors="pt")
+        merge = int(getattr(self._images, "merge_size", 2))
+        counts = (out["image_grid_thw"].prod(dim=-1) // merge**2).tolist()
+        return out["pixel_values"], out["image_grid_thw"], {image["id"]: int(n) for image, n in zip(images, counts)}
+
+
+def image_sizes(images: list[dict]) -> dict[str, tuple[int, int]]:
+    """Per picture id, its width and height as sent after its EXIF orientation (the sides swapped when the orientation
+    turns it a quarter, 5 to 8), read from the file's header: the size a point or box answer is given in."""
+    from PIL import Image
+
+    out = {}
+    for image in images:
+        source = image["path"] if image.get("path") is not None else io.BytesIO(base64.b64decode(image["data"]))
+        try:
+            with Image.open(source) as im:
+                width, height = im.size
+                try:
+                    orientation = int(im.getexif().get(0x0112, 1))
+                except Exception:
+                    orientation = 1
+        except (OSError, ValueError) as error:
+            raise requests.RequestError(f"Image {image['id']!r} could not be read: {error}") from None
+        out[image["id"]] = (height, width) if orientation in (5, 6, 7, 8) else (width, height)
+    return out