File size: 7,276 Bytes
995bd7d 6620bb4 995bd7d 6620bb4 995bd7d 6620bb4 995bd7d | 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 | """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)}}
|