needle_mlx.py

"""Small MLX inference implementation for Cactus Needle."""

from __future__ import annotations

import json
import math
from pathlib import Path
from typing import Any

import mlx.core as mx


class NeedleModel:
    def __init__(self, weights: dict[str, mx.array], config: dict[str, Any]):
        self.weights = weights
        self.config = config
        self.num_heads = int(config["num_heads"])
        self.num_kv_heads = int(config["num_kv_heads"])
        self.d_model = int(config["d_model"])
        self.head_dim = self.d_model // self.num_heads

    @classmethod
    def from_pretrained(cls, path: str | Path) -> "NeedleModel":
        path = Path(path)
        config = json.loads((path / "config.json").read_text())
        weights = mx.load(str(path / "model.safetensors"))
        return cls(weights, config)

    def _weight(self, name: str, layer: int | None = None) -> mx.array:
        value = self.weights[name]
        return value if layer is None else value[layer]

    @staticmethod
    def _linear(x: mx.array, kernel: mx.array, bias: mx.array | None = None) -> mx.array:
        output = mx.matmul(x, kernel)
        return output if bias is None else output + bias

    @staticmethod
    def _zcrms_norm(x: mx.array, scale: mx.array) -> mx.array:
        rms = mx.sqrt(mx.mean(mx.square(x.astype(mx.float32)), axis=-1, keepdims=True) + 1e-6)
        return ((1.0 + scale) * x / rms).astype(x.dtype)

    def _rope(self, x: mx.array) -> mx.array:
        sequence_length = x.shape[2]
        half_dim = self.head_dim // 2
        positions = mx.arange(sequence_length, dtype=mx.float32)
        inv_freq = 1.0 / (
            float(self.config["rope_theta"])
            ** (mx.arange(0, self.head_dim, 2, dtype=mx.float32) / self.head_dim)
        )
        angles = mx.outer(positions, inv_freq)
        cos = mx.reshape(mx.cos(angles), (1, 1, sequence_length, half_dim))
        sin = mx.reshape(mx.sin(angles), (1, 1, sequence_length, half_dim))
        first, second = x[..., :half_dim], x[..., half_dim:]
        return mx.concatenate([first * cos - second * sin, second * cos + first * sin], axis=-1)

    def _attention(
        self,
        q_input: mx.array,
        kv_input: mx.array,
        prefix: str,
        layer: int,
        mask: mx.array | None,
        rope: bool,
    ) -> mx.array:
        q = self._linear(q_input, self._weight(f"{prefix}/q_proj/kernel", layer))
        k = self._linear(kv_input, self._weight(f"{prefix}/k_proj/kernel", layer))
        v = self._linear(kv_input, self._weight(f"{prefix}/v_proj/kernel", layer))
        batch = q.shape[0]
        q = mx.transpose(mx.reshape(q, (batch, -1, self.num_heads, self.head_dim)), (0, 2, 1, 3))
        k = mx.transpose(mx.reshape(k, (batch, -1, self.num_kv_heads, self.head_dim)), (0, 2, 1, 3))
        v = mx.transpose(mx.reshape(v, (batch, -1, self.num_kv_heads, self.head_dim)), (0, 2, 1, 3))
        q = self._zcrms_norm(q, self._weight(f"{prefix}/q_norm/scale", layer))
        k = self._zcrms_norm(k, self._weight(f"{prefix}/k_norm/scale", layer))
        if self.num_heads != self.num_kv_heads:
            repeats = self.num_heads // self.num_kv_heads
            k = mx.repeat(k, repeats, axis=1)
            v = mx.repeat(v, repeats, axis=1)
        if rope:
            q, k = self._rope(q), self._rope(k)
        scores = mx.matmul(q.astype(mx.float32), mx.transpose(k.astype(mx.float32), (0, 1, 3, 2)))
        scores = scores / math.sqrt(self.head_dim)
        if mask is not None:
            scores = mx.where(mask, scores, mx.array(-1e9, dtype=mx.float32))
        output = mx.matmul(mx.softmax(scores, axis=-1).astype(v.dtype), v)
        output = mx.reshape(mx.transpose(output, (0, 2, 1, 3)), (batch, -1, self.d_model))
        return self._linear(output, self._weight(f"{prefix}/out_proj/kernel", layer))

    def encode(self, source_tokens: mx.array) -> tuple[mx.array, mx.array]:
        source_tokens = mx.array(source_tokens, dtype=mx.int32)
        padding_mask = source_tokens != int(self.config["pad_token_id"])
        mask = mx.reshape(padding_mask, (source_tokens.shape[0], 1, 1, source_tokens.shape[1]))
        x = self.weights["embedding/embedding"][source_tokens] * math.sqrt(self.d_model)
        for layer in range(int(self.config["num_encoder_layers"])):
            prefix = "encoder/layers/EncoderBlock_0"
            residual = x
            x = self._zcrms_norm(x, self._weight(f"{prefix}/ZCRMSNorm_0/scale", layer))
            x = self._attention(x, x, f"{prefix}/self_attn", layer, mask, rope=True)
            gate = mx.sigmoid(self._weight(f"{prefix}/attn_gate", layer))
            x = residual + gate * x
        return self._zcrms_norm(x, self.weights["encoder/final_norm/scale"]), mask

    def decode(self, target_tokens: mx.array, encoder_out: mx.array, cross_mask: mx.array) -> mx.array:
        target_tokens = mx.array(target_tokens, dtype=mx.int32)
        length = target_tokens.shape[1]
        causal = mx.tril(mx.ones((length, length), dtype=mx.bool_))
        target_valid = target_tokens != int(self.config["pad_token_id"])
        self_mask = mx.reshape(causal, (1, 1, length, length)) & mx.reshape(
            target_valid, (target_tokens.shape[0], 1, 1, length)
        )
        x = self.weights["embedding/embedding"][target_tokens] * math.sqrt(self.d_model)
        prefix = "decoder/layers/DecoderBlock_0"
        for layer in range(int(self.config["num_decoder_layers"])):
            residual = x
            x = self._zcrms_norm(x, self._weight(f"{prefix}/ZCRMSNorm_0/scale", layer))
            x = self._attention(x, x, f"{prefix}/self_attn", layer, self_mask, rope=True)
            x = residual + mx.sigmoid(self._weight(f"{prefix}/self_attn_gate", layer)) * x
            residual = x
            x = self._zcrms_norm(x, self._weight(f"{prefix}/ZCRMSNorm_1/scale", layer))
            x = self._attention(x, encoder_out, f"{prefix}/cross_attn", layer, cross_mask, rope=False)
            x = residual + mx.sigmoid(self._weight(f"{prefix}/cross_attn_gate", layer)) * x
        x = self._zcrms_norm(x, self.weights["decoder/ZCRMSNorm_0/scale"])
        return mx.matmul(x.astype(mx.float32), mx.transpose(self.weights["embedding/embedding"].astype(mx.float32)))

    def forward(self, source_tokens: mx.array, target_tokens: mx.array) -> mx.array:
        encoder_out, cross_mask = self.encode(source_tokens)
        return self.decode(target_tokens, encoder_out, cross_mask)