mlx_model.py

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

try:
    import mlx.core as mx
    import mlx.nn as nn
except ImportError:  # Linux export validation intentionally has no MLX runtime.
    mx = None
    nn = None


if nn is not None and mx is not None:

    class RMSNorm(nn.Module):
        def __init__(self, dimension: int, epsilon: float = 1e-6) -> None:
            super().__init__()
            self.weight = mx.ones((dimension,))
            self.epsilon = epsilon

        def __call__(self, values: Any) -> Any:
            normalized = values * mx.rsqrt(
                mx.mean(mx.square(values), axis=-1, keepdims=True) + self.epsilon
            )
            return normalized * self.weight

    def _rotate_half(values: Any) -> Any:
        first, second = mx.split(values, 2, axis=-1)
        return mx.concatenate((-second, first), axis=-1)

    def apply_rope(query: Any, key: Any) -> tuple[Any, Any]:
        sequence_length = query.shape[-2]
        dimension = query.shape[-1]
        positions = mx.arange(sequence_length, dtype=mx.float32)
        frequencies = 1.0 / (10000 ** (mx.arange(0, dimension, 2, dtype=mx.float32) / dimension))
        angles = positions[:, None] * frequencies[None, :]
        angles = mx.concatenate((angles, angles), axis=-1)[None, None, :, :]
        cosine = mx.cos(angles)
        sine = mx.sin(angles)
        return query * cosine + _rotate_half(query) * sine, key * cosine + _rotate_half(key) * sine

    class CausalSelfAttention(nn.Module):
        def __init__(self, config: dict[str, Any]) -> None:
            super().__init__()
            self.heads = int(config["attention_heads"])
            self.head_dimension = int(config["model_width"]) // self.heads
            width = int(config["model_width"])
            self.qkv = nn.Linear(width, 3 * width, bias=False)
            self.output = nn.Linear(width, width, bias=False)

        def __call__(self, values: Any) -> Any:
            batch, sequence, width = values.shape
            qkv = self.qkv(values).reshape(batch, sequence, 3, self.heads, self.head_dimension)
            qkv = qkv.transpose(2, 0, 3, 1, 4)
            query, key, value = qkv[0], qkv[1], qkv[2]
            query, key = apply_rope(query, key)
            mask = nn.MultiHeadAttention.create_additive_causal_mask(sequence)
            attended = mx.fast.scaled_dot_product_attention(
                query,
                key,
                value,
                scale=self.head_dimension**-0.5,
                mask=mask,
            )
            return self.output(attended.transpose(0, 2, 1, 3).reshape(batch, sequence, width))

    class SwiGLU(nn.Module):
        def __init__(self, config: dict[str, Any]) -> None:
            super().__init__()
            width = int(config["model_width"])
            ffn_width = int(config["ffn_width"])
            self.gate = nn.Linear(width, ffn_width, bias=False)
            self.up = nn.Linear(width, ffn_width, bias=False)
            self.down = nn.Linear(ffn_width, width, bias=False)

        def __call__(self, values: Any) -> Any:
            return self.down(nn.silu(self.gate(values)) * self.up(values))

    class TransformerBlock(nn.Module):
        def __init__(self, config: dict[str, Any]) -> None:
            super().__init__()
            width = int(config["model_width"])
            self.attention_norm = RMSNorm(width)
            self.attention = CausalSelfAttention(config)
            self.ffn_norm = RMSNorm(width)
            self.ffn = SwiGLU(config)

        def __call__(self, values: Any) -> Any:
            values = values + self.attention(self.attention_norm(values))
            return values + self.ffn(self.ffn_norm(values))

    class TransformerRouterMLX(nn.Module):
        def __init__(self, config: dict[str, Any]) -> None:
            super().__init__()
            width = int(config["model_width"])
            self.config = config
            self.embedding = nn.Embedding(int(config["vocab_size"]), width)
            self.blocks = [TransformerBlock(config) for _ in range(int(config["layers"]))]
            self.norm = RMSNorm(width)
            self.pass_head = nn.Linear(width, int(config["candidate_count"]))
            self.score_head = nn.Linear(width, int(config["candidate_count"]))

        def __call__(self, input_ids: Any) -> dict[str, Any]:
            hidden = self.embedding(input_ids)
            for block in self.blocks:
                hidden = block(hidden)
            hidden = self.norm(hidden)
            route_state = hidden[:, -1]
            return {
                "pass_logits": self.pass_head(route_state),
                "scores": mx.sigmoid(self.score_head(route_state)),
            }

    def load_mlx_router(directory: str | Path) -> Any:
        root = Path(directory)
        config = json.loads((root / "config.json").read_text(encoding="utf-8"))["model"]
        model = TransformerRouterMLX(config)
        weights = mx.load(str(root / "model.safetensors"))
        model.load_weights(list(weights.items()), strict=True)
        mx.eval(model.parameters())
        return model

else:

    class TransformerRouterMLX:  # type: ignore[no-redef]
        def __init__(self, config: dict[str, Any]) -> None:
            raise RuntimeError("MLX requires macOS on Apple Silicon")

    def load_mlx_router(directory: str | Path) -> Any:
        raise RuntimeError("MLX requires macOS on Apple Silicon")