| """Matilda Decision for 🤗 Transformers (`trust_remote_code=True`). |
| |
| `AutoModel.from_pretrained(repo, trust_remote_code=True)` loads the repository through its own runtime |
| (`matilda_decision_runtime.py`: bidirectional prompt attention and the 255-option decision readout) and returns a |
| model with `system_one(state=..., questions={...}, images=[...])`. Without `trust_remote_code` the repository loads |
| as the plain, causal `Qwen3_5Model` backbone, which does not give the model's decisions. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import os |
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
| from transformers import PreTrainedModel |
| from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5Config |
|
|
| try: |
| from .matilda_decision_runtime import MAX_IMAGE_PIXELS, DecisionModel, system_one |
| except ImportError: |
| from matilda_decision_runtime import MAX_IMAGE_PIXELS, DecisionModel, system_one |
|
|
| HUB_OPTIONS = ("revision", "cache_dir", "token", "local_files_only", "force_download") |
|
|
|
|
| class MatildaDecisionModel(PreTrainedModel): |
| """A Matilda Decision checkpoint behind `system_one(state=..., questions={...}, images=[...])`.""" |
|
|
| config_class = Qwen3_5Config |
| base_model_prefix = "decision" |
| _no_split_modules: list[str] = [] |
|
|
| def __init__(self, config: Qwen3_5Config): |
| super().__init__(config) |
| self.runtime: DecisionModel | None = None |
| self.post_init() |
|
|
| def _init_weights(self, module: Any) -> None: |
| """Every weight comes from the checkpoint.""" |
|
|
| @classmethod |
| def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike, *model_args: Any, |
| config: Qwen3_5Config | None = None, **kwargs: Any) -> MatildaDecisionModel: |
| """Load a Hub repository or a local download. |
| |
| `device` (or `device_map` naming one device) picks the device, default cuda:0 when a GPU is visible. |
| `max_image_pixels` caps each image's resized size (default 262,144 = 0.26 MP, as trained). |
| Numerics are fixed by the checkpoint (BF16 backbone, FP32 readout). |
| """ |
| hub = {key: kwargs.pop(key) for key in HUB_OPTIONS if key in kwargs} |
| device = kwargs.pop("device", None) |
| device_map = kwargs.pop("device_map", None) |
| if isinstance(device_map, dict): |
| if set(device_map) != {""}: |
| raise ValueError("Matilda Decision runs on one device: pass a device name or {'': device}") |
| device_map = device_map[""] |
| if device is None and device_map not in (None, "auto"): |
| device = f"cuda:{device_map}" if isinstance(device_map, int) else str(device_map) |
| max_image_pixels = int(kwargs.pop("max_image_pixels", MAX_IMAGE_PIXELS)) |
| for key in ("dtype", "torch_dtype"): |
| if kwargs.pop(key, None) not in (None, "auto", "bfloat16", torch.bfloat16): |
| raise ValueError(f"{key}: numerics are fixed by the checkpoint (BF16 backbone, FP32 readout)") |
| local = Path(os.fspath(pretrained_model_name_or_path)).expanduser() |
| if not local.is_dir(): |
| from huggingface_hub import snapshot_download |
|
|
| local = Path(snapshot_download(str(pretrained_model_name_or_path), |
| **{key: value for key, value in hub.items() if value is not None})) |
| if config is None: |
| config = cls.config_class.from_pretrained(str(local)) |
| model = cls(config) |
| model.runtime = DecisionModel(local, device=device, max_image_pixels=max_image_pixels) |
| model.name_or_path = str(pretrained_model_name_or_path) |
| model.eval() |
| return model |
|
|
| def _require(self) -> DecisionModel: |
| if self.runtime is None: |
| raise RuntimeError("Load the model with from_pretrained") |
| return self.runtime |
|
|
| @property |
| def decision_config(self) -> dict[str, Any]: |
| return self._require().config |
|
|
| def system_one(self, *, state: Any, questions: dict[str, Any], images: list[Any] | None = None, |
| batch_size: int = 64, token_budget: int = 131072, max_tokens: int = 131072) -> dict[str, Any]: |
| """Typed choice / noul / score answers about one state: `{"model", "answers", "usage"}`. |
| |
| `images`: any number of images every question sees (PIL images, local paths, http(s) URLs or base64 |
| `data:image/...` URLs), placed before the text and resized to at most 0.26 MP each. Nothing is truncated: |
| a request longer than `max_tokens` raises an error. |
| """ |
| return system_one(self._require(), state=state, questions=questions, images=images, |
| max_tokens=max_tokens, token_budget=token_budget, batch_size=batch_size) |
|
|
| def forward(self, state: Any = None, questions: dict[str, Any] | None = None, |
| images: list[Any] | None = None) -> dict[str, Any]: |
| return self.system_one(state=state, questions=questions or {}, images=images) |
|
|
| def to(self, *args: Any, **kwargs: Any) -> MatildaDecisionModel: |
| device, dtype, _, _ = torch._C._nn._parse_to(*args, **kwargs) |
| if dtype is not None: |
| raise TypeError("Numerics are fixed by the checkpoint; only the device can change") |
| if device is not None: |
| runtime = self._require() |
| runtime.to(device) |
| runtime.device_name = str(device) |
| return self |
|
|
| def train(self, mode: bool = True) -> MatildaDecisionModel: |
| if mode: |
| raise RuntimeError("Matilda Decision runs in inference mode only") |
| return super().train(False) |
|
|
| def save_pretrained(self, *args: Any, **kwargs: Any) -> None: |
| raise NotImplementedError("The repository itself is the package; copy it with huggingface_hub.snapshot_download") |
|
|
