fount-vlm-nell-02 0.3.11__py3-none-any.whl

mlx_vlm/models/moondream2/language.py

@@ -0,0 +1,267 @@
+from typing import Optional, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from ..base import (
+    LanguageModelOutput,
+    create_attention_mask,
+    scaled_dot_product_attention,
+)
+from .config import TextConfig
+
+
+class Attention(nn.Module):
+    """Multi-head attention with partial RoPE (32 dims)."""
+
+    def __init__(self, config: TextConfig):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.scale = self.head_dim**-0.5
+
+        # Combined QKV projection (like original moondream)
+        qkv_dim = self.hidden_size + 2 * (self.num_kv_heads * self.head_dim)
+        self.qkv_proj = nn.Linear(self.hidden_size, qkv_dim, bias=True)
+        self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=True)
+
+        # Partial RoPE: only apply to first half of head_dim (32 out of 64)
+        rope_dims = int(self.head_dim * config.partial_rotary_factor)
+        self.rope = nn.RoPE(
+            dims=rope_dims,
+            traditional=False,
+            base=config.rope_theta,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = None,
+    ) -> mx.array:
+        B, L, _ = x.shape
+
+        # Combined QKV projection
+        qkv = self.qkv_proj(x)
+
+        # Split into Q, K, V
+        q_dim = self.num_heads * self.head_dim
+        kv_dim = self.num_kv_heads * self.head_dim
+        q, k, v = mx.split(qkv, [q_dim, q_dim + kv_dim], axis=-1)
+
+        # Reshape for attention
+        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)
+
+        # Apply partial RoPE
+        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)
+
+        # Attention
+        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.o_proj(output)
+
+
+class MLP(nn.Module):
+    """Simple MLP with GELU activation (not gated like Phi3)."""
+
+    def __init__(self, config: TextConfig):
+        super().__init__()
+        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=True)
+        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=True)
+        self.activation = nn.GELU(approx="precise")
+
+    def __call__(self, x: mx.array) -> mx.array:
+        x = self.activation(self.fc1(x))
+        x = self.fc2(x)
+        return x
+
+
+class TransformerBlock(nn.Module):
+    """Transformer block with pre-norm using LayerNorm."""
+
+    def __init__(self, config: TextConfig):
+        super().__init__()
+        self.self_attn = Attention(config)
+        self.mlp = MLP(config)
+        # Moondream uses a single LayerNorm before both attention and MLP
+        # The residual pattern is: x + attn(ln(x)) + mlp(ln(x))
+        self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = None,
+    ) -> mx.array:
+        # Moondream uses parallel attention and MLP
+        # x = x + attn(ln(x)) + mlp(ln(x))
+        normalized = self.input_layernorm(x)
+        attn_out = self.self_attn(normalized, mask, cache)
+        mlp_out = self.mlp(normalized)
+        return x + attn_out + mlp_out
+
+
+class PhiModel(nn.Module):
+    """Core transformer model."""
+
+    def __init__(self, config: TextConfig):
+        super().__init__()
+        self.config = config
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = [TransformerBlock(config) for _ in range(config.num_hidden_layers)]
+        self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        inputs_embeds: Optional[mx.array] = None,
+        mask: Optional[mx.array] = None,
+        cache=None,
+    ):
+        if inputs_embeds is None:
+            h = self.embed_tokens(inputs)
+        else:
+            h = inputs_embeds
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
+            # Moondream uses a special "prefix attention" mask where the
+            # BOS+image patch tokens attend fully within the prefix.
+            prefix_len = getattr(self.config, "prefix_attn_len", None)
+            try:
+                # Only apply during prefill (offset == 0) and only if prefix fits.
+                cache0 = cache[0] if isinstance(cache, list) and len(cache) > 0 else None
+                offset0 = getattr(cache0, "offset", None)
+
+                if (
+                    prefix_len is not None
+                    and offset0 == 0
+                    and hasattr(mask, "ndim")
+                    and mask.ndim >= 4
+                    and h.shape[1] >= prefix_len
+                ):
+                    if str(mask.dtype) == "bool":
+                        mask[..., :prefix_len, :prefix_len] = True
+                    else:
+                        # For additive masks, 0 indicates "allowed"
+                        mask[..., :prefix_len, :prefix_len] = 0
+            except Exception:
+                pass
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, c)
+
+        return self.norm(h)
+
+
+class LanguageModel(nn.Module):
+    """Language model with LM head."""
+
+    def __init__(self, config: TextConfig):
+        super().__init__()
+        self.config = config
+        self.model = PhiModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=True)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        inputs_embeds: Optional[mx.array] = None,
+        mask: Optional[mx.array] = None,
+        cache=None,
+        **kwargs,
+    ):
+        # #region agent log
+        import json
+        log_file = "/Users/zekieldee/Desktop/code/mlx-vlm/.cursor/debug.log"
+        def log_lm(location, message, data, hypothesis_id):
+            try:
+                with open(log_file, "a") as f:
+                    f.write(json.dumps({"sessionId": "debug-session", "runId": "inference", "hypothesisId": hypothesis_id, "location": location, "message": message, "data": data, "timestamp": __import__("time").time_ns() // 1000000}) + "\n")
+            except: pass
+        
+        if inputs_embeds is not None:
+            log_lm("language.py:lm_input_embeds", "Language model input embeddings", {
+                "shape": str(inputs_embeds.shape),
+                "dtype": str(inputs_embeds.dtype),
+                "mean": float(mx.mean(inputs_embeds)),
+                "std": float(mx.std(inputs_embeds)),
+                "min": float(mx.min(inputs_embeds)),
+                "max": float(mx.max(inputs_embeds)),
+                "has_nan": bool(mx.any(mx.isnan(inputs_embeds))),
+                "has_inf": bool(mx.any(mx.isinf(inputs_embeds)))
+            }, "H9")
+        # #endregion
+        
+        h = self.model(inputs, inputs_embeds=inputs_embeds, mask=mask, cache=cache)
+        
+        # #region agent log
+        log_lm("language.py:lm_hidden_states", "Language model hidden states after layers", {
+            "shape": str(h.shape),
+            "dtype": str(h.dtype),
+            "mean": float(mx.mean(h)),
+            "std": float(mx.std(h)),
+            "min": float(mx.min(h)),
+            "max": float(mx.max(h)),
+            "last_token_mean": float(mx.mean(h[:, -1, :])),
+            "last_token_std": float(mx.std(h[:, -1, :]))
+        }, "H9")
+        # #endregion
+        
+        logits = self.lm_head(h)
+        
+        # #region agent log
+        last_logits = logits[0, -1, :]
+        top5_idx = mx.argsort(last_logits)[-5:].tolist()
+        top5_val = mx.sort(last_logits)[-5:].tolist()
+        log_lm("language.py:lm_logits", "Language model logits output", {
+            "shape": str(logits.shape),
+            "dtype": str(logits.dtype),
+            "mean": float(mx.mean(logits)),
+            "std": float(mx.std(logits)),
+            "min": float(mx.min(logits)),
+            "max": float(mx.max(logits)),
+            "last_token_logits_mean": float(mx.mean(logits[:, -1, :])),
+            "last_token_logits_std": float(mx.std(logits[:, -1, :])),
+            "last_token_top5_indices": top5_idx,
+            "last_token_top5_values": top5_val
+        }, "H9,H10")
+        # #endregion
+        
+        return LanguageModelOutput(logits=logits)
+
+    def sanitize(self, weights):
+        """Sanitize language model weights."""
+        return {
+            k: v for k, v in weights.items() if "self_attn.rotary_emb.inv_freq" not in k
+        }
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    @property
+    def head_dim(self):
+        return self.config.hidden_size // self.config.num_attention_heads
+
+    @property
+    def n_kv_heads(self):
+        return self.config.num_key_value_heads