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

mlx_vlm/models/deepseekocr/conversation.py

@@ -0,0 +1,264 @@
+"""
+From https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py
+"""
+
+import dataclasses
+from enum import IntEnum, auto
+from typing import Dict, List
+
+
+class SeparatorStyle(IntEnum):
+    """Separator styles."""
+
+    DeepSeek = auto()
+    DeepSeekV2 = auto()
+    PLAIN = auto()
+    ALIGNMENT = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that manages prompt templates and keeps all conversation history."""
+
+    # The name of this template
+    name: str
+    # The template of the system prompt
+    system_template: str = "{system_message}"
+    # The system message
+    system_message: str = ""
+    # The names of two roles
+    roles: List[str] = (("USER", "ASSISTANT"),)
+    # All messages. Each item is (role, message).
+    messages: List[List[str]] = ()
+    # The number of few shot examples
+    offset: int = 0
+    # The separator style and configurations
+    sep_style: SeparatorStyle = SeparatorStyle.DeepSeek
+    sep: str = "\n"
+    sep2: str = None
+    # Stop criteria (the default one is EOS token)
+    stop_str: str = None
+    # Stops generation if meeting any token in this list
+    stop_token_ids: List[int] = None
+
+    def get_prompt(self) -> str:
+        """Get the prompt for generation."""
+        system_prompt = self.system_template.format(system_message=self.system_message)
+        if self.sep_style == SeparatorStyle.DeepSeek:
+            seps = [self.sep, self.sep2]
+            if system_prompt == "" or system_prompt is None:
+                ret = ""
+            else:
+                ret = system_prompt + seps[0]
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ":"
+            return ret
+        elif self.sep_style == SeparatorStyle.DeepSeekV2:
+            seps = [self.sep, self.sep2]
+            if system_prompt == "" or system_prompt is None:
+                ret = ""
+            else:
+                ret = system_prompt + seps[0]
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    if role == "User":
+                        ret += "<｜sft▁begin｜>\n" + message + self.sep
+                    else:
+                        ret += message + self.sep2
+                else:
+                    ret = ret
+            return ret
+
+        elif self.sep_style == SeparatorStyle.PLAIN:
+            seps = [self.sep, self.sep2]
+            ret = ""
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    if i % 2 == 0:
+                        ret += message + seps[i % 2]
+                    else:
+                        ret += message + seps[i % 2]
+                else:
+                    ret += ""
+            return ret
+        elif self.sep_style == SeparatorStyle.ALIGNMENT:
+            seps = [self.sep, self.sep2]
+            ret = ""
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    if i % 2 == 0:
+                        ret += "<image>\n" + seps[i % 2]
+                    else:
+                        ret += message + seps[i % 2]
+                else:
+                    ret += ""
+            return ret
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+    def set_system_message(self, system_message: str):
+        """Set the system message."""
+        self.system_message = system_message
+
+    def append_message(self, role: str, message: str):
+        """Append a new message."""
+        self.messages.append([role, message])
+
+    def update_last_message(self, message: str):
+        """Update the last output.
+
+        The last message is typically set to be None when constructing the prompt,
+        so we need to update it in-place after getting the response from a model.
+        """
+        self.messages[-1][1] = message
+
+    def reset_message(self):
+        """Reset a new message."""
+        self.messages = []
+
+    def to_gradio_chatbot(self):
+        """Convert the conversation to gradio chatbot format."""
+        ret = []
+        for i, (role, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                ret.append([msg, None])
+            else:
+                ret[-1][-1] = msg
+        return ret
+
+    def to_openai_api_messages(self):
+        """Convert the conversation to OpenAI chat completion format."""
+        system_prompt = self.system_template.format(system_message=self.system_message)
+        ret = [{"role": "system", "content": system_prompt}]
+
+        for i, (_, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                ret.append({"role": "user", "content": msg})
+            else:
+                if msg is not None:
+                    ret.append({"role": "assistant", "content": msg})
+        return ret
+
+    def copy(self):
+        return Conversation(
+            name=self.name,
+            system_template=self.system_template,
+            system_message=self.system_message,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            stop_str=self.stop_str,
+            stop_token_ids=self.stop_token_ids,
+        )
+
+    def dict(self):
+        return {
+            "template_name": self.name,
+            "system_message": self.system_message,
+            "roles": self.roles,
+            "messages": self.messages,
+            "offset": self.offset,
+        }
+
+
+# A global registry for all conversation templates
+conv_templates: Dict[str, Conversation] = {}
+
+
+def register_conv_template(template: Conversation, override: bool = False):
+    """Register a new conversation template."""
+    if not override:
+        assert (
+            template.name not in conv_templates
+        ), f"{template.name} has been registered."
+
+    conv_templates[template.name] = template
+
+
+def get_conv_template(name: str) -> Conversation:
+    """Get a conversation template."""
+    return conv_templates[name].copy()
+
+
+register_conv_template(
+    Conversation(
+        name="deepseek",
+        system_template="{system_message}",
+        # system_message="You are a helpful assistant. Please answer truthfully and write out your "
+        # "thinking step by step to be sure you get the right answer.",
+        system_message="",
+        roles=("<|User|>", "<|Assistant|>"),
+        messages=(),
+        offset=0,
+        sep_style=SeparatorStyle.DeepSeek,
+        sep="\n\n",
+        sep2="<｜end▁of▁sentence｜>",
+        stop_token_ids=[100001],
+        stop_str=["User:", "<｜end▁of▁sentence｜>"],
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name="deepseekv2",
+        system_template="{system_message}",
+        system_message="",
+        roles=("|<User>|", "|<Assistant>|"),
+        messages=(),
+        offset=0,
+        sep_style=SeparatorStyle.DeepSeekV2,
+        sep="\n<｜sft▁end｜>",
+        sep2="<｜end▁of▁sentence｜>",
+        stop_token_ids=[100001],
+        stop_str=["User:", "<｜end▁of▁sentence｜>"],
+    )
+)
+
+
+register_conv_template(
+    Conversation(
+        name="plain",
+        system_template="",
+        system_message="",
+        roles=("", ""),
+        messages=(),
+        offset=0,
+        sep_style=SeparatorStyle.PLAIN,
+        sep="",
+        sep2="",
+        stop_token_ids=[100001],
+        stop_str=["</s>"],
+    )
+)
+
+
+register_conv_template(
+    Conversation(
+        name="alignment",
+        system_template="",
+        system_message="",
+        roles=("", ""),
+        messages=(),
+        offset=0,
+        sep_style=SeparatorStyle.ALIGNMENT,
+        sep="",
+        sep2="",
+        stop_token_ids=[100001],
+        stop_str=["</s>"],
+    )
+)
+
+
+if __name__ == "__main__":
+    print("deepseek template:")
+    conv = get_conv_template("deepseek")

mlx_vlm/models/deepseekocr/deepseekocr.py

@@ -0,0 +1,371 @@
+import math
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+from transformers import AutoProcessor
+
+from mlx_vlm.models.base import InputEmbeddingsFeatures
+
+from .config import ModelConfig, ProjectorConfig, SAMViTConfig
+from .language import LanguageModel
+from .processing_deepseekocr import DeepseekOCRProcessor
+from .sam import SAMEncoder
+from .vision import VisionModel
+
+AutoProcessor.register("deepseekocr", DeepseekOCRProcessor)
+
+
+class MlpProjector(nn.Module):
+    def __init__(self, config: ProjectorConfig):
+        super().__init__()
+        self.config = config
+
+        if config.projector_config.projector_type == "linear":
+            modules = nn.Linear(
+                config.projector_config.input_dim, config.projector_config.n_embed
+            )
+
+        elif config.projector_config.projector_type == "downsample_mlp_gelu":
+            mlp_depth = config.projector_config.depth
+            mlp_ratio = config.projector_config.mlp_ratio
+            modules = [
+                nn.Linear(
+                    config.projector_config.input_dim
+                    * config.projector_config.downsample_ratio
+                    * config.projector_config.downsample_ratio,
+                    config.projector_config.n_embed * mlp_ratio,
+                )
+            ]
+            for _ in range(1, mlp_depth - 1):
+                modules.append(nn.GELU())
+                modules.append(
+                    nn.Linear(
+                        config.projector_config.n_embed * mlp_ratio,
+                        config.projector_config.n_embed * mlp_ratio,
+                    )
+                )
+            modules.append(nn.GELU())
+            modules.append(
+                nn.Linear(
+                    config.projector_config.n_embed * mlp_ratio,
+                    config.projector_config.n_embed,
+                )
+            )
+        else:
+            raise ValueError(
+                f"Unknown projector type: {config.projector_config.projector_type}"
+            )
+
+        self.layers = modules
+
+    def __call__(self, x):
+        if self.config.projector_config.projector_type == "downsample_mlp_gelu":
+            bs, hw, input_dim = x.shape
+            h = w = int(math.sqrt(hw))
+
+            # Compute padding
+            pad = (
+                0
+                if h % self.config.projector_config.downsample_ratio == 0
+                else self.config.projector_config.downsample_ratio
+                - h % self.config.projector_config.downsample_ratio
+            )
+
+            x = mx.reshape(x, (bs, h, w, input_dim))
+            if pad > 0:
+                x = mx.pad(x, [(0, 0), (0, pad), (0, pad), (0, 0)], constant_values=0)
+
+            x = mx.transpose(x, (0, 3, 1, 2))  # B, C, H, W
+
+            # Manual implementation of unfold for downsampling
+            h_pad, w_pad = x.shape[2], x.shape[3]
+            ds = self.config.projector_config.downsample_ratio
+            patches = []
+
+            for i in range(0, h_pad - ds + 1, ds):
+                for j in range(0, w_pad - ds + 1, ds):
+                    patch = x[:, :, i : i + ds, j : j + ds]
+                    patches.append(mx.reshape(patch, (bs, -1)))
+
+            x = mx.stack(patches, axis=1)  # B, N_patches, C*ds*ds
+
+        if self.config.projector_config.projector_type == "linear":
+            x = self.layers(x)
+        else:
+            for layer in self.layers:
+                x = layer(x)
+        return x
+
+
+class Model(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+        self.vision_model = VisionModel(config.vision_config)
+        sam_config = SAMViTConfig()
+        self.sam_model = SAMEncoder(
+            img_size=sam_config.image_size,
+            patch_size=sam_config.patch_size,
+            embed_dim=sam_config.width,
+            depth=sam_config.layers,
+            num_heads=sam_config.heads,
+            window_size=sam_config.window_size,
+            global_attn_indexes=sam_config.global_attn_indexes,
+        )
+        self.language_model = LanguageModel(config.text_config)
+        self.projector = MlpProjector(config)
+
+        self.tile_tag = config.tile_tag
+        self.global_view_pos = config.global_view_pos
+        # 用于format image token sequence的特殊token
+        embed_std = 1 / mx.sqrt(
+            mx.array(config.projector_config.n_embed, dtype=mx.float32)
+        )
+
+        if self.tile_tag == "2D":
+
+            # <|view_separator|>, <|\n|>
+            self.image_newline = mx.array(
+                mx.random.normal((config.projector_config.n_embed,)) * embed_std
+            )
+            # fix the typo: view_seperater
+            self.view_separator = mx.array(
+                mx.random.normal((config.projector_config.n_embed,)) * embed_std
+            )
+        else:
+            raise ValueError(
+                f"Only 2D tile_tag is supported currently, got: {self.tile_tag}"
+            )
+
+    def get_input_embeddings(
+        self,
+        input_ids: Optional[mx.array] = None,
+        pixel_values: Optional[mx.array] = None,
+        images_spatial_crop: Optional[mx.array] = None,
+        images_seq_mask: Optional[mx.array] = None,
+        **kwargs,
+    ):
+        input_embeds = self.language_model.model.embed_tokens(input_ids)
+
+        if pixel_values is None:
+            return InputEmbeddingsFeatures(inputs_embeds=input_embeds)
+
+        # Only process images on prefill (input_ids.shape[1] != 1), not during autoregressive decoding
+        if (
+            self.sam_model is not None
+            and input_ids.shape[1] != 1
+            and mx.sum(pixel_values[1]).item() != 0
+        ):
+
+            idx = 0
+            patch_idx = 0  # Track patch offset for batch processing
+            all_patches = pixel_values[0]
+            all_image_ori = pixel_values[1]
+
+            for crop_shape in images_spatial_crop.tolist():
+                images_in_this_batch = []
+                width_crop_num, height_crop_num = int(crop_shape[0]), int(crop_shape[1])
+
+                # Calculate number of patches for this image
+                has_crops = width_crop_num > 1 or height_crop_num > 1
+                num_patches = width_crop_num * height_crop_num if has_crops else 0
+
+                # Extract patches for current image
+                if has_crops and num_patches > 0:
+                    patches = all_patches[patch_idx : patch_idx + num_patches]
+                    patch_idx += num_patches
+                else:
+                    patches = None
+
+                # Extract global image for current image (one per batch item)
+                image_ori = all_image_ori[idx : idx + 1]
+
+                if patches is not None and mx.sum(patches).item() != 0:
+                    local_features_1 = self.sam_model(patches.transpose(0, 2, 3, 1))
+
+                    local_features_2 = self.vision_model(
+                        patches.transpose(0, 2, 3, 1), patch_embeds=local_features_1
+                    )
+
+                    local_features = mx.concatenate(
+                        (
+                            local_features_2[:, 1:],
+                            local_features_1.flatten(start_axis=1, end_axis=2),
+                        ),
+                        axis=-1,
+                    )
+
+                    local_features = self.projector(local_features)
+
+                    global_features_1 = self.sam_model(image_ori.transpose(0, 2, 3, 1))
+                    global_features_2 = self.vision_model(
+                        image_ori.transpose(0, 2, 3, 1), global_features_1
+                    )
+
+                    global_features = mx.concatenate(
+                        (
+                            global_features_2[:, 1:],
+                            global_features_1.flatten(start_axis=1, end_axis=2),
+                        ),
+                        axis=-1,
+                    )
+                    global_features = self.projector(global_features)
+
+                    # Remove batch dimension for single image processing
+                    global_features = global_features[0]  # (hw, n_dim)
+                    hw, n_dim = global_features.shape
+                    h = w = int(hw**0.5)
+
+                    _, hw2, n_dim2 = local_features.shape
+                    h2 = w2 = int(hw2**0.5)
+
+                    global_features = global_features.reshape(h, w, n_dim)
+
+                    global_features = mx.concatenate(
+                        [
+                            global_features,
+                            mx.broadcast_to(
+                                self.image_newline[None, None, :], (h, 1, n_dim)
+                            ),
+                        ],
+                        axis=1,
+                    )
+
+                    global_features = global_features.reshape(-1, n_dim)
+
+                    local_features = (
+                        local_features.reshape(
+                            height_crop_num, width_crop_num, h2, w2, n_dim2
+                        )
+                        .transpose(0, 2, 1, 3, 4)
+                        .reshape(height_crop_num * h2, width_crop_num * w2, n_dim2)
+                    )
+                    local_features = mx.concatenate(
+                        [
+                            local_features,
+                            mx.broadcast_to(
+                                self.image_newline[None, None, :],
+                                (height_crop_num * h2, 1, n_dim2),
+                            ),
+                        ],
+                        axis=1,
+                    )
+                    local_features = local_features.reshape(-1, n_dim2)
+
+                    global_local_features = mx.concatenate(
+                        [local_features, global_features, self.view_separator[None, :]],
+                        axis=0,
+                    )
+
+                else:
+                    global_features_1 = self.sam_model(image_ori.transpose(0, 2, 3, 1))
+                    global_features_2 = self.vision_model(
+                        image_ori.transpose(0, 2, 3, 1), global_features_1
+                    )
+                    global_features = mx.concatenate(
+                        (
+                            global_features_2[:, 1:],
+                            global_features_1.flatten(start_axis=1, end_axis=2),
+                        ),
+                        axis=-1,
+                    )
+                    global_features = self.projector(global_features)
+
+                    # Remove batch dimension for single image processing
+                    global_features = global_features[0]  # (hw, n_dim)
+                    hw, n_dim = global_features.shape
+                    h = w = int(hw**0.5)
+
+                    global_features = global_features.reshape(h, w, n_dim)
+
+                    global_features = mx.concatenate(
+                        [
+                            global_features,
+                            mx.broadcast_to(
+                                self.image_newline[None, None, :], (h, 1, n_dim)
+                            ),
+                        ],
+                        axis=1,
+                    )
+
+                    global_features = global_features.reshape(-1, n_dim)
+
+                    global_local_features = mx.concatenate(
+                        [global_features, self.view_separator[None, :]], axis=0
+                    )
+
+                images_in_this_batch.append(global_local_features)
+
+                if images_in_this_batch:
+                    images_in_this_batch = mx.concatenate(images_in_this_batch, axis=0)
+                    # Find positions where images should be placed
+                    image_indices = np.where(images_seq_mask[idx])[0].tolist()
+                    # Directly assign the image features to those positions
+                    input_embeds[idx, image_indices] = images_in_this_batch
+
+                idx += 1
+
+        return InputEmbeddingsFeatures(inputs_embeds=input_embeds)
+
+    @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,
+    ):
+
+        images_spatial_crop = kwargs.get("images_spatial_crop", None)
+        images_seq_mask = kwargs.get("images_seq_mask", None)
+
+        input_embeddings = self.get_input_embeddings(
+            input_ids, pixel_values, images_spatial_crop, images_seq_mask
+        )
+
+        logits = self.language_model(
+            input_ids, cache=cache, inputs_embeds=input_embeddings.inputs_embeds
+        )
+        return logits
+
+    @staticmethod
+    def sanitize(weights):
+        def transform_key(key):
+            if "model.layers" in key and "language_model" not in key:
+                key = key.replace("model.layers", "language_model.model.layers")
+
+            if "model.embed_tokens" in key and "language_model" not in key:
+                key = key.replace(
+                    "model.embed_tokens", "language_model.model.embed_tokens"
+                )
+
+            if "model.norm" in key and "language_model" not in key:
+                key = key.replace("model.norm", "language_model.model.norm")
+
+            if "model.vision_model" in key:
+                key = key.replace("model.vision_model", "vision_model")
+
+            if "model.sam_model" in key:
+                key = key.replace("model.sam_model", "sam_model")
+
+            if "model.projector" in key:
+                key = key.replace("model.projector", "projector")
+
+            if "model.view_seperator" in key:
+                key = key.replace("model.view_seperator", "view_separator")
+
+            if "model.image_newline" in key:
+                key = key.replace("model.image_newline", "image_newline")
+
+            if "lm_head.weight" in key and "language_model" not in key:
+                key = key.replace("lm_head.weight", "language_model.lm_head.weight")
+
+            return key
+
+        return {transform_key(k): v for k, v in weights.items()}