| """DecisionTune 1.0 inference engine. Standalone: needs only the model files plus torch and transformers (backend "torch"), |
| onnxruntime (backend "onnx") or laya-mlx on Apple silicon (backend "mlx"). |
| |
| The rendering and readout are the ones used for the published Decision Index scores: |
| one sequence per question, [CLS] question [SEP] ([MASK] option)+ [SEP] state [SEP]. A linear head scores the hidden state |
| at each [MASK]; probabilities are the softmax over those scores. A yes/no question uses the options "yes" and "no" and |
| returns P(yes). No calibration temperature, no option filtering, no truncation: a sequence longer than the 8,192-token |
| context is refused. |
| |
| from engine import DecisionModel # run from a directory that holds the model files |
| m = DecisionModel(".") |
| m.choose("Order arrived broken.", "Which team handles this?", ["billing", "shipping", "tech support"]) |
| m.yes_no("Order arrived broken.", "Is the customer asking for a refund?") |
| """ |
| import json |
| import math |
| import os |
| import sys |
|
|
| NAME = "DecisionTune 1.0" |
| SPECIAL = ("[MASK]", "[CLS]", "[SEP]", "[PAD]", "[UNK]") |
| FILES = { |
| "torch": ("config.json", "tokenizer.json", "model.safetensors", "head.safetensors"), |
| "onnx": ("config.json", "tokenizer.json", "onnx/model.onnx"), |
| } |
|
|
|
|
| def clean(s): |
| for t in SPECIAL: |
| s = s.replace(t, " ") |
| return s |
|
|
|
|
| def state_text(s): |
| if s is None: |
| return "" |
| if isinstance(s, str): |
| return s |
| if isinstance(s, dict): |
| return " | ".join(f"{k}: {v if isinstance(v, str) else json.dumps(v, ensure_ascii=False, separators=(',', ':'))}" for k, v in s.items()) |
| return json.dumps(s, ensure_ascii=False, separators=(",", ":")) |
|
|
|
|
| def as_text(x): |
| return x if isinstance(x, str) else json.dumps(x, ensure_ascii=False, separators=(",", ":")) |
|
|
|
|
| def option_text(key, desc): |
| if isinstance(desc, str) and desc.strip(): |
| return desc |
| if desc is not None and desc != "": |
| return as_text(desc) |
| return str(key).replace("_", " ") |
|
|
|
|
| class TorchBackend: |
| def __init__(self, path, device=None, dtype="float32"): |
| import torch |
| from safetensors.torch import load_file |
| from transformers import ModernBertConfig, ModernBertModel |
|
|
| self.torch = torch |
| self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
| dt = getattr(torch, dtype) |
| cfg = ModernBertConfig.from_pretrained(path) |
| self.model = ModernBertModel(cfg) |
| self.model.load_state_dict(load_file(f"{path}/model.safetensors"), strict=True) |
| hw = load_file(f"{path}/head.safetensors") |
| self.head = torch.nn.Linear(cfg.hidden_size, 1) |
| self.head.load_state_dict({"weight": hw["head.weight"], "bias": hw["head.bias"]}) |
| self.model.to(self.device, dt).eval() |
| self.head.to(self.device, dt).eval() |
|
|
| def logits(self, ids, pos): |
| torch = self.torch |
| with torch.inference_mode(): |
| x = torch.tensor([ids], device=self.device) |
| h = self.model(input_ids=x, attention_mask=torch.ones_like(x)).last_hidden_state |
| return self.head(h[0, pos]).squeeze(-1).float().cpu().tolist() |
|
|
|
|
| class OnnxBackend: |
| def __init__(self, path, device=None, dtype="float32"): |
| import numpy as np |
| import onnxruntime as ort |
|
|
| self.np, self.device = np, "cpu" |
| self.sess = ort.InferenceSession(f"{path}/onnx/model.onnx", providers=["CPUExecutionProvider"]) |
|
|
| def logits(self, ids, pos): |
| np = self.np |
| x = np.array([ids], dtype=np.int64) |
| out = self.sess.run(["logits"], {"input_ids": x, "attention_mask": np.ones_like(x), "positions": np.array(pos, dtype=np.int64)})[0] |
| return [float(v) for v in out] |
|
|
|
|
| class MlxBackend: |
| """Apple silicon GPU via the ModernBERT encoder in laya-mlx (Apache-2.0), reading the same safetensors.""" |
|
|
| def __init__(self, path, device=None, dtype="float32"): |
| import mlx.core as mx |
| import numpy as np |
| from laya_mlx.model import EncoderConfig, ModernBert |
| from safetensors.numpy import load_file |
|
|
| mx.set_default_device(mx.gpu) |
| mx.set_cache_limit(4 * 2**30) |
| self.mx, self.np, self.device = mx, np, "mlx" |
| with open(f"{path}/config.json") as f: |
| self.model = ModernBert(EncoderConfig.from_dict(json.load(f))) |
| self.model.load_weights([(k, mx.array(v)) for k, v in load_file(f"{path}/model.safetensors").items()], strict=True) |
| mx.eval(self.model.parameters()) |
| hw = load_file(f"{path}/head.safetensors") |
| self.W, self.b = mx.array(hw["head.weight"][0]), mx.array(hw["head.bias"][0]) |
|
|
| def logits(self, ids, pos): |
| mx, np = self.mx, self.np |
| x = mx.array(np.array([ids], dtype=np.int32)) |
| h = self.model(x, mx.ones_like(x)).astype(mx.float32) |
| return np.array(h[0, mx.array(pos)] @ self.W + self.b).tolist() |
|
|
|
|
| BACKENDS = {"torch": TorchBackend, "onnx": OnnxBackend, "mlx": MlxBackend} |
| FILES["mlx"] = FILES["torch"] |
|
|
|
|
| def pick_backend(backend="auto"): |
| """"auto" = mlx on Apple silicon when laya-mlx is installed, else torch.""" |
| if backend != "auto": |
| return backend |
| import platform |
| import importlib.util |
|
|
| if sys.platform == "darwin" and platform.machine() == "arm64" and importlib.util.find_spec("laya_mlx"): |
| return "mlx" |
| return "torch" |
|
|
|
|
| class DecisionModel: |
| def __init__(self, path, backend="auto", device=None, dtype="float32"): |
| from tokenizers import Tokenizer |
|
|
| backend = pick_backend(backend) |
| if backend not in BACKENDS: |
| raise ValueError(f"backend must be one of {sorted(BACKENDS)}, not {backend!r}") |
| self.path, self.backend_name = path, backend |
| self.tok = Tokenizer.from_file(f"{path}/tokenizer.json") |
| with open(f"{path}/config.json") as f: |
| self.limit = json.load(f)["max_position_embeddings"] |
| self.CLS, self.SEP, self.MARK = (self.tok.token_to_id(t) for t in ("[CLS]", "[SEP]", "[MASK]")) |
| self.backend = BACKENDS[backend](path, device=device, dtype=dtype) |
| self.device = self.backend.device |
|
|
| @classmethod |
| def from_pretrained(cls, path, **kw): |
| """Local model directory. The decision_tune package version also downloads from the Hub and checks SHA-256s.""" |
| return cls(os.path.expanduser(path), **kw) |
|
|
| def _t(self, s): |
| return self.tok.encode(clean(s), add_special_tokens=False).ids |
|
|
| def encode(self, state, instructions, opts): |
| ids, pos = [self.CLS] + self._t(as_text(instructions)) + [self.SEP], [] |
| for o in opts: |
| pos.append(len(ids)) |
| ids += [self.MARK] + self._t(o) |
| ids += [self.SEP] + self._t(state_text(state)) + [self.SEP] |
| if len(ids) > self.limit: |
| raise ValueError(f"sequence of {len(ids)} tokens exceeds the {self.limit}-token context limit") |
| return ids, pos |
|
|
| def logits(self, state, instructions, opts): |
| return self.backend.logits(*self.encode(state, instructions, opts)) |
|
|
| @staticmethod |
| def _softmax(z): |
| top = max(z) |
| w = [math.exp(v - top) for v in z] |
| tot = sum(w) |
| return [v / tot for v in w] |
|
|
| def choose(self, state, question, options): |
| """options: a list of strings, or {key: description}. Returns {"choice", "probabilities", "confidence"}.""" |
| crit = options if isinstance(options, dict) else {o: o for o in options} |
| if len(crit) < 2: |
| raise ValueError("need at least 2 options") |
| p = self._softmax(self.logits(state, question, [option_text(k, d) for k, d in crit.items()])) |
| probs = dict(zip(crit, p)) |
| best = max(probs, key=probs.get) |
| return {"choice": best, "probabilities": probs, "confidence": probs[best]} |
|
|
| def yes_no(self, state, question): |
| """P(yes) for a yes/no question.""" |
| return self._softmax(self.logits(state, question, ["yes", "no"]))[0] |
|
|
