"""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 = { # file name, size on disk, accuracy on the same 300 Kev development questions "phone": "aubin_mini_nb4_g8.onnx", # 4-bit block weights + 8-bit embeddings, 267 MB, locked Kev test 0.847 "nb4": "aubin_mini_nb4.onnx", # 4-bit block weights, 422 MB, locked Kev test 0.846 } 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") # the state is cut, never the question/option self.tok.no_padding() self.tok_long = Tokenizer.from_file(os.path.join(model_dir, "tokenizer.json")) # only when question+option alone exceed max_len 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) # ---- scoring ------------------------------------------------------------------------------------------- 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)): # plain list: the answer is the option text itself 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)}}