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

mlx_vlm/models/jina_vlm/language.py

@@ -0,0 +1,272 @@
+"""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]
+        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)

mlx_vlm/models/jina_vlm/processing_jinavlm.py

@@ -0,0 +1,266 @@
+"""Processor for Jina VLM in MLX-VLM."""
+
+from typing import Dict, List, Literal, Optional, Union
+
+import mlx.core as mx
+import numpy as np
+from PIL import Image
+from transformers.processing_utils import ProcessorMixin
+
+from .image_processor import ImageProcessor
+
+
+class JinaVLMProcessor(ProcessorMixin):
+    """Processor for Jina VLM that combines tokenizer and image processor."""
+
+    tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast")
+    attributes = ["tokenizer"]
+
+    def __init__(
+        self,
+        tokenizer,
+        image_token: str = "<|image|>",
+        chat_template: Optional[str] = None,
+        **kwargs,
+    ):
+        self.tokenizer = tokenizer
+        self.image_token = image_token
+        self._image_proc = ImageProcessor()  # Internal, not exposed as image_processor
+
+        # Get image token ID
+        self.image_token_id = self.tokenizer.convert_tokens_to_ids(image_token)
+
+        super().__init__(tokenizer, **kwargs)
+
+        # Set chat template AFTER super().__init__ - always set the default if not already set
+        default_chat_template = (
+            "{% for message in messages %}"
+            "{% if message['role'] == 'system' %}{{ '<|system|>\n' + message['content'] + '\n' }}"
+            "{% elif message['role'] == 'user' %}{{ '<|user|>\n' + message['content'] + '\n' }}"
+            "{% elif message['role'] == 'assistant' %}{{ '<|assistant|>\n' + message['content'] + '\n' }}"
+            "{% endif %}{% endfor %}"
+            "{% if add_generation_prompt %}{{ '<|assistant|>\n' }}{% endif %}"
+        )
+        if chat_template is not None:
+            self.tokenizer.chat_template = chat_template
+        elif not self.tokenizer.chat_template:
+            self.tokenizer.chat_template = default_chat_template
+
+    @property
+    def chat_template(self):
+        return self.tokenizer.chat_template
+
+    @chat_template.setter
+    def chat_template(self, value):
+        self.tokenizer.chat_template = value
+
+    @property
+    def pad_token(self):
+        return self.tokenizer.pad_token
+
+    @property
+    def pad_token_id(self):
+        return self.tokenizer.pad_token_id
+
+    @property
+    def eos_token(self):
+        return self.tokenizer.eos_token
+
+    @property
+    def eos_token_id(self):
+        return self.tokenizer.eos_token_id
+
+    @property
+    def bos_token(self):
+        return self.tokenizer.bos_token
+
+    @property
+    def bos_token_id(self):
+        return self.tokenizer.bos_token_id
+
+    def encode(self, text: str, add_special_tokens: bool = False) -> List[int]:
+        return self.tokenizer.encode(text, add_special_tokens=add_special_tokens)
+
+    def decode(self, token_ids: List[int], **kwargs) -> str:
+        return self.tokenizer.decode(token_ids, **kwargs)
+
+    def batch_decode(self, token_ids, **kwargs) -> List[str]:
+        return self.tokenizer.batch_decode(token_ids, **kwargs)
+
+    def process_one(
+        self,
+        prompt: str,
+        images: Optional[List[Image.Image]] = None,
+        inference_mode: bool = True,
+    ) -> Dict:
+        """Process a single prompt with images."""
+        if images is None:
+            images = []
+
+        # Process images
+        if images:
+            image_outputs = self._image_proc.preprocess(images)
+            pixel_values_list = image_outputs["pixel_values"]
+            image_tokens = image_outputs["image_tokens"]
+            image_input_idx_list = image_outputs["image_input_idx"]
+            image_masks_list = image_outputs["image_masks"]
+        else:
+            pixel_values_list = None
+            image_tokens = []
+            image_input_idx_list = None
+            image_masks_list = None
+
+        # Split prompt by image token
+        text_splits = prompt.split(self.image_token)
+
+        # Build input_ids with image tokens interleaved
+        input_ids = []
+        current_image_idx = 0
+        updated_image_input_idx = []
+
+        for i, text_part in enumerate(text_splits):
+            # Encode text part
+            if text_part:
+                text_tokens = self.encode(text_part, add_special_tokens=False)
+                input_ids.extend(text_tokens)
+
+            # Add image tokens if not the last split and we have images
+            if i < len(text_splits) - 1 and current_image_idx < len(image_tokens):
+                # Get image tokens for this image
+                img_tokens = image_tokens[current_image_idx]
+                # Offset image_input_idx by current position
+                if image_input_idx_list is not None and current_image_idx < len(
+                    image_input_idx_list
+                ):
+                    offset_idx = image_input_idx_list[current_image_idx] + len(
+                        input_ids
+                    )
+                    updated_image_input_idx.append(offset_idx)
+                input_ids.extend(img_tokens.tolist())
+                current_image_idx += 1
+
+        input_ids = mx.array(input_ids)
+
+        result = {
+            "input_ids": input_ids[None, :],  # Add batch dimension
+            "attention_mask": mx.ones_like(input_ids)[None, :],
+        }
+
+        if pixel_values_list is not None and len(pixel_values_list) > 0:
+            # Stack pixel values: (n_crops, n_patches, patch_dim)
+            result["pixel_values"] = mx.array(np.stack(pixel_values_list))
+            # Stack image_input_idx: (n_images, tokens_per_image)
+            result["image_input_idx"] = mx.array(np.stack(updated_image_input_idx))
+            # Stack image_masks: (n_crops, n_patches)
+            result["image_masks"] = mx.array(np.stack(image_masks_list))
+
+        return result
+
+    def __call__(
+        self,
+        text: Optional[Union[str, List[str]]] = None,
+        images: Optional[Union[Image.Image, List[Image.Image]]] = None,
+        inference_mode: bool = True,
+        return_tensors: Literal["np", "mx", "pt"] = "mx",
+        **kwargs,
+    ) -> Dict:
+        """Process text and images for Jina VLM.
+
+        When called with just text (like a tokenizer), returns tokenizer output.
+        When called with text and images, returns full processed inputs.
+
+        Args:
+            text: Input text or list of texts
+            images: Input image or list of images
+            inference_mode: Whether in inference mode
+            return_tensors: Type of tensors to return
+
+        Returns:
+            Dictionary containing processed inputs
+        """
+        # If called with just text (like a tokenizer), delegate to tokenizer
+        if text is not None and images is None:
+            return self.tokenizer(text, **kwargs)
+
+        if text is None:
+            raise ValueError("Text must be provided")
+
+        # Normalize inputs
+        if isinstance(text, str):
+            texts = [text]
+        else:
+            texts = text
+
+        if images is None:
+            images_list = [None] * len(texts)
+        elif isinstance(images, Image.Image):
+            images_list = [[images]]
+        elif isinstance(images, list) and len(images) > 0:
+            if isinstance(images[0], Image.Image):
+                # Single list of images for single prompt
+                images_list = [images]
+            else:
+                images_list = images
+        else:
+            images_list = [None] * len(texts)
+
+        # Process each text-image pair
+        batch_results = []
+        for prompt, imgs in zip(texts, images_list):
+            result = self.process_one(prompt, imgs, inference_mode)
+            batch_results.append(result)
+
+        # Collate results
+        if len(batch_results) == 1:
+            return batch_results[0]
+        else:
+            return self._collate_batch(batch_results)
+
+    def _collate_batch(self, batch_results: List[Dict]) -> Dict:
+        """Collate multiple results into a batch."""
+        # Get max sequence length
+        max_len = max(r["input_ids"].shape[1] for r in batch_results)
+
+        padded_input_ids = []
+        padded_attention_mask = []
+
+        for r in batch_results:
+            seq_len = r["input_ids"].shape[1]
+            pad_len = max_len - seq_len
+
+            if pad_len > 0:
+                input_ids = mx.concatenate(
+                    [mx.full((1, pad_len), self.pad_token_id), r["input_ids"]], axis=1
+                )
+                attention_mask = mx.concatenate(
+                    [mx.zeros((1, pad_len)), r["attention_mask"]], axis=1
+                )
+            else:
+                input_ids = r["input_ids"]
+                attention_mask = r["attention_mask"]
+
+            padded_input_ids.append(input_ids)
+            padded_attention_mask.append(attention_mask)
+
+        result = {
+            "input_ids": mx.concatenate(padded_input_ids, axis=0),
+            "attention_mask": mx.concatenate(padded_attention_mask, axis=0),
+        }
+
+        # Combine pixel values if present
+        all_pixel_values = []
+        all_image_input_idx = []
+        all_image_masks = []
+
+        for r in batch_results:
+            if "pixel_values" in r:
+                all_pixel_values.append(r["pixel_values"])
+                all_image_input_idx.append(r["image_input_idx"])
+                all_image_masks.append(r["image_masks"])
+
+        if all_pixel_values:
+            result["pixel_values"] = mx.concatenate(all_pixel_values, axis=0)
+            result["image_input_idx"] = mx.concatenate(all_image_input_idx, axis=0)
+            result["image_masks"] = mx.concatenate(all_image_masks, axis=0)
+
+        return result