File size: 8,181 Bytes
6d40b78 | 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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 | """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 = { # model files each backend reads, relative to the model directory
"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): # "key: value | key: value"; non-string values as compact JSON
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") # mps only by explicit device="mps"
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) # an uncapped MLX cache grows without bound on long runs
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]
|