DecisionNode Edge Mini 0.9.2: point and box questions, the /v1/systemone route as hosted, a box calibration
0e77dfd verified | """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 | |
| 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}) | |
| def device(self) -> torch.device: | |
| return self.readout.device | |
| 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) | |
| 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 | |