| """Decision Index kit engine (`--engine matilda_decision_engine:MatildaDecisionEngine`, repository root on PYTHONPATH). |
| |
| PYTHONPATH=. python -m decision_index pipeline --engine matilda_decision_engine:MatildaDecisionEngine \ |
| --option checkpoint=. --out runs/Matilda-Decision-v1-Edge |
| |
| Images go in as `engine(state, questions, images=[...])`. Options: checkpoint, device, max_image_pixels |
| (default 262,144), max_tokens, token_budget, batch_size. |
| """ |
|
|
| import hashlib |
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
| from decision_index.engines import Engine, Unsupported |
|
|
| from matilda_decision_runtime import MAX_IMAGE_PIXELS, CapacityError, DecisionModel, decide |
|
|
|
|
| class MatildaDecisionEngine(Engine): |
| name = "matilda-decision" |
| latency = "In-process request wall time: prompt rendering, image decoding, tokenization and one forward pass per question batch." |
|
|
| def __init__(self, checkpoint: str = ".", device: str | None = None, max_image_pixels: int = MAX_IMAGE_PIXELS, |
| max_tokens: int = 131072, token_budget: int = 131072, batch_size: int = 64, **options: Any) -> None: |
| super().__init__(**options) |
| self.model = DecisionModel(checkpoint, device=device, max_image_pixels=int(max_image_pixels)) |
| self.max_tokens, self.token_budget, self.batch_size = int(max_tokens), int(token_budget), int(batch_size) |
| config = Path(checkpoint) / "decision_config.json" |
| self.label = str(self.model.config.get("name", "matilda-decision")) |
| self.provenance = {"checkpoint": str(Path(checkpoint).resolve()), "revision": self.model.revision, |
| "checkpoint_config_sha256": hashlib.sha256(config.read_bytes()).hexdigest(), |
| "attention": self.model.attention_mode, "temperature": self.model.temperature, |
| "max_image_pixels": self.model.max_image_pixels, "max_tokens": self.max_tokens} |
|
|
| def runtime(self) -> dict[str, object]: |
| if not torch.cuda.is_available(): |
| return {"device": "cpu", "torch": torch.__version__} |
| properties = torch.cuda.get_device_properties(0) |
| return {"device": properties.name, "memory_gib": round(properties.total_memory / 2**30), "torch": torch.__version__} |
|
|
| def synchronize(self) -> None: |
| if torch.cuda.is_available(): |
| torch.cuda.synchronize() |
|
|
| def __call__(self, state: object, questions: dict[str, dict[str, Any]], images: Any = None) -> tuple[dict[str, object], None]: |
| for question in questions.values(): |
| if question["type"] not in ("choice", "noul"): |
| raise Unsupported(f"question type {question['type']!r}") |
| try: |
| distributions, input_tokens = decide(self.model, state, questions, images=list(images or ()), |
| max_tokens=self.max_tokens, token_budget=self.token_budget, |
| batch_size=self.batch_size) |
| except CapacityError as error: |
| raise Unsupported(str(error)) from error |
| answers = {key: self._answer(questions[key], values) for key, values in distributions.items()} |
| return {"model": self.label, "answers": answers, "usage": {"input_tokens": input_tokens}}, None |
|
|
| @staticmethod |
| def _answer(question: dict[str, Any], values: list[float]) -> dict[str, object]: |
| """The kit's answer shapes: noul -> P(true); choice -> argmax key plus every option's probability.""" |
| if question["type"] == "noul": |
| return {"type": "noul", "noul": values[1]} |
| keys = list(question["criteria"]) |
| return {"type": "choice", "choice": keys[max(range(len(keys)), key=values.__getitem__)], |
| "probabilities": dict(zip(keys, values, strict=True))} |
|
|
