| """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 |
|
|
| |
| 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: |
| 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): |
| 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) |
|
|
