mlx-decision 0.3.0__py3-none-any.whl

mlx_decision/__init__.py

@@ -0,0 +1,32 @@
+from .errors import DecisionError
+from .model import DecisionModel, load
+from .types import (
+    Choice,
+    ChoiceAnswer,
+    Noul,
+    NoulAnswer,
+    Request,
+    Response,
+    Result,
+    Score,
+    ScoreAnswer,
+    Usage,
+)
+
+__version__ = "0.3.0"
+
+__all__ = [
+    "Choice",
+    "ChoiceAnswer",
+    "DecisionError",
+    "DecisionModel",
+    "Noul",
+    "NoulAnswer",
+    "Request",
+    "Response",
+    "Result",
+    "Score",
+    "ScoreAnswer",
+    "Usage",
+    "load",
+]

mlx_decision/answers.py

@@ -0,0 +1,50 @@
+"""Turn per-option probabilities into answers. Shared by every backend."""
+
+from collections.abc import Mapping, Sequence
+
+from .types import Answer, Choice, ChoiceAnswer, Noul, NoulAnswer, Score, ScoreAnswer
+
+
+def choice_confidence(probabilities: Sequence[float]) -> float:
+    """How far the top probability sits above an even split, from 0 to 1."""
+    n = len(probabilities)
+    if n < 2:
+        return 1.0
+    return max(0.0, (max(probabilities) - 1 / n) / (1 - 1 / n))
+
+
+def score_confidence(probabilities: Sequence[float]) -> float:
+    """One minus the spread around the most likely level, relative to an even spread."""
+    n = len(probabilities)
+    if n < 2:
+        return 1.0
+    peak = max(range(n), key=probabilities.__getitem__)
+    spread = sum(p * abs(i - peak) for i, p in enumerate(probabilities))
+    even_spread = sum(abs(i - (n - 1) / 2) for i in range(n)) / n
+    return max(0.0, 1 - spread / even_spread)
+
+
+def build_answer(question: Noul | Choice | Score, probabilities: Mapping[str, float]) -> Answer:
+    """Build the answer for one question from its option probabilities.
+
+    ``probabilities`` is keyed by option id: ``"true"``/``"false"`` for a noul,
+    the criteria keys for a choice, ``"0"``, ``"1"``, ... for a score.
+    """
+    if isinstance(question, Noul):
+        return NoulAnswer(noul=probabilities["true"])
+    if isinstance(question, Choice):
+        options = list(question.criteria)
+        ordered = {option: probabilities[option] for option in options}
+        return ChoiceAnswer(
+            choice=max(options, key=ordered.__getitem__),
+            confidence=choice_confidence(list(ordered.values())),
+            probabilities=ordered,
+        )
+    levels = [str(index) for index in range(len(question.criteria))]
+    values = [probabilities[level] for level in levels]
+    return ScoreAnswer(
+        score=sum(index * p for index, p in enumerate(values)),
+        confidence=score_confidence(values),
+        legend=dict(zip(levels, question.criteria, strict=True)),
+        probabilities=dict(zip(levels, values, strict=True)),
+    )

mlx_decision/backbones/qwen3_5/__init__.py

@@ -0,0 +1,3 @@
+from .qwen3_5 import Model, ModelArgs
+
+__all__ = ["Model", "ModelArgs"]

mlx_decision/backbones/qwen3_5/activations.py

@@ -0,0 +1,74 @@
+# Copyright © 2023 Apple Inc.
+#
+# Copied from mlx-lm (https://github.com/ml-explore/mlx-lm), file
+# mlx_lm/models/activations.py at commit 5cfec4cb39deba54210b3ff4d86f2337c7bc10b5,
+# under the MIT License. See LICENSES/mlx-lm-MIT.txt.
+
+from functools import partial
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+@partial(mx.compile, shapeless=True)
+def swiglu(gate, x):
+    return nn.silu(gate) * x
+
+
+@partial(mx.compile, shapeless=True)
+def precise_swiglu(h, gate, x):
+    gate = nn.silu(gate.astype(mx.float32))
+    x = x.astype(mx.float32)
+    return (gate * x).astype(h.dtype)
+
+
+@partial(mx.compile, shapeless=True)
+def swiglu_oai(
+    gate: mx.array, x: mx.array, alpha: float = 1.702, limit: float = 7.0
+) -> mx.array:
+    # GPT-OSS variant: clip the gate from above only, and add 1 to the linear part.
+    gate = mx.minimum(gate, limit)
+    x = mx.clip(x, -limit, limit)
+    return (x + 1.0) * (gate * mx.sigmoid(alpha * gate))
+
+
+class SwigluOAI(nn.Module):
+    def __init__(self, alpha: float = 1.702, limit: float = 7.0):
+        super().__init__()
+        self._alpha = alpha
+        self._limit = limit
+
+    def __call__(self, x: mx.array, gate: mx.array) -> mx.array:
+        return swiglu_oai(gate, x, alpha=self._alpha, limit=self._limit)
+
+
+@partial(mx.compile, shapeless=True)
+def xielu(x, alpha_p, alpha_n, beta, eps):
+    alpha_p = nn.softplus(alpha_p)
+    alpha_n = beta + nn.softplus(alpha_n)
+    return mx.where(
+        x > 0,
+        alpha_p * mx.square(x) + beta * x,
+        (mx.expm1(mx.minimum(x, eps)) - x) * alpha_n + beta * x,
+    )
+
+
+class XieLU(nn.Module):
+    def __init__(
+        self,
+        alpha_p_init=0.8,
+        alpha_n_init=0.8,
+        beta=0.5,
+        eps=-1e-6,
+    ):
+        super().__init__()
+        alpha_p_tensor = mx.array(alpha_p_init)
+        alpha_n_tensor = mx.array(alpha_n_init - beta)
+        self.alpha_p = mx.log(mx.exp(alpha_p_tensor) - 1)
+        self.alpha_n = mx.log(mx.exp(alpha_n_tensor) - 1)
+
+        self.beta = mx.array(beta)
+        self.eps = mx.array(eps)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return xielu(x, self.alpha_p, self.alpha_n, self.beta, self.eps)

mlx_decision/backbones/qwen3_5/base.py

@@ -0,0 +1,143 @@
+# Copyright © 2023 Apple Inc.
+#
+# Copied from mlx-lm (https://github.com/ml-explore/mlx-lm), file
+# mlx_lm/models/base.py at commit 5cfec4cb39deba54210b3ff4d86f2337c7bc10b5,
+# under the MIT License. See LICENSES/mlx-lm-MIT.txt.
+
+import inspect
+from dataclasses import dataclass
+from typing import Optional
+
+import mlx.core as mx
+from mlx.utils import tree_map
+
+
+@dataclass
+class BaseModelArgs:
+    @classmethod
+    def from_dict(cls, params):
+        return cls(
+            **{
+                k: v
+                for k, v in params.items()
+                if k in inspect.signature(cls).parameters
+            }
+        )
+
+
+def create_causal_mask(
+    N: int,
+    offset: int = 0,
+    window_size: Optional[int] = None,
+    right_padding: Optional[mx.array] = None,
+    left_padding: Optional[mx.array] = None,
+):
+    rinds = mx.arange(offset + N)
+    linds = mx.arange(offset, offset + N) if offset else rinds
+    linds = linds[:, None]
+    rinds = rinds[None]
+    mask = linds >= rinds
+    if window_size is not None:
+        mask = mask & (linds < rinds + window_size)
+    if right_padding is not None:
+        mask = mask & (rinds < mx.expand_dims((offset + N) - right_padding, (1, 2, 3)))
+    if left_padding is not None:
+        mask = mask & (mx.expand_dims(left_padding, (1, 2, 3)) <= rinds)
+    return mask
+
+
+def create_attention_mask(
+    h, cache=None, window_size: Optional[int] = None, return_array: bool = False
+):
+    N = h.shape[1]
+    if cache and hasattr(cache, "make_mask"):
+        return cache.make_mask(N, return_array=return_array, window_size=window_size)
+    if N == 1:
+        return None
+    if return_array or (window_size and N > window_size):
+        return create_causal_mask(N, window_size=window_size)
+    return "causal"
+
+
+def create_ssm_mask(h, cache=None):
+    if cache and hasattr(cache, "make_mask"):
+        return cache.make_mask(h.shape[1])
+    return None
+
+
+def quantized_scaled_dot_product_attention(
+    queries: mx.array,
+    q_keys: tuple[mx.array, mx.array, mx.array],
+    q_values: tuple[mx.array, mx.array, mx.array],
+    scale: float,
+    mask: Optional[mx.array],
+    group_size: int = 64,
+    bits: int = 8,
+) -> mx.array:
+    B, n_q_heads, L, D = queries.shape
+    n_kv_heads = q_keys[0].shape[-3]
+    n_repeats = n_q_heads // n_kv_heads
+
+    queries *= scale
+
+    if n_repeats > 1:
+        queries = mx.reshape(queries, (B, n_kv_heads, n_repeats, L, D))
+        q_keys = tree_map(lambda x: mx.expand_dims(x, axis=-3), q_keys)
+        q_values = tree_map(lambda x: mx.expand_dims(x, axis=-3), q_values)
+
+    scores = mx.quantized_matmul(
+        queries, *q_keys, transpose=True, group_size=group_size, bits=bits
+    )
+    if mask is not None:
+        if isinstance(mask, str):
+            qL, kL = scores.shape[-2:]
+            q_indices = mx.arange(kL - qL, kL)
+            k_indices = mx.arange(kL)
+            mask = q_indices[:, None] >= k_indices[None]
+        if n_repeats > 1 and mask.ndim > 3:
+            mask = mx.expand_dims(mask, -3)
+        if mask.dtype == mx.bool_:
+            scores = mx.where(mask, scores, mx.finfo(scores.dtype).min)
+        else:
+            scores += mask
+    scores = mx.softmax(scores, axis=-1, precise=True)
+    out = mx.quantized_matmul(
+        scores, *q_values, transpose=False, group_size=group_size, bits=bits
+    )
+
+    if n_repeats > 1:
+        out = mx.reshape(out, (B, n_q_heads, L, D))
+
+    return out
+
+
+def scaled_dot_product_attention(
+    queries,
+    keys,
+    values,
+    cache,
+    scale: float,
+    mask: Optional[mx.array],
+    sinks: Optional[mx.array] = None,
+) -> mx.array:
+    if hasattr(cache, "bits"):
+        if sinks is not None:
+            raise ValueError("Quantized SDPA does not support attention sinks.")
+        return quantized_scaled_dot_product_attention(
+            queries,
+            keys,
+            values,
+            scale=scale,
+            mask=mask,
+            group_size=cache.group_size,
+            bits=cache.bits,
+        )
+    else:
+        return mx.fast.scaled_dot_product_attention(
+            queries,
+            keys,
+            values,
+            scale=scale,
+            mask=mask,
+            sinks=sinks,
+        )