anima_mlx/models/text_encoder.py

"""Qwen3 0.6B text encoder for Anima MLX parity."""

from __future__ import annotations

import math
from pathlib import Path
from typing import Any, Mapping

from .primitives import apply_rotary_pos_emb, linear, rms_norm, silu


class Qwen3Attention:
    def __init__(
        self,
        weights: Mapping[str, Any],
        prefix: str,
        *,
        hidden_size: int = 1024,
        num_heads: int = 16,
        num_key_value_heads: int = 8,
        head_dim: int = 128,
    ) -> None:
        self.weights = weights
        self.prefix = prefix
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.num_key_value_heads = num_key_value_heads
        self.head_dim = head_dim
        self.num_key_value_groups = num_heads // num_key_value_heads

    def _weight(self, name: str) -> Any:
        return self.weights[f"{self.prefix}.{name}"]

    def _linear(self, name: str, x: Any) -> Any:
        return linear(x, self._weight(f"{name}.weight"))

    def __call__(self, x: Any, cos: Any, sin: Any, attention_mask: Any | None = None) -> Any:
        import mlx.core as mx

        batch, seq_len, _ = x.shape
        q = self._linear("q_proj", x).reshape(batch, seq_len, self.num_heads, self.head_dim)
        k = self._linear("k_proj", x).reshape(batch, seq_len, self.num_key_value_heads, self.head_dim)
        v = self._linear("v_proj", x).reshape(batch, seq_len, self.num_key_value_heads, self.head_dim)

        q = rms_norm(q, self._weight("q_norm.weight"))
        k = rms_norm(k, self._weight("k_norm.weight"))

        q = mx.swapaxes(q, 1, 2)
        k = mx.swapaxes(k, 1, 2)
        v = mx.swapaxes(v, 1, 2)

        q = apply_rotary_pos_emb(q, cos, sin)
        k = apply_rotary_pos_emb(k, cos, sin)

        if self.num_key_value_groups != 1:
            k = mx.repeat(k, self.num_key_value_groups, axis=1)
            v = mx.repeat(v, self.num_key_value_groups, axis=1)

        scores = (q @ mx.swapaxes(k, -1, -2)) * (1.0 / math.sqrt(self.head_dim))
        scores = _apply_attention_mask(scores, attention_mask)
        probs = mx.softmax(scores, axis=-1)
        out = probs @ v
        out = mx.swapaxes(out, 1, 2).reshape(batch, seq_len, self.num_heads * self.head_dim)
        return self._linear("o_proj", out)


class Qwen3MLP:
    def __init__(self, weights: Mapping[str, Any], prefix: str) -> None:
        self.weights = weights
        self.prefix = prefix

    def _weight(self, name: str) -> Any:
        return self.weights[f"{self.prefix}.{name}"]

    def __call__(self, x: Any) -> Any:
        gate = linear(x, self._weight("gate_proj.weight"))
        up = linear(x, self._weight("up_proj.weight"))
        return linear(silu(gate) * up, self._weight("down_proj.weight"))


class Qwen3DecoderLayer:
    def __init__(self, weights: Mapping[str, Any], index: int) -> None:
        self.weights = weights
        self.prefix = f"layers.{index}"
        self.self_attn = Qwen3Attention(weights, f"{self.prefix}.self_attn")
        self.mlp = Qwen3MLP(weights, f"{self.prefix}.mlp")

    def _weight(self, name: str) -> Any:
        return self.weights[f"{self.prefix}.{name}"]

    def __call__(self, x: Any, cos: Any, sin: Any, attention_mask: Any | None = None) -> Any:
        residual = x
        hidden = rms_norm(x, self._weight("input_layernorm.weight"))
        x = residual + self.self_attn(hidden, cos, sin, attention_mask)

        residual = x
        hidden = rms_norm(x, self._weight("post_attention_layernorm.weight"))
        return residual + self.mlp(hidden)


class Qwen3TextEncoder:
    """MLX implementation of the Qwen3 0.6B text encoder used by Anima."""

    def __init__(
        self,
        weights: Mapping[str, Any],
        *,
        layer_count: int = 28,
        head_dim: int = 128,
        rope_theta: float = 1_000_000.0,
    ) -> None:
        self.weights = self._normalize_keys(weights)
        self.layer_count = layer_count
        self.head_dim = head_dim
        self.rope_theta = rope_theta
        self.layers = [Qwen3DecoderLayer(self.weights, index) for index in range(layer_count)]

    @classmethod
    def from_safetensors(cls, path: str | Path, *, dtype: str = "float32") -> "Qwen3TextEncoder":
        from anima_mlx.utils.weights import load_mlx_safetensors_subset

        path = _resolve_text_encoder_path(path)
        weights = load_mlx_safetensors_subset(path, prefix="model.", strip_prefix="model.", dtype=dtype)
        return cls(weights)

    @staticmethod
    def _normalize_keys(weights: Mapping[str, Any]) -> dict[str, Any]:
        normalized: dict[str, Any] = {}
        for key, value in weights.items():
            if key.startswith("model."):
                key = key.removeprefix("model.")
            normalized[key] = value
        return normalized

    def __call__(
        self,
        input_ids: Any,
        *,
        attention_mask: Any | None = None,
        output_hidden_states: bool = False,
    ) -> Any | tuple[Any, tuple[Any, ...]]:
        import mlx.core as mx

        x = mx.take(self.weights["embed_tokens.weight"], input_ids.astype(mx.int64), axis=0)
        cos, sin = self._position_embeddings(x, input_ids.shape[-1])
        hidden_states: list[Any] = []
        if output_hidden_states:
            hidden_states.append(x)
        for layer in self.layers:
            x = layer(x.astype(mx.float32), cos, sin, attention_mask)
            if output_hidden_states:
                hidden_states.append(x)
        x = rms_norm(x.astype(mx.float32), self.weights["norm.weight"])
        if output_hidden_states:
            hidden_states[-1] = x
            return x, tuple(hidden_states)
        return x

    def _position_embeddings(self, x: Any, seq_len: int) -> tuple[Any, Any]:
        import mlx.core as mx

        inv_freq = 1.0 / (self.rope_theta ** (mx.arange(0, self.head_dim, 2).astype(mx.float32) / self.head_dim))
        position_ids = mx.arange(seq_len).astype(mx.float32)[None, :]
        freqs = position_ids[..., None] * inv_freq
        emb = mx.concatenate([freqs, freqs], axis=-1)
        return mx.cos(emb).astype(x.dtype), mx.sin(emb).astype(x.dtype)


def _apply_attention_mask(scores: Any, attention_mask: Any | None) -> Any:
    import mlx.core as mx

    seq_len = scores.shape[-1]
    positions = mx.arange(seq_len)
    causal_mask = positions[None, :] > positions[:, None]
    min_value = mx.array(-1e9, dtype=scores.dtype)
    scores = mx.where(causal_mask[None, None, :, :], min_value, scores)
    if attention_mask is not None:
        padding_mask = attention_mask.astype(mx.bool_)
        scores = mx.where(padding_mask[:, None, None, :], scores, min_value)
    return scores


def _resolve_text_encoder_path(path: str | Path) -> Path:
    resolved = Path(path)
    if resolved.is_dir():
        return resolved / "text_encoder.safetensors"
    return resolved