| """AUBIN Mini by Norovox — on-device typed decisions (choice / yes-no) with calibrated probabilities.
|
|
|
| A 395M-parameter cross-encoder: every option is read together with the question and the state, scored once,
|
| and the option scores of a question are turned into probabilities with a fitted temperature.
|
| Runs on CPU with ONNX Runtime only (no torch, no GPU).
|
|
|
| from aubin_mini import AubinMini
|
| m = AubinMini.from_pretrained("emrevrg/AUBIN-Mini") # phone build, 267 MB
|
| m.decide({"ticket": "Package arrived broken, customer wants money back"},
|
| {"route": {"type": "choice", "instructions": "Which team?",
|
| "criteria": {"billing": "refunds", "shipping": "delivery", "tech": "software"}},
|
| "urgent": {"type": "noul", "instructions": "Is this urgent?"}})
|
| """
|
| import json
|
| import os
|
|
|
| import numpy as np
|
|
|
| __version__ = "0.1.0"
|
| __all__ = ["AubinMini", "options_of"]
|
|
|
| VARIANTS = {
|
| "phone": "aubin_mini_nb4_g8.onnx",
|
| "nb4": "aubin_mini_nb4.onnx",
|
| }
|
|
|
|
|
| def options_of(q):
|
| """Question → (answer keys, option texts). Same rendering the model was trained with."""
|
| t, c = q.get("type", "choice"), q.get("criteria")
|
| if t == "noul":
|
| if isinstance(c, dict):
|
| return [False, True], [f"false — {c.get('false', 'no')}", f"true — {c.get('true', 'yes')}"]
|
| return [False, True], ["false — no", "true — yes"]
|
| if isinstance(c, list):
|
| return list(range(len(c))), [str(v) for v in c]
|
| if not isinstance(c, dict) or not c:
|
| raise ValueError("a 'choice' question needs a non-empty 'criteria' dict {option: description}")
|
| return list(c.keys()), [f"{k}: {v}" for k, v in c.items()]
|
|
|
|
|
| def _state_text(state):
|
| return state if isinstance(state, str) else json.dumps(state, ensure_ascii=False)
|
|
|
|
|
| class AubinMini:
|
| def __init__(self, model_dir, variant="phone", threads=None, max_len=512, batch=32):
|
| import onnxruntime as ort
|
| from tokenizers import Tokenizer
|
| cfg_path = os.path.join(model_dir, "aubin_mini.json")
|
| self.cfg = json.load(open(cfg_path, encoding="utf-8")) if os.path.exists(cfg_path) else {}
|
| self.temperature = float(self.cfg.get("temperature", 1.0))
|
| self.max_len, self.batch = max_len, batch
|
| self.tok = Tokenizer.from_file(os.path.join(model_dir, "tokenizer.json"))
|
| self.tok.enable_truncation(max_length=max_len, strategy="only_second")
|
| self.tok.no_padding()
|
| self.tok_long = Tokenizer.from_file(os.path.join(model_dir, "tokenizer.json"))
|
| self.tok_long.enable_truncation(max_length=max_len, strategy="longest_first")
|
| self.tok_long.no_padding()
|
| pad = self.cfg.get("pad_token", "[PAD]")
|
| self.pad_id = self.tok.token_to_id(pad) or 0
|
| so = ort.SessionOptions()
|
| so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
| if threads:
|
| so.intra_op_num_threads = int(threads)
|
| path = os.path.join(model_dir, VARIANTS.get(variant, variant))
|
| self.sess = ort.InferenceSession(path, so, providers=["CPUExecutionProvider"])
|
| self.variant = variant
|
|
|
| @classmethod
|
| def from_pretrained(cls, repo_or_dir="emrevrg/AUBIN-Mini", variant="phone", token=None, **kw):
|
| if os.path.isdir(repo_or_dir):
|
| return cls(repo_or_dir, variant, **kw)
|
| from huggingface_hub import snapshot_download
|
| name = VARIANTS.get(variant, variant)
|
| d = snapshot_download(repo_or_dir, token=token,
|
| allow_patterns=[name, "tokenizer.json", "tokenizer_config.json", "aubin_mini.json"])
|
| return cls(d, variant, **kw)
|
|
|
|
|
| def _scores(self, pairs):
|
| """[(question+option, state)] → raw scores, length-sorted batches (less padding)."""
|
| try:
|
| enc = self.tok.encode_batch(pairs)
|
| except Exception:
|
| enc = []
|
| for p in pairs:
|
| try:
|
| enc.append(self.tok.encode(*p))
|
| except Exception:
|
| enc.append(self.tok_long.encode(*p))
|
| order = sorted(range(len(enc)), key=lambda i: len(enc[i].ids))
|
| out = np.zeros(len(enc), dtype=np.float32)
|
| for k in range(0, len(order), self.batch):
|
| idx = order[k:k + self.batch]
|
| T = max(len(enc[i].ids) for i in idx)
|
| ids = np.full((len(idx), T), self.pad_id, dtype=np.int64)
|
| am = np.zeros((len(idx), T), dtype=np.int64)
|
| for b, i in enumerate(idx):
|
| n = len(enc[i].ids)
|
| ids[b, :n] = enc[i].ids
|
| am[b, :n] = 1
|
| s = self.sess.run(["score"], {"input_ids": ids, "attention_mask": am})[0]
|
| out[idx] = np.asarray(s, dtype=np.float32).reshape(-1)
|
| return out
|
|
|
| def log_probs(self, question, option_texts, state=""):
|
| """Calibrated log-probabilities over option_texts (any number of options)."""
|
| st = "STATE: " + _state_text(state)
|
| s = self._scores([(f"QUESTION: {question}\nOPTION: {t}", st) for t in option_texts]) / self.temperature
|
| s = s - s.max()
|
| return s - np.log(np.exp(s).sum())
|
|
|
| def choose(self, question, options, state=""):
|
| """options: dict {key: description} or list of texts → {answer, confidence, probabilities}."""
|
| if isinstance(options, (list, tuple)):
|
| keys = texts = [str(o) for o in options]
|
| else:
|
| keys, texts = options_of({"type": "choice", "criteria": options})
|
| return self._pack(keys, self.log_probs(question, texts, state))
|
|
|
| def decide(self, state, questions):
|
| """Kev/AUBIN typed-decision JSON: {qid: {type: choice|noul, instructions, criteria}} → {qid: result}."""
|
| flat, spans = [], {}
|
| st = "STATE: " + _state_text(state)
|
| for qid, q in questions.items():
|
| keys, texts = options_of(q)
|
| ins = q.get("instructions") or qid
|
| spans[qid] = (len(flat), keys)
|
| flat += [(f"QUESTION: {ins}\nOPTION: {t}", st) for t in texts]
|
| s = self._scores(flat) / self.temperature
|
| res = {}
|
| for qid, (a, keys) in spans.items():
|
| x = s[a:a + len(keys)]
|
| x = x - x.max()
|
| res[qid] = self._pack(keys, x - np.log(np.exp(x).sum()))
|
| return res
|
|
|
| @staticmethod
|
| def _pack(keys, lp):
|
| p = np.exp(lp)
|
| j = int(p.argmax())
|
| return {"answer": keys[j], "confidence": round(float(p[j]), 4),
|
| "probabilities": {str(k).lower() if isinstance(k, bool) else str(k): round(float(v), 4) for k, v in zip(keys, p)}}
|
|
|
