nexaai 1.0.19rc7__cp310-cp310-macosx_14_0_universal2.whl → 1.0.19rc9__cp310-cp310-macosx_14_0_universal2.whl

nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3n/config.py

@@ -1,139 +0,0 @@
-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/binds/nexa_mlx/py-lib/vlm/modeling/models/gemma3n/gemma3n.py

@@ -1,322 +0,0 @@
-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