"""Language model decoder for Jina VLM in MLX.""" from typing import List, Optional import mlx.core as mx import mlx.nn as nn from mlx_lm.models.base import create_attention_mask, scaled_dot_product_attention from mlx_lm.models.cache import KVCache from ..base import LanguageModelOutput from .config import TextConfig class RMSNorm(nn.Module): """RMS Layer Normalization.""" def __init__(self, dims: int, eps: float = 1e-6): super().__init__() self.eps = eps self.weight = mx.ones((dims,)) def __call__(self, x: mx.array) -> mx.array: rms = mx.sqrt(mx.mean(x * x, axis=-1, keepdims=True) + self.eps) return self.weight * (x / rms) class RoPE(nn.Module): """Rotary Positional Embeddings.""" def __init__(self, dims: int, theta: float = 1000000.0): super().__init__() self.dims = dims self.theta = theta inv_freq = 1.0 / (theta ** (mx.arange(0, dims, 2).astype(mx.float32) / dims)) self._inv_freq = inv_freq def __call__(self, x: mx.array, offset: int = 0) -> mx.array: seq_len = x.shape[2] # Handle array-valued offset from BatchKVCache if isinstance(offset, mx.array): offset = offset.max().item() positions = mx.arange(offset, offset + seq_len).astype(mx.float32) freqs = positions[:, None] * self._inv_freq[None, :] emb = mx.concatenate([freqs, freqs], axis=-1) cos = mx.cos(emb)[None, None, :, :] sin = mx.sin(emb)[None, None, :, :] x1 = x[..., : self.dims // 2] x2 = x[..., self.dims // 2 :] rotated = mx.concatenate([-x2, x1], axis=-1) return (x * cos + rotated * sin).astype(x.dtype) class Attention(nn.Module): """Multi-head attention with GQA and RoPE - matches weight naming: attn.qkv, attn.out""" def __init__(self, config: TextConfig): super().__init__() self.config = config self.num_heads = config.num_attention_heads self.num_kv_heads = config.num_key_value_heads self.head_dim = config.head_dim self.num_kv_groups = self.num_heads // self.num_kv_heads self.scale = self.head_dim**-0.5 # Fused QKV projection - named to match weights qkv_size = ( config.num_attention_heads + 2 * config.num_key_value_heads ) * config.head_dim self.qkv = nn.Linear(config.hidden_size, qkv_size, bias=False) self.out = nn.Linear( config.num_attention_heads * config.head_dim, config.hidden_size, bias=False ) # QK normalization - named to match weights if config.use_qk_norm: self.q_norm = RMSNorm(config.head_dim, eps=config.rms_norm_eps) self.k_norm = RMSNorm(config.head_dim, eps=config.rms_norm_eps) else: self.q_norm = None self.k_norm = None self.rope = RoPE(config.head_dim, theta=config.rope_theta) def __call__( self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[KVCache] = None, ) -> mx.array: B, L, _ = x.shape # Compute fused QKV qkv = self.qkv(x) q_size = self.num_heads * self.head_dim kv_size = self.num_kv_heads * self.head_dim q = qkv[..., :q_size] k = qkv[..., q_size : q_size + kv_size] v = qkv[..., q_size + kv_size :] q = q.reshape(B, L, self.num_heads, self.head_dim).transpose(0, 2, 1, 3) k = k.reshape(B, L, self.num_kv_heads, self.head_dim).transpose(0, 2, 1, 3) v = v.reshape(B, L, self.num_kv_heads, self.head_dim).transpose(0, 2, 1, 3) if self.q_norm is not None: q = self.q_norm(q) k = self.k_norm(k) if cache is not None: q = self.rope(q, offset=cache.offset) k = self.rope(k, offset=cache.offset) k, v = cache.update_and_fetch(k, v) else: q = self.rope(q) k = self.rope(k) output = scaled_dot_product_attention( q, k, v, cache=cache, scale=self.scale, mask=mask ) output = output.transpose(0, 2, 1, 3).reshape(B, L, -1) return self.out(output) class MLP(nn.Module): """MLP with SwiGLU - matches weight naming: ffn.gate_up, ffn.down""" def __init__(self, config: TextConfig): super().__init__() # Fused gate and up projection - named to match weights self.gate_up = nn.Linear( config.hidden_size, 2 * config.intermediate_size, bias=False ) self.down = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) def __call__(self, x: mx.array) -> mx.array: gate_up = self.gate_up(x) # Jina VLM convention: first half is value, second half is gate (activated) up, gate = mx.split(gate_up, 2, axis=-1) return self.down(nn.silu(gate) * up) class TransformerBlock(nn.Module): """Transformer block - matches weight naming: attn_norm, ffn_norm""" def __init__(self, config: TextConfig, layer_idx: int = 0): super().__init__() # Named to match weights self.attn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.attn = Attention(config) self.ffn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.ffn = MLP(config) def __call__( self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[KVCache] = None, ) -> mx.array: h = self.attn(self.attn_norm(x), mask=mask, cache=cache) x = x + h x = x + self.ffn(self.ffn_norm(x)) return x class ExtendedEmbedding(nn.Module): """Embedding with additional tokens - matches weight naming: embedding, new_embedding""" def __init__(self, vocab_size: int, additional_size: int, dims: int): super().__init__() self.vocab_size = vocab_size self.additional_size = additional_size # Named to match weights self.embedding = mx.zeros((vocab_size, dims)) self.new_embedding = mx.zeros((additional_size, dims)) def __call__(self, x: mx.array) -> mx.array: full_embedding = mx.concatenate([self.embedding, self.new_embedding], axis=0) return full_embedding[x] class TextModel(nn.Module): """Text decoder model - matches weight naming structure""" def __init__(self, config: TextConfig): super().__init__() self.config = config # Named to match weights: language_model.embedding if config.additional_vocab_size > 0: self.embedding = ExtendedEmbedding( config.vocab_size, config.additional_vocab_size, config.hidden_size ) else: self.embedding = nn.Embedding(config.vocab_size, config.hidden_size) self.layers = [ TransformerBlock(config, layer_idx=i) for i in range(config.num_hidden_layers) ] # Named to match weights: language_model.ln_f self.ln_f = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def __call__( self, input_ids: mx.array, inputs_embeds: Optional[mx.array] = None, mask: Optional[mx.array] = None, cache: Optional[List[KVCache]] = None, ) -> mx.array: if inputs_embeds is None: x = self.embedding(input_ids) else: x = inputs_embeds for i, layer in enumerate(self.layers): layer_cache = cache[i] if cache is not None else None x = layer(x, mask=mask, cache=layer_cache) return self.ln_f(x) class LanguageModel(nn.Module): """Language model wrapper - the TextModel is accessed as language_model in weights""" def __init__(self, config: TextConfig): super().__init__() self.config = config self.model_type = config.model_type # This will be loaded under "language_model" prefix self.embedding = None # Handled by sanitize self.layers = None # Handled by sanitize self.ln_f = None # Handled by sanitize # Build the actual model components directly here # They'll be found via language_model.embedding, language_model.layers, etc. if config.additional_vocab_size > 0: self.embedding = ExtendedEmbedding( config.vocab_size, config.additional_vocab_size, config.hidden_size ) else: self.embedding = nn.Embedding(config.vocab_size, config.hidden_size) self.layers = [ TransformerBlock(config, layer_idx=i) for i in range(config.num_hidden_layers) ] self.ln_f = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def __call__( self, inputs: mx.array, inputs_embeds: Optional[mx.array] = None, mask: Optional[mx.array] = None, cache: Optional[List[KVCache]] = None, **kwargs, ) -> LanguageModelOutput: if inputs_embeds is None: x = self.embedding(inputs) else: x = inputs_embeds # Initialize cache if needed if cache is None: cache = [None] * len(self.layers) # Create causal attention mask mask = create_attention_mask(x, cache) for i, layer in enumerate(self.layers): x = layer(x, mask=mask, cache=cache[i]) hidden_states = self.ln_f(x) logits = self.lm_head(hidden_states) return LanguageModelOutput(logits=logits)