Add Dinah-0 Apache: Dinah-0 retrained without non-commercial data

3265ba1 verified 11 days ago

13.4 kB

"""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?"))

	"""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?"))