File size: 11,210 Bytes
45507bb | 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 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 | """Deployment API: the same `decide()` over an ONNX model (numpy + onnxruntime
only, no torch) or a PyTorch checkpoint.
from minijev import MiniJEV
model = MiniJEV.load("model.onnx")
decision = model.decide(state=state, question="Choose the next move",
choices=["UP", "DOWN", "LEFT", "RIGHT"])
decision.choice, decision.confidence, decision.scores
"""
from __future__ import annotations
import json
from functools import lru_cache
from pathlib import Path
from typing import Any, Sequence
import numpy as np
from minijev.keys import STATE_INPUTS
from minijev.tokenizers import ByteTokenizer
from minijev.tasks import Task, get_task
from minijev.types import BooleanDecision, Decision
META_KEY = "minijev"
class Decider:
"""Common front-end: tokenisation, state encoding, Decision objects.
`run` returns *calibrated* probabilities/logits (logits / T); the raw, uncalibrated
distribution is softmax(logits * T) and is reported next to it in every Decision.
"""
task: Task
tokenizer = ByteTokenizer()
max_question_len: int = 32
max_choice_len: int = 16
temperature: float = 1.0
# -- to implement: probs, logits, embedding from model-ready arrays
def _run(self, state: np.ndarray, question_ids: np.ndarray, choice_ids: np.ndarray):
raise NotImplementedError
@lru_cache(maxsize=256)
def _tokens(self, question: str, choices: tuple[str, ...]):
dyn = self.tokenizer.kind != "byte"
q = self.tokenizer.encode_batch([question], self.max_question_len, dynamic=dyn)
c = self.tokenizer.encode_batch(list(choices), self.max_choice_len, dynamic=dyn)[None]
return q, c
def encode(self, states: Sequence[Any]) -> np.ndarray:
enc = self.task.encode_states(states)
return next(iter(enc.values()))
def probs(self, states: Sequence[Any], question: str, choices: Sequence[str],
encoded: np.ndarray | None = None) -> np.ndarray:
return self.run(states, question, choices, encoded)[0]
def run(self, states, question, choices, encoded=None):
q, c = self._tokens(question, tuple(choices))
x = self.encode(states) if encoded is None else encoded
B = x.shape[0]
return self._run(x, np.repeat(q, B, 0), np.repeat(c, B, 0))
def decide_batch(self, states: Sequence[Any], question: str, choices: Sequence[str],
return_embedding: bool = False) -> list[Decision]:
if len(set(choices)) != len(choices):
raise ValueError(f"duplicate choices: {list(choices)}")
probs, logits, emb = self.run(states, question, choices)
raw = _softmax(logits * self.temperature)
out = []
for i in range(len(probs)):
k = int(probs[i].argmax())
out.append(Decision(choice=choices[k], index=k, confidence=float(probs[i, k]),
scores={ch: float(p) for ch, p in zip(choices, probs[i])},
logits=[float(v) for v in logits[i]],
embedding=emb[i].tolist() if return_embedding else None,
raw_scores={ch: float(p) for ch, p in zip(choices, raw[i])},
temperature=self.temperature))
return out
def decide(self, state: Any, question: str | None = None, choices: Sequence[str] | None = None,
return_embedding: bool = False) -> Decision:
question = question if question is not None else self.task.default_question
choices = list(choices) if choices is not None else list(self.task.default_choices)
return self.decide_batch([state], question, choices, return_embedding)[0]
def choice(self, state: Any, question: str, choices: Sequence[str]) -> str:
return self.decide(state, question, choices).choice
def boolean(self, state: Any, question: str, labels: tuple[str, str] = ("yes", "no")) -> BooleanDecision:
"""Yes/no question, answered internally as a 2-choice decision."""
d = self.decide(state, question, list(labels))
p_yes = d.scores[labels[0]]
return BooleanDecision(answer=p_yes >= 0.5, confidence=max(p_yes, 1 - p_yes), p_yes=p_yes,
raw_p_yes=d.raw_scores[labels[0]])
def _softmax(x: np.ndarray) -> np.ndarray:
z = x - x.max(-1, keepdims=True)
e = np.exp(z)
return e / e.sum(-1, keepdims=True)
class OnnxDecider(Decider):
def __init__(self, path: str | Path, threads: int = 1, providers: list[str] | None = None):
import onnxruntime as ort
so = ort.SessionOptions()
so.intra_op_num_threads = threads
so.inter_op_num_threads = 1
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
self.session = ort.InferenceSession(str(path), so, providers=providers or ["CPUExecutionProvider"])
meta = json.loads(self.session.get_modelmeta().custom_metadata_map[META_KEY])
self.meta = meta
self.task = get_task(meta["task"], meta.get("task_cfg", {}))
self.max_question_len = meta["max_question_len"]
self.max_choice_len = meta["max_choice_len"]
self.state_input = meta["state_input"]
self.path = str(path)
def _run(self, state, question_ids, choice_ids):
probs, logits, emb = self.session.run(None, {self.state_input: state, "question_ids": question_ids,
"choice_ids": choice_ids})
return probs, logits, emb
class TorchDecider(Decider):
def __init__(self, checkpoint: str | Path | None = None, model=None, cfg: dict | None = None,
device: str = "cpu", threads: int | None = None, task: str | None = None):
import torch
from minijev.choices import EmbeddingCache
from minijev.tokenizers import make_tokenizer
from minijev.training.supervised import make_model, make_task_for, task_specs
if threads:
torch.set_num_threads(threads)
task_T = {}
if model is None:
ck = torch.load(checkpoint, map_location=device, weights_only=False)
cfg = ck["config"]
task_T = ck.get("task_temperatures") or {}
model = make_model(cfg)
model.load_state_dict(ck["model"])
self.specs = {s["key"]: s["task"] for s in task_specs(cfg)}
self.tokenizer = make_tokenizer(cfg["model"].get("text_encoder"))
self.torch, self.model, self.cfg, self.device = torch, model.to(device).eval(), cfg, device
self.max_question_len = cfg["model"].get("max_question_len", 32)
self.max_choice_len = cfg["model"].get("max_choice_len", 16)
self.global_temperature = float(model.log_temperature.exp())
self.task_temperatures = task_T
self.cache = EmbeddingCache() if model.backbone is not None and model.backbone.frozen else None
self._make_task = lambda key: make_task_for(cfg, self.specs[key])
self.use_task(task or next(iter(self.specs)))
@classmethod
def from_pretrained(cls, path_or_repo: str | Path, device: str = "cpu", threads: int | None = None,
task: str | None = None) -> "TorchDecider":
"""Load a folder written by `minijev export-hf` (local path or Hugging Face repo id)."""
from minijev.hub import load_pretrained
model, cfg, task_T = load_pretrained(path_or_repo, device)
dec = cls(model=model, cfg=cfg, device=device, threads=threads, task=task)
dec.task_temperatures = task_T
return dec.use_task(dec.task_key)
def use_task(self, key: str) -> "TorchDecider":
"""Switch the state format (and per-task temperature) for a multi-task model."""
self.task_key = key
self.task = self._make_task(key)
self.state_key = STATE_INPUTS[self.task.state_encoder]
self.temperature = float(self.task_temperatures.get(key, self.global_temperature))
return self
def _bb(self, texts: list[str], max_len: int):
torch = self.torch
dev = self.device
def compute(missing):
ids = torch.as_tensor(self.tokenizer.encode_batch(missing, max_len, dynamic=True), device=dev)
return self.model.backbone(ids)
tok, mask, pooled = self.cache.get_many(texts, compute)
return tok.to(dev), mask.to(dev), pooled.to(dev)
def forward(self, states, question: str, choices: Sequence[str], encoded=None):
torch = self.torch
x = self.encode(states) if encoded is None else encoded
B = x.shape[0]
state = {self.state_key: torch.as_tensor(x, device=self.device)}
with torch.inference_mode():
if self.cache is not None: # frozen backbone: each question / label is encoded once, ever
qt, qm, qp = self._bb([question], self.max_question_len)
_, _, cp = self._bb(list(choices), self.max_choice_len)
cm = torch.ones(B, len(choices), dtype=torch.bool, device=self.device)
return self.model(state, question_bb=(qt.expand(B, -1, -1), qm.expand(B, -1), qp.expand(B, -1)),
choice_bb=(cp[None].expand(B, -1, -1), cm))
q, c = self._tokens(question, tuple(choices))
return self.model(state, torch.as_tensor(np.repeat(q, B, 0), device=self.device),
torch.as_tensor(np.repeat(c, B, 0), device=self.device))
def run(self, states, question, choices, encoded=None):
torch = self.torch
out = self.forward(states, question, choices, encoded)
logits = out["logits"] / self.temperature
probs = torch.softmax(logits, -1)
return probs.cpu().numpy(), logits.cpu().numpy(), out["embedding"].float().cpu().numpy()
def _run(self, state, question_ids, choice_ids): # kept for callers passing model-ready arrays
torch = self.torch
with torch.inference_mode():
out = self.model({self.state_key: torch.as_tensor(state, device=self.device)},
torch.as_tensor(question_ids, device=self.device),
torch.as_tensor(choice_ids, device=self.device))
logits = out["logits"] / self.temperature
return torch.softmax(logits, -1).cpu().numpy(), logits.cpu().numpy(), out["embedding"].float().cpu().numpy()
def score(self, state: Any, question: str) -> float:
"""Continuous 0..1 score from the scalar head (meaningful once trained with `score` targets)."""
out = self.forward([state], question, ["yes", "no"])
return float(self.torch.sigmoid(out["score"])[0])
class MiniJEV:
"""Entry point: `MiniJEV.load(path)` for .onnx or .pt files."""
@staticmethod
def load(path: str | Path, threads: int = 1, **kw) -> Decider:
p = Path(path)
if p.suffix == ".onnx":
return OnnxDecider(p, threads=threads, **kw)
if p.suffix == ".pt":
return TorchDecider(p, threads=threads, **kw)
return TorchDecider.from_pretrained(path, threads=threads, **kw) # export-hf folder or hub repo id
|