"""Minimal inference code for Dinah-0 Apache. Dinah-0 Apache is a 150M-parameter encoder (ModernBERT architecture) that answers typed decisions: choice pick one label out of a set -> {"choice", "confidence", "probabilities"} noul probability that a statement is true -> {"noul"} score expected value on an ordered rubric -> {"score", "confidence", "probabilities"} Every question is read as one sequence [CLS] instructions [SEP] state [SEP] [OPT] option_1 [OPT] option_2 ... [OPT] option_n and each option is scored from the hidden state at its [OPT] marker. Nothing is truncated: a request longer than `max_length` (8,192 tokens) raises `RequestTooLong`. Two backends with the same API: from dinah import Dinah, DinahONNX model = Dinah.from_pretrained("Lukitaduarte/dinah-0-apache") # PyTorch, fp32 weights model = DinahONNX.from_pretrained("Lukitaduarte/dinah-0-apache") # onnxruntime, int8, CPU model.choice(state="The package arrived broken.", criteria={"refund": "Customer wants money back", "replace": "Customer wants a new unit"}, instructions="What does the customer want?") """ from __future__ import annotations import json import math from pathlib import Path from typing import Any import numpy as np QUESTION_TYPES = {"noul": 0, "choice": 1, "score": 2} NEG = -1e4 class RequestTooLong(ValueError): """The request does not fit in `max_length` tokens. Nothing is ever truncated.""" def render(content: Any) -> str: """Strings are stripped; objects and arrays become compact JSON with sorted keys (as in training).""" if content is None: return "" if isinstance(content, str): return content.strip() return json.dumps(content, ensure_ascii=False, sort_keys=True, separators=(",", ":")) def _resolve(path_or_repo: str | Path, files: list[str], revision: str | None = None) -> Path: path = Path(path_or_repo) if path.is_dir(): return path from huggingface_hub import snapshot_download return Path(snapshot_download(str(path_or_repo), allow_patterns=files, revision=revision)) class _Base: """Tokenization, batching and the typed-decision API shared by both backends.""" def __init__(self, directory: Path): from transformers import AutoTokenizer self.config = json.loads((directory / "config.json").read_text()) extra = self.config["dinah"] self.tokenizer = AutoTokenizer.from_pretrained(str(directory)) self.opt_id = self.tokenizer.convert_tokens_to_ids(extra["opt_token"]) assert self.opt_id == extra["opt_token_id"], "tokenizer does not match the model" self.max_length = extra["max_length"] self.pad_multiple = extra["pad_multiple"] cal = extra.get("calibration") self.temperature = np.array(cal["temperature"], dtype=np.float64) if cal else None self.platt = tuple(cal["platt"]) if cal else None def _ids(self, text: str) -> list[int]: return self.tokenizer(text, add_special_tokens=False)["input_ids"] if text else [] def _encode(self, instructions: Any, state: Any, options: list[Any]) -> tuple[list[int], list[int]]: t = self.tokenizer ids = [t.cls_token_id, *self._ids(render(instructions)), t.sep_token_id, *self._ids(render(state)), t.sep_token_id] positions = [] for option in options: positions.append(len(ids)) ids += [self.opt_id, *self._ids(render(option))] if len(ids) > self.max_length: raise RequestTooLong(f"request needs {len(ids)} tokens, above max_length={self.max_length} (no truncation)") return ids, positions def _batch(self, items: list[tuple[list[int], list[int], int]]) -> dict[str, np.ndarray]: b = len(items) length = max(len(ids) for ids, _, _ in items) length = -(-length // self.pad_multiple) * self.pad_multiple n = max(len(pos) for _, pos, _ in items) batch = {"ids": np.full((b, length), self.tokenizer.pad_token_id, dtype=np.int64), "pad_mask": np.zeros((b, length), dtype=bool), "opt_positions": np.zeros((b, n), dtype=np.int64), # padding points at [CLS], masked below "opt_mask": np.zeros((b, n), dtype=bool), "qtype": np.array([q for _, _, q in items], dtype=np.int64)} for i, (ids, pos, _) in enumerate(items): batch["ids"][i, :len(ids)] = ids batch["pad_mask"][i, :len(ids)] = True batch["opt_positions"][i, :len(pos)] = pos batch["opt_mask"][i, :len(pos)] = True return batch def _forward(self, batch: dict[str, np.ndarray]) -> tuple[np.ndarray, np.ndarray]: raise NotImplementedError def predict(self, questions: list[dict[str, Any]], batch_size: int = 16) -> list[dict[str, Any]]: """Answer a list of questions. Each question is a dict with `type` ("choice", "noul" or "score"), `state`, optional `instructions`, and `criteria`: choice: {label: description or None} (None means the label is its own description) noul: {"true": ..., "false": ...} describing both poles (strongly recommended) score: list of descriptions, ordered from lowest (0) to highest """ encoded, meta = [], [] for q in questions: kind = q["type"] if kind == "choice": labels = list(q["criteria"]) options = [d if d is not None else label for label, d in q["criteria"].items()] elif kind == "noul": c = q.get("criteria") or {} options = [c.get("false") or "The statement above is false.", c.get("true") or q.get("instructions")] labels = None elif kind == "score": options, labels = list(q["criteria"]), None else: raise ValueError(f"unknown question type {kind!r}") if len(options) < 2: raise ValueError("a question needs at least 2 options") ids, pos = self._encode(q.get("instructions"), q.get("state"), options) encoded.append((ids, pos, QUESTION_TYPES[kind])) meta.append((kind, labels, options)) answers = [] for start in range(0, len(encoded), batch_size): logits, conf_logits = self._forward(self._batch(encoded[start:start + batch_size])) for j, (kind, labels, options) in enumerate(meta[start:start + batch_size]): z = logits[j, :len(options)].astype(np.float64) if self.temperature is not None: # one temperature per (question type, number of options 2..9) z = z / self.temperature[QUESTION_TYPES[kind], min(max(len(options), 2), 9) - 2] p = np.exp(z - z.max()); p /= p.sum() c = float(conf_logits[j]) if self.platt is not None: c = self.platt[0] * c + self.platt[1] confidence = 1.0 / (1.0 + math.exp(-max(min(c, 30.0), -30.0))) if kind == "noul": answers.append({"type": "noul", "noul": float(p[1])}) elif kind == "choice": best = int(p.argmax()) answers.append({"type": "choice", "choice": labels[best], "confidence": confidence, "probabilities": {l: float(x) for l, x in zip(labels, p)}}) else: answers.append({"type": "score", "score": float(sum(i * x for i, x in enumerate(p))), "confidence": confidence, "probabilities": {str(i): float(x) for i, x in enumerate(p)}}) return answers def choice(self, state: Any, criteria: dict[str, Any], instructions: Any = None) -> dict[str, Any]: return self.predict([{"type": "choice", "state": state, "criteria": criteria, "instructions": instructions}])[0] def noul(self, state: Any, instructions: Any, true: Any = None, false: Any = None) -> dict[str, Any]: return self.predict([{"type": "noul", "state": state, "instructions": instructions, "criteria": {"true": true, "false": false}}])[0] def score(self, state: Any, criteria: list[Any], instructions: Any = None) -> dict[str, Any]: return self.predict([{"type": "score", "state": state, "criteria": criteria, "instructions": instructions}])[0] def _confidence_features(probs, mask, qtype): """Shape of the option distribution: top-1, margin, normalized entropy, 1/k and the question type.""" import torch k = mask.sum(-1).clamp(min=2).float() top2 = torch.where(mask, probs, torch.zeros_like(probs)).sort(dim=-1, descending=True).values[:, :2] entropy = -(probs * probs.clamp(min=1e-9).log()).sum(-1) / k.log() onehot = (qtype[:, None] == torch.arange(len(QUESTION_TYPES), device=qtype.device)[None, :]).float() return torch.cat([top2[:, :1], top2[:, :1] - top2[:, 1:2], entropy[:, None], (1.0 / k)[:, None], onehot], dim=-1) class Dinah(_Base): """PyTorch backend (fp32 weights, model.safetensors). On CUDA it runs under bfloat16 autocast by default.""" FILES = ["config.json", "model.safetensors", "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json"] def __init__(self, directory: Path, device: str | None = None, dtype: str | None = None): super().__init__(directory) import torch import torch.nn as nn from safetensors.torch import load_file from transformers import ModernBertConfig, ModernBertModel cfg = ModernBertConfig(**{k: v for k, v in self.config.items() if k not in ("dinah", "architectures")}) hidden = cfg.hidden_size net = nn.Module() net.encoder = ModernBertModel(cfg) net.type_embedding = nn.Embedding(len(QUESTION_TYPES), hidden) net.scorer = nn.Sequential(nn.LayerNorm(hidden), nn.Dropout(0.0), nn.Linear(hidden, hidden), nn.GELU(), nn.Linear(hidden, 1)) net.confidence = nn.Module() net.confidence.net = nn.Sequential(nn.Linear(7, 64), nn.GELU(), nn.Linear(64, 64), nn.GELU(), nn.Linear(64, 1)) net.load_state_dict(load_file(str(directory / "model.safetensors")), strict=True) self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") self.dtype = dtype or ("bfloat16" if self.device == "cuda" else "float32") self.net = net.to(self.device).eval() self.torch = torch @classmethod def from_pretrained(cls, path_or_repo: str | Path = "Lukitaduarte/dinah-0-apache", device: str | None = None, dtype: str | None = None, revision: str | None = None) -> "Dinah": return cls(_resolve(path_or_repo, cls.FILES, revision), device, dtype) def _forward(self, batch): torch = self.torch t = {k: torch.as_tensor(v, device=self.device) for k, v in batch.items()} autocast = self.dtype != "float32" and self.device.startswith("cuda") with torch.no_grad(), torch.autocast("cuda", dtype=getattr(torch, self.dtype), enabled=autocast): h = self.net.encoder(input_ids=t["ids"], attention_mask=t["pad_mask"]).last_hidden_state h = h + self.net.type_embedding(t["qtype"])[:, None, :] idx = t["opt_positions"].unsqueeze(-1).expand(-1, -1, h.size(-1)) logits = self.net.scorer(h.gather(1, idx)).squeeze(-1).float() logits = torch.where(t["opt_mask"], logits, torch.full_like(logits, NEG)) probs = logits.log_softmax(-1).exp() * t["opt_mask"] conf = self.net.confidence.net(_confidence_features(probs, t["opt_mask"], t["qtype"])).squeeze(-1) return logits.float().cpu().numpy(), conf.float().cpu().numpy() class DinahONNX(_Base): """onnxruntime backend (int8 weights, block-local attention). Made for CPU.""" FILES = ["config.json", "onnx/model_int8.onnx", "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json"] def __init__(self, directory: Path, threads: int | None = None): super().__init__(directory) import onnxruntime as ort options = ort.SessionOptions() if threads: options.intra_op_num_threads = threads self.session = ort.InferenceSession(str(directory / "onnx/model_int8.onnx"), options, providers=["CPUExecutionProvider"]) @classmethod def from_pretrained(cls, path_or_repo: str | Path = "Lukitaduarte/dinah-0-apache", threads: int | None = None, revision: str | None = None) -> "DinahONNX": return cls(_resolve(path_or_repo, cls.FILES, revision), threads) def _forward(self, batch): logits, conf = self.session.run(["logits", "conf_logit"], batch) return logits, conf if __name__ == "__main__": model = DinahONNX.from_pretrained(".") print(model.choice(state="The package arrived broken and I need it for tomorrow.", criteria={"refund": "The customer wants their money back", "replacement": "The customer wants a new unit sent", "complaint": "The customer only wants to complain"}, instructions="What does the customer want?"))