File size: 3,832 Bytes
c0f7144 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | """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-Nano
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): # type: ignore[misc]
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))}
|