decision-toolkit 0.1.0__py3-none-any.whl

decision_toolkit/adapters/julia.py

@@ -0,0 +1,168 @@
+"""`decision_toolkit.adapters.julia` — Julia-1's single-graph layout.
+
+Contract (ported from the official `julia/data.py` sequence builder and the
+ONNX repository's browser runtime):
+
+    [CLS] <type> question: <instructions> [SEP] ([MASK] opt0) ([MASK] opt1) … [SEP] <state> [SEP]
+    one ONNX graph: input_ids, attention_mask, marker_pos, marker_mask, qtype -> logits
+    probabilities = softmax(logits); noul options are [false, true], so p(true) = p[1]
+
+Unlike Laya there is no separate head and no temperature: the exported graph
+scores the options directly.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Mapping
+from typing import Any
+
+from ..engine import OnnxEngine, encode
+from ..errors import AdapterError
+from .base import (
+    build_choice_answer,
+    build_noul_answer,
+    build_score_answer,
+    question_options,
+    softmax,
+    state_to_text,
+)
+
+logger = logging.getLogger(__name__)
+
+QTYPES = {"choice": 0, "score": 1, "noul": 2}
+
+
+class JuliaAdapter:
+    """System One over a Julia-1 ONNX export (one graph, marker/qtype inputs)."""
+
+    name = "julia"
+
+    def __init__(self, *, max_length: int = 1024, head_length: int = 256) -> None:
+        self.max_length = int(max_length)
+        self.head_length = int(head_length)
+
+    # ------------------------------------------------------------------ setup
+    @staticmethod
+    def _special_ids(tokenizer: Any) -> tuple[int, int, int]:
+        for mask_name, bos_name, eos_name in (
+            ("<mask>", "<bos>", "<eos>"),
+            ("[MASK]", "[CLS]", "[SEP]"),
+        ):
+            ids = (
+                tokenizer.token_to_id(mask_name),
+                tokenizer.token_to_id(bos_name),
+                tokenizer.token_to_id(eos_name),
+            )
+            if all(token is not None for token in ids):
+                return ids  # type: ignore[return-value]
+        raise AdapterError(
+            "the tokenizer exposes neither <mask>/<bos>/<eos> nor [MASK]/[CLS]/[SEP]; "
+            "this is not a Julia-1 tokenizer"
+        )
+
+    @staticmethod
+    def _render_options(question: Mapping[str, Any]) -> list[str]:
+        criteria = question.get("criteria")
+        kind = question.get("type")
+        if kind == "choice":
+            return [f"{key}: {value}" if value else key for key, value in criteria.items()]
+        if kind == "score":
+            return [str(level) for level in criteria]
+        criteria = criteria or {}
+        return [
+            "false: " + (criteria.get("false") or "false"),
+            "true: " + (criteria.get("true") or "true"),
+        ]
+
+    def _sequence(
+        self,
+        tokenizer: Any,
+        state_text: str,
+        question: Mapping[str, Any],
+        options: list[str],
+    ) -> tuple[list[int], list[int], bool]:
+        mask, bos, eos = self._special_ids(tokenizer)
+        truncated = False
+        instructions = str(question.get("instructions") or "")
+        head_ids = encode(tokenizer, f"{question['type']} question: {instructions}")
+        option_ids: list[list[int]] = []
+        for option in options:
+            encoded = encode(tokenizer, " " + option)
+            if len(encoded) > 48:
+                truncated = True
+            option_ids.append([mask, *encoded[:48]])
+
+        budget = self.head_length - sum(len(option) for option in option_ids)
+        if budget < 16:
+            truncated = True
+            per = max(4, (self.head_length - 16) // max(1, len(option_ids)))
+            option_ids = [option[:per] for option in option_ids]
+            budget = self.head_length - sum(len(option) for option in option_ids)
+        if len(head_ids) > max(8, budget):
+            truncated = True
+        head_ids = head_ids[: max(8, budget)]
+
+        ids = [bos, *head_ids, eos]
+        markers: list[int] = []
+        for option in option_ids:
+            markers.append(len(ids))
+            ids.extend(option)
+        ids.append(eos)
+
+        room = max(1, self.max_length - len(ids) - 1)
+        state_ids = encode(tokenizer, state_text)
+        if len(state_ids) > room:
+            truncated = True
+        state_ids = state_ids[:room]
+        all_ids = (ids + state_ids + [eos])[: self.max_length]
+        if len(markers) > sum(1 for marker in markers if marker < self.max_length):
+            truncated = True
+        return all_ids, [marker for marker in markers if marker < self.max_length], truncated
+
+    # -------------------------------------------------------------- inference
+    def system_one(
+        self, engine: OnnxEngine, state: Any, questions: Mapping[str, Mapping[str, Any]]
+    ) -> Mapping[str, Mapping[str, Any]]:
+        import numpy as np
+
+        tokenizer = engine.tokenizer
+        state_text = state_to_text(state)
+        answers: dict[str, dict[str, Any]] = {}
+
+        for name, question in questions.items():
+            kind, labels, _ = question_options(question)
+            options = self._render_options(question)
+            ids, markers, truncated = self._sequence(tokenizer, state_text, question, options)
+            if len(markers) != len(options):
+                raise AdapterError(
+                    f"{name!r}: the options do not fit in {self.head_length} tokens"
+                )
+
+            input_ids = np.asarray([ids], dtype=np.int64)
+            attention = np.ones_like(input_ids)
+            logits = engine.session("model").run(None, {
+                "input_ids": input_ids,
+                "attention_mask": attention,
+                "marker_pos": np.asarray([markers], dtype=np.int64),
+                "marker_mask": np.ones((1, len(markers)), dtype=bool),
+                "qtype": np.asarray([QTYPES[kind]], dtype=np.int64),
+            })[0]
+
+            probabilities = softmax(np.asarray(logits)[0][: len(markers)].tolist())
+            if kind == "choice":
+                answer = build_choice_answer(labels, probabilities)
+            elif kind == "noul":
+                answer = build_noul_answer(probabilities[1], probabilities[0])
+            else:
+                answer = build_score_answer(labels, probabilities)
+            if truncated:
+                answer["truncated"] = True
+                answer["warnings"] = [
+                    f"input was truncated to fit max_length={self.max_length}"
+                ]
+            answers[name] = answer
+        return answers
+
+
+__all__ = ["JuliaAdapter"]

decision_toolkit/adapters/laya.py

@@ -0,0 +1,225 @@
+"""`decision_toolkit.adapters.laya` — Laya's split encoder+head layout.
+
+Contract (ported from the reference implementation and validated against the
+mininvox export, where it reproduced the torch runtime's answers):
+
+    [CLS] <type> question: <instructions> [SEP] ([MASK] opt0) ([MASK] opt1) … [SEP] <state> [SEP]
+    encoder(input_ids, attention_mask) -> last_hidden_state
+    head(hidden_states, marker_pos, marker_mask, qtype, attention_mask)
+        -> logits per marker + act_logits
+    choice / score: softmax(logits / temperature)
+    noul: options are [false, true], so p(true) = p[1]
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+from collections.abc import Mapping
+from typing import Any
+
+from ..engine import OnnxEngine, encode
+from ..errors import AdapterError
+from .base import (
+    build_choice_answer,
+    build_noul_answer,
+    build_score_answer,
+    question_options,
+    softmax,
+    state_to_text,
+)
+
+logger = logging.getLogger(__name__)
+
+QTYPES = {"choice": 0, "score": 1, "noul": 2}
+QTYPE_NAMES = ["choice", "score", "noul"]
+
+
+class LayaAdapter:
+    """System One over a Laya split export (encoder + head + rl_agent_config)."""
+
+    name = "laya"
+
+    # ------------------------------------------------------------------ setup
+    @staticmethod
+    def _special_ids(tokenizer: Any) -> tuple[int, int, int]:
+        # Laya checkpoints ship two tokenizer serializations of the same contract:
+        #   multilingual: <mask>/<bos>/<eos>   ·   English: [MASK]/[CLS]/[SEP]
+        for mask_name, bos_name, eos_name in (
+            ("<mask>", "<bos>", "<eos>"),
+            ("[MASK]", "[CLS]", "[SEP]"),
+        ):
+            ids = (
+                tokenizer.token_to_id(mask_name),
+                tokenizer.token_to_id(bos_name),
+                tokenizer.token_to_id(eos_name),
+            )
+            if all(token is not None for token in ids):
+                return ids  # type: ignore[return-value]
+        raise AdapterError(
+            "the tokenizer exposes neither <mask>/<bos>/<eos> nor [MASK]/[CLS]/[SEP]; "
+            "this is not a Laya tokenizer"
+        )
+
+    @staticmethod
+    def _render_options(question: Mapping[str, Any]) -> list[str]:
+        criteria = question.get("criteria")
+        kind = question.get("type")
+        if kind == "choice":
+            return [f"{key}: {value}" if value else key for key, value in criteria.items()]
+        if kind == "score":
+            return [f"level {index}: {level}" for index, level in enumerate(criteria)]
+        criteria = criteria or {}
+        return [
+            "false: " + (criteria.get("false") or "no, the statement does not hold"),
+            "true: " + (criteria.get("true") or "yes, the statement holds"),
+        ]
+
+    def _sequence(
+        self,
+        tokenizer: Any,
+        state_text: str,
+        question: Mapping[str, Any],
+        options: list[str],
+        config: Mapping[str, Any],
+    ) -> tuple[list[int], list[int], bool]:
+        max_len = int(config.get("max_len", 1024))
+        head_max_len = int(config.get("head_max_len", 256))
+        mask, bos, eos = self._special_ids(tokenizer)
+        truncated = False
+
+        instructions = str(question.get("instructions") or "")
+        head_ids = encode(tokenizer, f"{question['type']} question: {instructions}")
+        option_ids: list[list[int]] = []
+        for option in options:
+            encoded = encode(tokenizer, " " + option)
+            if len(encoded) > 48:
+                truncated = True
+            option_ids.append([mask, *encoded[:48]])
+
+        budget = head_max_len - sum(len(option) for option in option_ids)
+        if budget < 16:
+            truncated = True
+            per = max(4, (head_max_len - 16) // max(1, len(option_ids)))
+            option_ids = [option[:per] for option in option_ids]
+            budget = head_max_len - sum(len(option) for option in option_ids)
+        if len(head_ids) > max(8, budget):
+            truncated = True
+        head_ids = head_ids[: max(8, budget)]
+
+        ids = [bos, *head_ids, eos]
+        markers: list[int] = []
+        for option in option_ids:
+            markers.append(len(ids))
+            ids.extend(option)
+        ids.append(eos)
+
+        room = max(0, max_len - len(ids) - 1)
+        state_ids = encode(tokenizer, state_text)
+        if len(state_ids) > room:
+            truncated = True
+        state_ids = state_ids[:room]
+        all_ids = (ids + state_ids + [eos])[:max_len]
+        if len(markers) > sum(1 for marker in markers if marker < max_len):
+            truncated = True
+        return all_ids, [marker for marker in markers if marker < max_len], truncated
+
+    @staticmethod
+    def _temperature(config: Mapping[str, Any], qtype: int, count: int) -> float:
+        size = "2" if count <= 2 else "3-5" if count <= 5 else "6-10" if count <= 10 else "11+"
+        bucket = f"{QTYPE_NAMES[qtype]}:{size}"
+        by_options = config.get("temperature_by_options") or {}
+        temperatures = config.get("temperature") or [1.0, 1.0, 1.0]
+        value = float(by_options.get(bucket, temperatures[qtype]))
+        # the torch runtime clamps temperatures outside [0.5, 5] before applying them
+        return min(5.0, max(0.5, value))
+
+    @staticmethod
+    def _confidence(probabilities: list[float]) -> float:
+        count = len(probabilities)
+        if count < 2:
+            return 1.0
+        entropy = -sum(p * math.log(max(p, 1e-12)) for p in probabilities)
+        return 1 - entropy / math.log(count)
+
+    # -------------------------------------------------------------- inference
+    def system_one(
+        self, engine: OnnxEngine, state: Any, questions: Mapping[str, Mapping[str, Any]]
+    ) -> Mapping[str, Mapping[str, Any]]:
+        import numpy as np
+
+        config = engine.read_json("config")
+        tokenizer = engine.tokenizer
+        state_text = state_to_text(state)
+        answers: dict[str, dict[str, Any]] = {}
+
+        for name, question in questions.items():
+            kind, labels, _ = question_options(question)
+            options = self._render_options(question)
+            ids, markers, truncated = self._sequence(
+                tokenizer, state_text, question, options, config
+            )
+            if len(markers) != len(options):
+                raise AdapterError(
+                    f"{name!r}: the options do not fit in {config.get('head_max_len')} tokens"
+                )
+            qtype = QTYPES[kind]
+
+            input_ids = np.asarray([ids], dtype=np.int64)
+            attention = np.ones_like(input_ids)
+            hidden = engine.session("encoder").run(
+                None, {"input_ids": input_ids, "attention_mask": attention}
+            )[0]
+
+            head = engine.session("head")
+            inputs = {item.name: item for item in head.get_inputs()}
+            hidden_name = next(
+                (name for name in ("hidden_states", "hidden", "last_hidden_state") if name in inputs),
+                None,
+            )
+            if hidden_name is None:
+                raise AdapterError(
+                    f"the head exposes no hidden-state input; inputs: {sorted(inputs)}"
+                )
+            qtype_value = (
+                np.asarray([qtype], dtype=np.int64)
+                if len(inputs["qtype"].shape) == 1
+                else np.asarray([[qtype]], dtype=np.int64)
+            )
+            logits, act_logits = head.run(None, {
+                hidden_name: hidden,
+                "attention_mask": attention,
+                "marker_pos": np.asarray([markers], dtype=np.int64),
+                "marker_mask": np.ones((1, len(markers)), dtype=bool),
+                "qtype": qtype_value,
+            })
+
+            raw = np.asarray(logits)[0][: len(markers)]
+            scaled = (raw / self._temperature(config, qtype, len(markers))).tolist()
+            probabilities = softmax(scaled)
+            act_probability = softmax(np.asarray(act_logits)[0].tolist())[0]
+
+            if kind == "choice":
+                answer = build_choice_answer(labels, probabilities)
+            elif kind == "noul":
+                # Laya's options are [false, true]: p[1] is the probability of true
+                answer = build_noul_answer(probabilities[1], probabilities[0])
+            else:
+                answer = build_score_answer(labels, probabilities)
+
+            if kind in {"choice", "score"}:
+                answer["answer_confidence"] = answer["confidence"]
+                answer["distribution_confidence"] = self._confidence(probabilities)
+            answer["confidence_kind"] = "probability"
+            answer["action"] = {"act_probability": float(act_probability)}
+            if truncated:
+                answer["truncated"] = True
+                answer["warnings"] = [
+                    f"input was truncated to fit max_len={config.get('max_len', 1024)}"
+                ]
+            answers[name] = answer
+
+        return answers
+
+
+__all__ = ["LayaAdapter"]

decision_toolkit/adapters/opendecider.py

@@ -0,0 +1,144 @@
+"""`decision_toolkit.adapters.opendecider` — OpenDecider nano's single-graph layout.
+
+Contract (from the project's `opendecider/prompt.py::nano_ids` and the export's
+`opendecider.json`):
+
+    [CLS] question: <instructions> [SEP] ([MASK] opt) … [SEP] input: <state> [SEP]
+    one ONNX graph: (input_ids, attention_mask) -> logits per token
+    the answer is softmax over the logits at the [MASK] positions
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Mapping
+from typing import Any
+
+from ..engine import OnnxEngine, encode
+from ..errors import AdapterError
+from .base import (
+    build_choice_answer,
+    build_noul_answer,
+    build_score_answer,
+    softmax,
+    state_to_text,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class OpendeciderAdapter:
+    """System One over the official OpenDecider nano ONNX build."""
+
+    name = "opendecider"
+
+    # ------------------------------------------------------------------ setup
+    @staticmethod
+    def _render_options(question: Mapping[str, Any]) -> dict[str, str | None]:
+        kind = question.get("type")
+        criteria = question.get("criteria")
+        if kind == "choice":
+            return {str(key): (None if value is None else str(value)) for key, value in criteria.items()}
+        if kind == "score":
+            return {str(level): None for level in criteria}
+        criteria = criteria or {}
+        return {
+            "false": None if criteria.get("false") is None else str(criteria.get("false")),
+            "true": None if criteria.get("true") is None else str(criteria.get("true")),
+        }
+
+    @staticmethod
+    def _special_ids(config: Mapping[str, Any]) -> tuple[int, int, int]:
+        ids = config.get("special_ids") or {}
+        try:
+            return int(ids["cls"]), int(ids["sep"]), int(ids["mask"])
+        except (KeyError, TypeError, ValueError) as error:
+            raise AdapterError(
+                "opendecider.json must declare special_ids.cls/sep/mask"
+            ) from error
+
+    def _sequence(
+        self,
+        tokenizer: Any,
+        state_text: str,
+        instructions: str,
+        options: Mapping[str, str | None],
+        special_ids: tuple[int, int, int],
+        max_len: int,
+    ) -> tuple[list[int], list[int], bool]:
+        cls_id, sep_id, mask_id = special_ids
+        truncated = False
+        ids = [cls_id, *encode(tokenizer, f"question: {instructions}"), sep_id]
+        markers: list[int] = []
+        for key, value in options.items():
+            markers.append(len(ids))
+            ids.append(mask_id)
+            text = f" {key}: {value}" if value and value != key else f" {key}"
+            ids.extend(encode(tokenizer, text))
+        ids.append(sep_id)
+
+        state_ids = encode(tokenizer, f"input: {state_text}")
+        room = max(max_len - len(ids) - 1, 0)
+        if len(state_ids) > room:
+            truncated = True
+        ids = ids + state_ids[:room] + [sep_id]
+        if len(ids) > max_len:
+            truncated = True
+        if len(markers) > sum(1 for marker in markers if marker < max_len):
+            truncated = True
+        return ids[:max_len], [marker for marker in markers if marker < max_len], truncated
+
+    def _scores(
+        self,
+        engine: OnnxEngine,
+        state_text: str,
+        instructions: str,
+        options: Mapping[str, str | None],
+        special_ids: tuple[int, int, int],
+        max_len: int,
+    ) -> tuple[list[float], bool]:
+        import numpy as np
+
+        ids, markers, truncated = self._sequence(
+            engine.tokenizer, state_text, instructions, options, special_ids, max_len
+        )
+        input_ids = np.asarray([ids], dtype=np.int64)
+        logits = engine.session("model").run(None, {
+            "input_ids": input_ids,
+            "attention_mask": np.ones_like(input_ids),
+        })[0][0]
+        return [float(logits[position]) for position in markers], truncated
+
+    # -------------------------------------------------------------- inference
+    def system_one(
+        self, engine: OnnxEngine, state: Any, questions: Mapping[str, Mapping[str, Any]]
+    ) -> Mapping[str, Mapping[str, Any]]:
+        config = engine.read_json("config")
+        special_ids = self._special_ids(config)
+        max_len = int(config.get("max_len", 2048))
+        state_text = state_to_text(state)
+        answers: dict[str, dict[str, Any]] = {}
+
+        for name, question in questions.items():
+            options = self._render_options(question)
+            scores, truncated = self._scores(
+                engine, state_text, str(question.get("instructions") or ""),
+                options, special_ids, max_len,
+            )
+            probabilities = softmax(scores)
+            kind = question.get("type")
+            labels = list(options)
+            if kind == "choice":
+                answer = build_choice_answer(labels, probabilities)
+            elif kind == "noul":
+                answer = build_noul_answer(probabilities[1], probabilities[0])
+            else:
+                answer = build_score_answer(labels, probabilities)
+            if truncated:
+                answer["truncated"] = True
+                answer["warnings"] = [f"input was truncated to fit max_len={max_len}"]
+            answers[name] = answer
+        return answers
+
+
+__all__ = ["OpendeciderAdapter"]

decision_toolkit/calibration.py

@@ -0,0 +1,79 @@
+"""`decision_toolkit.calibration` — do the probabilities mean what they say?
+
+Decision models answer with probabilities; whether those probabilities are
+*calibrated* (a 0.8 answer is right ~80% of the time) varies by model, by
+quantization and by domain. The toolkit keeps this small but first-class:
+measure before you threshold.
+
+A *pair* is ``(probability_given_to_the_chosen_answer, was_it_right)`` — build
+it from your own labelled rows with :func:`pairs_from_rows`.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Sequence
+
+Pair = tuple[float, bool]
+
+
+def pairs_from_rows(rows: Iterable[tuple[float, bool]]) -> list[Pair]:
+    """Normalize/validate ``(probability, correct)`` rows."""
+    pairs: list[Pair] = []
+    for probability, correct in rows:
+        p = float(probability)
+        if not 0.0 <= p <= 1.0:
+            raise ValueError(f"probability out of [0, 1]: {p}")
+        pairs.append((p, bool(correct)))
+    return pairs
+
+
+def brier_score(pairs: Sequence[Pair]) -> float:
+    """Mean squared error of the probabilities (lower is better; 0 is perfect)."""
+    if not pairs:
+        raise ValueError("brier_score needs at least one pair")
+    return sum((probability - float(correct)) ** 2 for probability, correct in pairs) / len(pairs)
+
+
+def expected_calibration_error(pairs: Sequence[Pair], n_bins: int = 10) -> float:
+    """Expected calibration error: |accuracy - confidence| averaged over bins."""
+    if not pairs:
+        raise ValueError("expected_calibration_error needs at least one pair")
+    bins: list[list[Pair]] = [[] for _ in range(n_bins)]
+    for probability, correct in pairs:
+        index = min(int(probability * n_bins), n_bins - 1)
+        bins[index].append((probability, correct))
+
+    total = len(pairs)
+    error = 0.0
+    for bucket in bins:
+        if not bucket:
+            continue
+        confidence = sum(p for p, _ in bucket) / len(bucket)
+        accuracy = sum(1.0 for _, correct in bucket if correct) / len(bucket)
+        error += (len(bucket) / total) * abs(accuracy - confidence)
+    return error
+
+
+def threshold_sweep(pairs: Sequence[Pair], thresholds: Sequence[float] = (0.5, 0.6, 0.7, 0.8, 0.9)
+                    ) -> dict[float, dict[str, float]]:
+    """Coverage and kept accuracy at each confidence threshold.
+
+    ``coverage`` is the share of items the model answers at or above the
+    threshold; ``accuracy_kept`` is how often those answers are right. An item
+    counts as right when its probability side (>= 0.5) matches its label —
+    pass ``correct=True/False`` directly to decide a different rule.
+    """
+    if not pairs:
+        raise ValueError("threshold_sweep needs at least one pair")
+    report: dict[float, dict[str, float]] = {}
+    for threshold in thresholds:
+        kept = [(p, c) for p, c in pairs if p >= threshold]
+        report[float(threshold)] = {
+            "coverage": len(kept) / len(pairs),
+            "accuracy_kept": (sum(1.0 for _, c in kept) / len(kept)) if kept else 0.0,
+            "n_kept": float(len(kept)),
+        }
+    return report
+
+
+__all__ = ["Pair", "brier_score", "expected_calibration_error", "pairs_from_rows", "threshold_sweep"]