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

mlx_vlm/models/glm_ocr/config.py

@@ -0,0 +1,93 @@
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional
+
+from ..base import BaseModelConfig
+
+
+@dataclass
+class TextConfig(BaseModelConfig):
+    model_type: str = "glm_ocr_text"
+    vocab_size: int = 59392
+    hidden_size: int = 1536
+    eos_token_id: List[int] = field(default_factory=lambda: [59246, 59253])
+    intermediate_size: int = 4608
+    max_position_embeddings: int = 131072
+    num_attention_heads: int = 16
+    num_hidden_layers: int = 16
+    num_key_value_heads: int = 8
+    head_dim: int = 128
+    rms_norm_eps: float = 1e-05
+    attention_bias: bool = False
+    attention_dropout: float = 0.0
+    hidden_act: str = "silu"
+    initializer_range: float = 0.02
+    rope_parameters: Dict = field(
+        default_factory=lambda: {
+            "rope_type": "default",
+            "mrope_section": [16, 24, 24],
+            "partial_rotary_factor": 1.0,
+            "rope_theta": 10000,
+        }
+    )
+    pad_token_id: int = 59246
+    use_cache: bool = True
+    tie_word_embeddings: bool = False
+    num_nextn_predict_layers: int = 1
+    dtype: str = "bfloat16"
+
+    @property
+    def rope_theta(self) -> float:
+        return self.rope_parameters.get("rope_theta", 10000)
+
+    @property
+    def partial_rotary_factor(self) -> float:
+        return self.rope_parameters.get("partial_rotary_factor", 1.0)
+
+
+@dataclass
+class VisionConfig(BaseModelConfig):
+    model_type: str = "glm_ocr_vision"
+    depth: int = 24
+    hidden_size: int = 1024
+    intermediate_size: int = 4096
+    num_heads: int = 16
+    patch_size: int = 14
+    image_size: int = 336
+    in_channels: int = 3
+    rms_norm_eps: float = 1e-05
+    attention_bias: bool = True
+    attention_dropout: float = 0.0
+    hidden_act: str = "silu"
+    hidden_dropout_prob: float = 0.0
+    initializer_range: float = 0.02
+    out_hidden_size: int = 1536
+    spatial_merge_size: int = 2
+    temporal_patch_size: int = 2
+
+
+@dataclass
+class ModelConfig(BaseModelConfig):
+    text_config: TextConfig = None
+    vision_config: VisionConfig = None
+    model_type: str = "glm_ocr"
+    vocab_size: int = 59392
+    ignore_index: int = -100
+    image_token_id: int = 59280
+    video_token_id: int = 59281
+    image_start_token_id: int = 59256
+    image_end_token_id: int = 59257
+    video_start_token_id: int = 59258
+    video_end_token_id: int = 59259
+    hidden_size: int = 1536
+    pad_token_id: int = 59246
+    eos_token_id: Optional[List[int]] = None
+    tie_word_embeddings: bool = False
+
+    def __post_init__(self):
+        if isinstance(self.text_config, dict):
+            self.text_config = TextConfig(**self.text_config)
+        if isinstance(self.vision_config, dict):
+            self.vision_config = VisionConfig(**self.vision_config)
+
+        if self.eos_token_id is None and self.text_config is not None:
+            self.eos_token_id = self.text_config.eos_token_id

mlx_vlm/models/glm_ocr/glm_ocr.py

@@ -0,0 +1,180 @@
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from ..base import InputEmbeddingsFeatures
+from .config import ModelConfig
+from .language import LanguageModel
+from .vision import VisionModel
+
+
+class Model(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+        self.vision_tower = VisionModel(config.vision_config)
+        self.language_model = LanguageModel(config.text_config, config)
+
+    def get_input_embeddings(
+        self,
+        input_ids: Optional[mx.array] = None,
+        pixel_values: Optional[mx.array] = None,
+        **kwargs,
+    ):
+        image_grid_thw = kwargs.get("image_grid_thw", None)
+        video_grid_thw = kwargs.get("video_grid_thw", None)
+        mask = kwargs.get("mask", None)
+        grid_thw = image_grid_thw if image_grid_thw is not None else video_grid_thw
+
+        if pixel_values is None:
+            # Reset position state for text-only generation
+            self.language_model._position_ids = None
+            self.language_model._rope_deltas = None
+            return InputEmbeddingsFeatures(
+                inputs_embeds=self.language_model.model.embed_tokens(input_ids)
+            )
+
+        dtype = self.vision_tower.patch_embed.proj.weight.dtype
+        pixel_values = pixel_values.astype(dtype)
+
+        inputs_embeds = self.language_model.model.embed_tokens(input_ids)
+
+        hidden_states = self.vision_tower(
+            pixel_values, grid_thw, output_hidden_states=False
+        )
+
+        split_sizes = (
+            image_grid_thw.prod(-1) // self.vision_tower.spatial_merge_size**2
+        ).tolist()
+
+        if len(split_sizes) > 1:
+            split_indices = []
+            cumsum = 0
+            for size in split_sizes[:-1]:
+                cumsum += size
+                split_indices.append(cumsum)
+            hidden_states = mx.split(hidden_states, split_indices, axis=0)
+            hidden_states = mx.concatenate(hidden_states, axis=0).astype(
+                hidden_states[0].dtype
+            )
+
+        final_inputs_embeds = self.merge_input_ids_with_image_features(
+            self.config.image_token_id,
+            self.config.video_token_id,
+            hidden_states,
+            inputs_embeds,
+            input_ids,
+        )
+
+        # Pre-calculate position_ids for chunked prefill (like Qwen3-VL)
+        if image_grid_thw is not None or video_grid_thw is not None:
+            position_ids, rope_deltas = self.language_model.get_rope_index(
+                input_ids, image_grid_thw, video_grid_thw, mask
+            )
+            self.language_model._position_ids = position_ids
+            self.language_model._rope_deltas = rope_deltas
+
+        return InputEmbeddingsFeatures(inputs_embeds=final_inputs_embeds)
+
+    @staticmethod
+    def merge_input_ids_with_image_features(
+        image_token_id,
+        video_token_id,
+        image_features,
+        inputs_embeds,
+        input_ids,
+    ):
+        """Merge image features into input embeddings at image token positions."""
+        image_positions = input_ids == image_token_id
+        if mx.sum(image_positions) == 0:
+            image_positions = input_ids == video_token_id
+
+        batch_size, seq_len = input_ids.shape
+
+        batch_outputs = []
+        feature_start_idx = 0
+
+        for batch_idx in range(batch_size):
+            image_mask = image_positions[batch_idx]
+            num_positions = mx.sum(image_mask).item()
+
+            if num_positions > 0:
+                batch_features = image_features[
+                    feature_start_idx : feature_start_idx + num_positions
+                ]
+
+                if batch_features.shape[0] != num_positions:
+                    raise ValueError(
+                        f"Number of image token positions ({num_positions}) does not match "
+                        f"number of image features ({batch_features.shape[0]}) for batch {batch_idx}"
+                    )
+
+                cumsum = mx.cumsum(image_mask.astype(mx.int32))
+                feature_indices = mx.where(image_mask, cumsum - 1, 0)
+
+                gathered_features = batch_features[feature_indices]
+
+                image_mask_expanded = mx.expand_dims(image_mask, axis=-1)
+                batch_output = mx.where(
+                    image_mask_expanded, gathered_features, inputs_embeds[batch_idx]
+                )
+
+                feature_start_idx += num_positions
+            else:
+                batch_output = inputs_embeds[batch_idx]
+
+            batch_outputs.append(batch_output)
+
+        return mx.stack(batch_outputs, axis=0)
+
+    @property
+    def layers(self):
+        return self.language_model.model.layers
+
+    def __call__(
+        self,
+        input_ids: mx.array,
+        pixel_values: Optional[mx.array] = None,
+        mask: Optional[mx.array] = None,
+        cache=None,
+        **kwargs,
+    ):
+
+        input_embeddings_features = self.get_input_embeddings(
+            input_ids, pixel_values, **kwargs
+        )
+
+        logits = self.language_model(
+            input_ids,
+            input_embeddings_features.inputs_embeds,
+            mask=mask,
+            cache=cache,
+            **kwargs,
+        )
+
+        return logits
+
+    def sanitize(self, weights):
+        def transform_key(key):
+            if "visual" in key:
+                if "vision_tower" not in key:
+                    key = key.replace("model.", "").replace("visual", "vision_tower")
+            if "model.language_model" in key:
+                key = key.replace("model.language_model", "language_model.model")
+            if "lm_head" in key and not key.startswith("language_model"):
+                key = key.replace("lm_head", "language_model.lm_head")
+            return key
+
+        sanitized = {}
+        for k, v in weights.items():
+            new_key = transform_key(k)
+
+            if "layers.16" in new_key:
+                continue
+
+            sanitized[new_key] = v
+
+        sanitized = self.vision_tower.sanitize(sanitized)
+
+        return sanitized