Upload folder using huggingface_hub

code/aubin_mini/__init__.py

@@ -0,0 +1,142 @@
+"""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, 255 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, 255 MB, dev 0.833
+    "nb4": "aubin_mini_nb4.onnx",        # 4-bit block weights, 402 MB, dev 0.840
+}
+
+
+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}."""
+        q = {"type": "choice", "criteria": options}
+        keys, texts = options_of(q)
+        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)}}

code/aubin_mini/__main__.py

@@ -0,0 +1,59 @@
+"""aubin-mini command line.
+
+    aubin-mini decide request.json            # {"state": ..., "questions": {...}} → JSON result
+    aubin-mini serve --port 8009              # POST /decide with the same JSON, GET /health
+"""
+import argparse
+import json
+import sys
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+from . import AubinMini, __version__
+
+
+def main(argv=None):
+    ap = argparse.ArgumentParser(prog="aubin-mini")
+    ap.add_argument("--model", default="emrevrg/AUBIN-Mini", help="HF repo or local directory")
+    ap.add_argument("--variant", default="phone", choices=["phone", "nb4"])
+    ap.add_argument("--threads", type=int, default=None)
+    sub = ap.add_subparsers(dest="cmd", required=True)
+    d = sub.add_parser("decide"); d.add_argument("request", help="JSON file, or - for stdin")
+    s = sub.add_parser("serve"); s.add_argument("--host", default="127.0.0.1"); s.add_argument("--port", type=int, default=8009)
+    a = ap.parse_args(argv)
+    m = AubinMini.from_pretrained(a.model, a.variant, threads=a.threads)
+
+    if a.cmd == "decide":
+        req = json.load(sys.stdin if a.request == "-" else open(a.request, encoding="utf-8"))
+        print(json.dumps(m.decide(req.get("state", ""), req["questions"]), ensure_ascii=False, indent=2))
+        return
+
+    class H(BaseHTTPRequestHandler):
+        def _send(self, code, obj):
+            b = json.dumps(obj, ensure_ascii=False).encode("utf-8")
+            self.send_response(code); self.send_header("Content-Type", "application/json; charset=utf-8")
+            self.send_header("Content-Length", str(len(b))); self.end_headers(); self.wfile.write(b)
+
+        def do_GET(self):
+            if self.path == "/health":
+                self._send(200, {"ok": True, "model": "AUBIN Mini", "version": __version__, "variant": m.variant})
+            else:
+                self._send(404, {"error": "not found"})
+
+        def do_POST(self):
+            if self.path != "/decide":
+                return self._send(404, {"error": "not found"})
+            try:
+                req = json.loads(self.rfile.read(int(self.headers.get("Content-Length", 0))) or b"{}")
+                self._send(200, m.decide(req.get("state", ""), req["questions"]))
+            except (KeyError, ValueError, TypeError) as e:
+                self._send(400, {"error": str(e)})
+
+        def log_message(self, *args):
+            pass
+
+    print(f"AUBIN Mini {__version__} ({m.variant}) on http://{a.host}:{a.port}  POST /decide", flush=True)
+    ThreadingHTTPServer((a.host, a.port), H).serve_forever()
+
+
+if __name__ == "__main__":
+    main()

code/eval_kev.py

@@ -0,0 +1,92 @@
+"""Reproduce AUBIN Mini's numbers on Kev's public locked suites (jaredpalmer/kev-suites, Apache-2.0).
+
+    pip install ./code && python code/eval_kev.py --model emrevrg/AUBIN-Mini --variant phone --out kev_phone.json
+    python code/eval_kev.py --shard 0/4 ...     # split the work across processes / machines, then --merge a.json b.json ...
+
+Same items, same metrics as the AUBIN/Kev protocol: accuracy, Brier, NLL, 15-bin ECE (overall and per source).
+decision-v7 test = sources Kev trained on; transfer-v4 test = 764 records from sources Kev never trained on.
+"""
+import argparse
+import json
+import math
+import time
+
+import numpy as np
+
+from aubin_mini import AubinMini, options_of
+
+SUITES = {"kev_test": "v7/decision-v7/test.jsonl", "kev_transfer_test": "v4/transfer-v4/test.jsonl",
+          "kev_dev": "v7/decision-v7/development.jsonl", "kev_transfer_dev": "v4/transfer-v4/development.jsonl"}
+
+
+def load(name):
+    from huggingface_hub import hf_hub_download
+    p = hf_hub_download("jaredpalmer/kev-suites", SUITES[name], repo_type="dataset")
+    items = []
+    for i, line in enumerate(open(p, encoding="utf-8")):
+        r = json.loads(line)
+        for qid, q in r["questions"].items():
+            keys, texts = options_of(q)
+            lab = q["label"]
+            lab = str(lab).lower() if isinstance(lab, bool) else str(lab)
+            skeys = [str(k).lower() if isinstance(k, bool) else str(k) for k in keys]
+            if lab not in skeys:
+                continue
+            items.append({"src": q.get("src") or r.get("_meta", {}).get("source") or name, "state": r["state"],
+                          "q": q.get("instructions") or qid, "texts": texts, "y": skeys.index(lab)})
+    return items
+
+
+def metrics(rows):
+    by, allr = {}, []
+    for src, y, lp in rows:
+        p = np.exp(np.asarray(lp)); oh = np.zeros(len(p)); oh[y] = 1
+        r = (int(p.argmax()) == y, float(((p - oh) ** 2).sum()), float(p.max()), float(-math.log(max(p[y], 1e-12))))
+        allr.append(r); by.setdefault(src, []).append(r)
+
+    def m(v):
+        ece = 0.0
+        for j in range(15):
+            bb = [x for x in v if j / 15 < x[2] <= (j + 1) / 15]
+            if bb:
+                ece += len(bb) / len(v) * abs(np.mean([x[2] for x in bb]) - np.mean([x[0] for x in bb]))
+        return {"n": len(v), "accuracy": round(float(np.mean([x[0] for x in v])), 4), "brier": round(float(np.mean([x[1] for x in v])), 4),
+                "nll": round(float(np.mean([x[3] for x in v])), 4), "ece15": round(float(ece), 4)}
+    return {"all": m(allr), "by_source": {k: m(v) for k, v in sorted(by.items())}}
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", default="emrevrg/AUBIN-Mini"); ap.add_argument("--variant", default="phone")
+    ap.add_argument("--suites", default="kev_test,kev_transfer_test"); ap.add_argument("--threads", type=int, default=None)
+    ap.add_argument("--shard", default="0/1"); ap.add_argument("--out", default="aubin_mini_kev.json")
+    ap.add_argument("--merge", nargs="*", default=None, help="merge shard outputs instead of running")
+    a = ap.parse_args()
+    if a.merge:
+        rows = {}
+        for f in a.merge:
+            for s, v in json.load(open(f))["rows"].items():
+                rows.setdefault(s, []).extend(v)
+        R = {s: metrics(v) for s, v in rows.items()}
+        json.dump(R, open(a.out, "w"), indent=1); print(json.dumps({s: R[s]["all"] for s in R}))
+        return
+    k, n = map(int, a.shard.split("/"))
+    m = AubinMini.from_pretrained(a.model, a.variant, threads=a.threads)
+    out, ms = {"variant": a.variant, "rows": {}}, []
+    for s in a.suites.split(","):
+        its = load(s)[k::n]
+        rows = []
+        for it in its:
+            t = time.perf_counter()
+            lp = m.log_probs(it["q"], it["texts"], it["state"])
+            ms.append((time.perf_counter() - t) * 1000)
+            rows.append((it["src"], it["y"], [round(float(x), 6) for x in lp]))
+        out["rows"][s] = rows
+        print(s, json.dumps(metrics(rows)["all"]), flush=True)
+    out["ms_per_question"] = {"median": round(float(np.median(ms)), 1), "p95": round(float(np.percentile(ms, 95)), 1)}
+    json.dump(out, open(a.out, "w"))
+    print("ms/question", out["ms_per_question"], flush=True)
+
+
+if __name__ == "__main__":
+    main()

code/pyproject.toml

@@ -0,0 +1,17 @@
+[build-system]
+requires = ["setuptools>=61"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "aubin-mini"
+version = "0.1.0"
+description = "AUBIN Mini by Norovox — on-device typed decisions with calibrated probabilities (ONNX Runtime, CPU)"
+requires-python = ">=3.9"
+license = {text = "Apache-2.0"}
+dependencies = ["onnxruntime>=1.18", "tokenizers>=0.19", "numpy", "huggingface_hub"]
+
+[project.scripts]
+aubin-mini = "aubin_mini.__main__:main"
+
+[tool.setuptools]
+packages = ["aubin_mini"]