| """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. |
| |
| 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, 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() |
| |
| 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): |
| self.backbone = backbone |
| self.tokenizer = tokenizer |
| self.readout = readout |
| self.config = config |
| self.calibration = calibration |
| self.path = path |
| self.autocast = autocast |
| self.runner = runner |
| self.torch_options = torch_options or {} |
| 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": |
| """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 |
| 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) |
| |
| 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" |
|
|
| |
|
|
| def decide(self, body: dict) -> dict: |
| """Answer a `/v1/decide` request body. Returns {"model", "answers", "usage"} as the hosted API does.""" |
| parsed = requests.parse(body) |
| probs, input_tokens = self._predict(parsed) |
| out = {} |
| for q in parsed.questions: |
| p = probs[q.id] |
| if q.type == "choice": |
| out[q.id] = answers.choice_answer(q.keys, p) |
| elif q.type == "score": |
| out[q.id] = answers.score_answer(q.levels, p) |
| 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.""" |
| return self._predict(requests.parse(body))[0] |
|
|
| |
|
|
| def _predict(self, parsed: requests.Request) -> tuple[dict[str, list[float]], int]: |
| 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 |
| for q in parsed.questions: |
| answered = [] |
| for row in q.rows: |
| 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 {q.id!r} needs {len(encoded.input_ids):,} tokens; the maximum context length is {self.limits['max_input_tokens']:,} tokens", |
| error_type="max_tokens_exceeded") |
| logits = self._logits(encoded, images) |
| answered.append(answers.calibrated(logits, row["type"], self.calibration, bool(image_ids))) |
| |
| total += len(encoded.input_ids) - (encoded.state_tokens if state_tokens is not None else 0) |
| state_tokens = encoded.state_tokens |
| 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] |
| return out, total |
|
|
| @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() |
| 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 _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)} |
|
|
