nexaai 1.0.29__cp310-cp310-macosx_14_0_universal2.whl

nexaai/mlx_backend/vlm/modeling/models/gemma3n/config.py

@@ -0,0 +1,139 @@
+import inspect
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Union
+
+
+@dataclass
+class AudioConfig:
+    input_feat_size: int = 80
+    hidden_size: int = 1536
+    conf_attention_chunk_size: int = 12
+    conf_attention_context_left: int = 13
+    conf_attention_context_right: int = 0
+    conf_attention_invalid_logits_value: float = -1e9
+    conf_attention_logit_cap: float = 50.0
+    conf_num_attention_heads: int = 8
+    conf_num_hidden_layers: int = 12
+    conf_conv_kernel_size: int = 5
+    conf_positional_bias_size: int = 256
+    conf_reduction_factor: int = 4
+    conf_residual_weight: float = 0.5
+    sscp_conv_channel_size: tuple[int, int] = (128, 32)
+    sscp_conv_group_norm_eps: float = 1e-3
+    sscp_conv_kernel_size: tuple[tuple[int, int], tuple[int, int]] = ((3, 3), (3, 3))
+    sscp_conv_stride_size: tuple[tuple[int, int], tuple[int, int]] = ((2, 2), (2, 2))
+    vocab_size: int = 128
+    sscp_conv_eps: float = 1e-3
+    rms_norm_eps: float = 1e-6
+    gradient_clipping: float = 10000000000.0
+    vocab_offset: int = 262_144 + 128  # text vocab size + vision vocab size
+
+    @classmethod
+    def from_dict(cls, params):
+        return cls(
+            **{
+                k: v
+                for k, v in params.items()
+                if k in inspect.signature(cls).parameters
+            }
+        )
+
+
+@dataclass
+class VisionConfig:
+    model_type: str = "gemma3n_vision"
+    num_hidden_layers: int = 12
+    hidden_size: int = 2048
+    intermediate_size: int = 8192
+    num_attention_heads: int = 16
+    patch_size: int = 16
+    image_size: int = 224
+    num_channels: int = 3
+    rms_norm_eps: float = 1e-6
+    vocab_size: int = 128
+    vocab_offset: int = 262_144
+
+    @classmethod
+    def from_dict(cls, params):
+        return cls(
+            **{
+                k: v
+                for k, v in params.items()
+                if k in inspect.signature(cls).parameters
+            }
+        )
+
+
+@dataclass
+class TextConfig:
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int = 2
+    head_dim: int = 256
+    rms_norm_eps: float = 1.0e-6
+    vocab_size: int = 262400
+    vocab_size_per_layer_input: int = 262144
+    num_key_value_heads: int = 4
+    laurel_rank: int = 64
+    frac_shared_layers: float = 0.5
+    altup_active_idx: int = 0
+    pad_token_id: int = 0
+    altup_num_inputs: int = 4
+    altup_coef_clip: Optional[float] = None
+    altup_correct_scale: bool = True
+    hidden_size_per_layer_input: int = 1024
+    rope_local_base_freq: float = 10000.0
+    rope_traditional: bool = False
+    rope_theta: float = 1000000.0
+    query_pre_attn_scalar: float = 0.0625
+    sliding_window: int = 1024
+    rope_scaling: Optional[Dict[str, Union[float, List[float]]]] = None
+    mm_tokens_per_image: int = 256
+    sliding_window_pattern: int = 5
+    activation_sparsity_pattern: Optional[List[float]] = None
+    final_logit_softcapping: float = 30.0
+    query_rescale_scalar: float = 1.0
+    num_kv_shared_layers: int = 0
+    max_position_embeddings: int = 32768
+    attn_logit_softcapping: float = 0.0
+    layer_types: List[str] = None
+
+    @classmethod
+    def from_dict(cls, params):
+        return cls(
+            **{
+                k: v
+                for k, v in params.items()
+                if k in inspect.signature(cls).parameters
+            }
+        )
+
+
+@dataclass
+class ModelConfig:
+    text_config: TextConfig
+    vision_config: VisionConfig
+    audio_config: AudioConfig
+    model_type: str
+    vocab_size: int = 257152
+    ignore_index: int = -100
+    image_token_index: int = 262145
+    audio_token_id: int = 262273
+    image_token_id: int = 262145
+    hidden_size: int = 2048
+    pad_token_id: int = 0
+    vision_soft_tokens_per_image: int = 256
+    audio_soft_tokens_per_image: int = 188
+    eos_token_id: Optional[List[int]] = None
+
+    @classmethod
+    def from_dict(cls, params):
+        return cls(
+            **{
+                k: v
+                for k, v in params.items()
+                if k in inspect.signature(cls).parameters
+            }
+        )

nexaai/mlx_backend/vlm/modeling/models/gemma3n/gemma3n.py

@@ -0,0 +1,322 @@
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .audio import AudioModel
+from .config import ModelConfig, TextConfig
+from .language import Gemma3nRMSNorm, LanguageModel
+from .vision import VisionModel
+
+
+def masked_scatter(input_tensor, mask, source):
+    """MLX implementation of PyTorch's masked_scatter"""
+
+    # Convert mask to boolean once
+    mask = mask.astype(mx.bool_)
+
+    # Early exit
+    if not mask.any():
+        return mx.broadcast_to(input_tensor, mask.shape)
+
+    # Flatten everything once
+    input_shape = mask.shape
+    result_flat = mx.broadcast_to(input_tensor, input_shape).flatten()
+    mask_flat = mask.flatten()
+    source_flat = source.flatten()
+
+    # Create selection indices using cumulative sum
+    selection_mask = mx.cumsum(mask_flat.astype(mx.int32)) - 1
+
+    # Bound check and create source selection
+    source_len = len(source_flat)
+    bounded_indices = selection_mask % source_len
+
+    # Vectorized selection from source
+    selected_values = source_flat[bounded_indices]
+
+    result_flat = mx.where(mask_flat, selected_values, result_flat)
+
+    return result_flat.reshape(input_shape)
+
+
+class Gemma3nMultimodalEmbedder(nn.Module):
+    """Embeds token ids or soft tokens into language model space."""
+
+    def __init__(self, multimodal_config: ModelConfig, text_config: TextConfig):
+        super().__init__()
+
+        self.multimodal_hidden_size = multimodal_config.hidden_size
+        self.eps = multimodal_config.rms_norm_eps
+        self.vocab_offset = multimodal_config.vocab_offset
+        self.vocab_size = multimodal_config.vocab_size
+        self.text_hidden_size = text_config.hidden_size
+
+        self.embedding = nn.Embedding(self.vocab_size, self.multimodal_hidden_size)
+        self.hard_embedding_norm = Gemma3nRMSNorm(
+            self.multimodal_hidden_size, eps=self.eps
+        )
+        self.soft_embedding_norm = Gemma3nRMSNorm(
+            self.multimodal_hidden_size, eps=self.eps
+        )
+        self.embedding_projection = nn.Linear(
+            self.multimodal_hidden_size, self.text_hidden_size, bias=False
+        )
+        self.embedding_post_projection_norm = Gemma3nRMSNorm(
+            self.text_hidden_size, eps=self.eps, with_scale=False
+        )
+
+    def __call__(
+        self, input_ids: mx.array = None, inputs_embeds: mx.array = None
+    ) -> mx.array:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You must specify exactly one of input_ids or inputs_embeds"
+            )
+
+        if inputs_embeds is not None:
+            emb_norm = self.soft_embedding_norm(inputs_embeds)
+        else:
+
+            hard_emb = self.embedding(input_ids - self.vocab_offset)
+            emb_norm = self.hard_embedding_norm(hard_emb)
+
+        emb_norm_proj = self.embedding_projection(emb_norm)
+        projected = self.embedding_post_projection_norm(emb_norm_proj)
+        return projected
+
+
+class Model(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.model_type = config.model_type
+        self.config = config
+
+        # Text
+        self.language_model = LanguageModel(config.text_config)
+        self.vocab_size = config.text_config.vocab_size
+        self.vocab_size_per_layer_input = config.text_config.vocab_size_per_layer_input
+
+        # Vision
+        self.vision_tower = VisionModel(config.vision_config)
+        self.embed_vision = Gemma3nMultimodalEmbedder(
+            config.vision_config, text_config=config.text_config
+        )
+
+        # Audio
+        self.audio_tower = AudioModel(config.audio_config)
+        self.embed_audio = Gemma3nMultimodalEmbedder(
+            config.audio_config, text_config=config.text_config
+        )
+
+    def get_input_embeddings(
+        self,
+        input_ids: Optional[mx.array] = None,
+        pixel_values: Optional[mx.array] = None,
+        input_features: Optional[mx.array] = None,
+        input_features_mask: Optional[mx.array] = None,
+        **kwargs,
+    ):
+
+        inputs_embeds = self.language_model.model.embed_tokens(input_ids)
+
+        per_layer_inputs_mask = mx.logical_and(
+            input_ids >= 0, input_ids < self.vocab_size_per_layer_input
+        )
+        per_layer_inputs_tokens = mx.where(
+            per_layer_inputs_mask, input_ids, mx.zeros_like(input_ids)
+        )
+        per_layer_inputs = self.language_model.model.get_per_layer_inputs(
+            per_layer_inputs_tokens
+        )
+        if pixel_values is None and input_features is None:
+            return inputs_embeds, per_layer_inputs
+
+        if input_ids is not None:
+
+            # Handle vision tokens (>= embed_vision.vocab_offset and < embed_audio.vocab_offset)
+            vision_mask = mx.logical_and(
+                input_ids >= self.embed_vision.vocab_offset,
+                input_ids < self.embed_audio.vocab_offset,
+            )
+            dummy_vision_token_id = (
+                self.embed_vision.vocab_offset + self.embed_vision.vocab_size - 1
+            )
+            vision_tokens = mx.where(vision_mask, input_ids, dummy_vision_token_id)
+            vision_embeds_flat = self.embed_vision(input_ids=vision_tokens)
+            inputs_embeds = mx.where(
+                vision_mask[..., None], vision_embeds_flat, inputs_embeds
+            )
+
+            # Handle audio tokens (>= embed_audio.vocab_offset)
+            audio_mask = input_ids >= self.embed_audio.vocab_offset
+            dummy_audio_token_id = (
+                self.embed_audio.vocab_offset + self.embed_audio.vocab_size - 1
+            )
+
+            audio_tokens = mx.where(audio_mask, input_ids, dummy_audio_token_id)
+            audio_embeds_flat = self.embed_audio(input_ids=audio_tokens)
+            inputs_embeds = mx.where(
+                audio_mask[..., None], audio_embeds_flat, inputs_embeds
+            )
+        else:
+            per_layer_inputs = None
+
+        # Vision features
+        if pixel_values is not None:
+            image_features = self.get_image_features(
+                pixel_values, self.vision_tower, self.config, self.embed_vision
+            )
+
+            modality = "image"
+            inputs_embeds = self.merge_multimodal_and_text(
+                inputs_embeds,
+                image_features,
+                self.construct_special_modality_mask(
+                    input_ids,
+                    inputs_embeds,
+                    self.config.image_token_id,
+                    modality=modality,
+                ),
+                modality=modality,
+            )
+
+        # Audio features
+        if input_features is not None:
+            audio_features, audio_mask = self.get_audio_features(
+                input_features, ~input_features_mask
+            )
+            audio_padding_ids = mx.array([[self.vocab_size - 1]])
+            audio_padding_embs = self.embed_audio(input_ids=audio_padding_ids)
+            audio_features = mx.where(
+                audio_mask[..., None], audio_padding_embs, audio_features
+            )
+
+            audio_batch_size, audio_seq_len, audio_embed_dim = audio_features.shape
+            extra_padding_tokens = (
+                self.config.audio_soft_tokens_per_image - audio_seq_len
+            )
+            extra_padding_features = mx.broadcast_to(
+                audio_padding_embs,
+                (audio_batch_size, extra_padding_tokens, audio_embed_dim),
+            )
+
+            audio_features = mx.concatenate(
+                (audio_features, extra_padding_features), axis=1
+            )
+            modality = "audio"
+            inputs_embeds = self.merge_multimodal_and_text(
+                inputs_embeds,
+                audio_features,
+                self.construct_special_modality_mask(
+                    input_ids,
+                    inputs_embeds,
+                    self.config.audio_token_id,
+                    modality=modality,
+                ),
+                modality=modality,
+            )
+
+        return inputs_embeds, per_layer_inputs
+
+    def get_audio_features(self, input_features, input_features_mask):
+        audio_outputs, audio_mask = self.audio_tower(
+            input_features, input_features_mask
+        )
+        return self.embed_audio(inputs_embeds=audio_outputs), audio_mask
+
+    @staticmethod
+    def get_image_features(pixel_values, vision_tower, config, embed_vision):
+        vision_outputs = vision_tower(
+            pixel_values,
+            output_hidden_states=True,
+        )
+        vision_outputs = vision_outputs.transpose(0, 3, 1, 2)
+        vision_outputs = vision_outputs.reshape(
+            vision_outputs.shape[0],
+            config.vision_config.hidden_size,
+            config.vision_soft_tokens_per_image,
+        ).transpose(0, 2, 1)
+
+        # Normalize and embed the soft tokens into language model space.
+        vision_outputs *= config.vision_config.hidden_size**0.5
+        return embed_vision(inputs_embeds=vision_outputs)
+
+    def construct_special_modality_mask(
+        self, input_ids, inputs_embeds, token_id, modality="image"
+    ):
+        if input_ids is None:
+            embed_fn = (
+                self.embed_audio
+                if modality == "audio"
+                else self.language_model.model.embed_tokens
+            )
+            special_modality_mask = inputs_embeds == embed_fn(
+                input_ids=mx.array([token_id])
+            )
+        else:
+            special_modality_mask = mx.expand_dims(input_ids == token_id, -1)
+            special_modality_mask = mx.broadcast_to(
+                special_modality_mask, inputs_embeds.shape
+            )
+        return special_modality_mask
+
+    @staticmethod
+    def merge_multimodal_and_text(
+        inputs_embeds, features, special_modality_mask, modality="image"
+    ):
+        # Count special tokens by summing the mask
+        modality_tokens_in_text = special_modality_mask.sum()
+        feature_tokens = features.size
+
+        if modality_tokens_in_text != feature_tokens:
+            raise ValueError(
+                f"Number of {modality}s does not match number of special {modality} tokens in the input text. "
+                f"Got {modality_tokens_in_text} {modality} tokens in the text and "
+                f"{feature_tokens} tokens from {modality} embeddings."
+            )
+        features = features.astype(inputs_embeds.dtype)
+
+        inputs_embeds = masked_scatter(inputs_embeds, special_modality_mask, features)
+        return inputs_embeds
+
+    def __call__(
+        self,
+        input_ids: mx.array,
+        pixel_values: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[mx.array] = None,
+        **kwargs,
+    ):
+        # Audio features
+        input_features = kwargs.pop("input_features", None)
+        input_features_mask = kwargs.pop("input_features_mask", None)
+        inputs_embeds, per_layer_inputs = self.get_input_embeddings(
+            input_ids=input_ids,
+            pixel_values=pixel_values,
+            input_features=input_features,
+            input_features_mask=input_features_mask,
+            **kwargs,
+        )
+
+        logits = self.language_model(
+            input_ids=None,
+            cache=cache,
+            inputs_embeds=inputs_embeds,
+            per_layer_inputs=per_layer_inputs,
+        )
+        return logits
+
+    def sanitize(self, weights):
+        sanitized_weights = {}
+        for k, v in weights.items():
+            # if "vision_tower" not in k and "embed_vision" not in k:
+            if k.startswith("model."):
+                sanitized_weights[".".join(k.split(".")[1:])] = v
+            else:
+                sanitized_weights[k] = v
+        return sanitized_weights
+
+    @property
+    def layers(self):
+        return self.language_model.model.layers