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/molmo/language.py

@@ -1,243 +0,0 @@
-import inspect
-from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union
-
-import mlx.core as mx
-import mlx.nn as nn
-
-from ..base import (
-    LanguageModelOutput,
-    create_attention_mask,
-    scaled_dot_product_attention,
-)
-from ..cache import KVCache
-
-
-@dataclass
-class TextConfig:
-    model_type: str = "molmo"
-    max_position_embeddings: int = 4096
-    d_model: int = 3584
-    n_heads: int = 28
-    n_kv_heads: int = 4
-    n_layers: int = 28
-    mlp_ratio: int = 4
-    max_sequence_length: int = 1024
-    act_output_multiplier: int = 0.5
-    mlp_hidden_size: int = 37888
-    vocab_size: int = 152064
-    embedding_size: Optional[int] = 152064
-    additional_vocab_size: Optional[int] = None
-    attention_dropout: float = 0.1
-    residual_dropout: float = 0.1
-    embedding_dropout: float = 0.1
-    layer_norm_eps: float = 1e-5
-    initializer_range: float = 0.02
-    pad_token_id: int = -1
-    rope: bool = True
-    rope_theta: float = 1000000.0
-    weight_tying: bool = False
-    rope_full_precision: bool = True
-    rope_impl: str = "interleave"
-    additional_vocab_size: Optional[int] = 128
-
-    @classmethod
-    def from_dict(cls, params):
-        return cls(
-            **{
-                k: v
-                for k, v in params.items()
-                if k in inspect.signature(cls).parameters
-            }
-        )
-
-
-class SwiGLU(nn.Module):
-    def __call__(self, x: mx.array) -> mx.array:
-        x, gate = mx.split(x, 2, axis=-1)
-        return nn.silu(gate) * x
-
-
-class MolmoBlock(nn.Module):
-    def __init__(self, config: TextConfig):
-        super().__init__()
-        self.attn_out = nn.Linear(config.d_model, config.d_model, bias=False)
-        self.ff_out = nn.Linear(
-            int(config.act_output_multiplier * config.mlp_hidden_size),
-            config.d_model,
-            bias=False,
-        )
-        self.attn_norm = nn.RMSNorm(config.d_model, eps=config.layer_norm_eps)
-        self.ff_norm = nn.RMSNorm(config.d_model, eps=config.layer_norm_eps)
-        self.ff_proj = nn.Linear(config.d_model, config.mlp_hidden_size, bias=False)
-        head_dim = config.d_model // config.n_heads
-        self.rotary_emb = nn.RoPE(head_dim, base=config.rope_theta)
-        self.scale = head_dim**-0.5
-        self.n_heads = config.n_heads
-        self.n_kv_heads = config.n_kv_heads
-        self.fused_dims = (
-            config.d_model,
-            config.n_kv_heads * head_dim,
-            config.n_kv_heads * head_dim,
-        )
-        self.att_proj = nn.Linear(config.d_model, sum(self.fused_dims), bias=True)
-        self.act = SwiGLU()
-
-    def __call__(self, x, mask=None, cache=None):
-        batch_size, seq_len, D = x.shape
-        attn_in = self.attn_norm(x)
-
-        qkv = self.att_proj(attn_in)
-
-        q, k, v = mx.split(
-            qkv, [self.fused_dims[0], self.fused_dims[0] + self.fused_dims[1]], axis=-1
-        )
-
-        q = q.reshape(batch_size, seq_len, self.n_heads, D // self.n_heads).transpose(
-            0, 2, 1, 3
-        )
-        k = k.reshape(
-            batch_size, seq_len, self.n_kv_heads, D // self.n_heads
-        ).transpose(0, 2, 1, 3)
-        v = v.reshape(
-            batch_size, seq_len, self.n_kv_heads, D // self.n_heads
-        ).transpose(0, 2, 1, 3)
-
-        if cache is not None:
-            q = self.rotary_emb(q, offset=cache.offset)
-            k = self.rotary_emb(k, offset=cache.offset)
-            k, v = cache.update_and_fetch(k, v)
-        else:
-            q = self.rotary_emb(q)
-            k = self.rotary_emb(k)
-
-        # Perform attention
-        att = scaled_dot_product_attention(q, k, v, cache, scale=self.scale, mask=mask)
-        att = att.transpose(0, 2, 1, 3).reshape(batch_size, seq_len, D)
-        att = self.attn_out(att)
-
-        # Add attention scores
-        # shape: (batch_size, seq_len, d_model)
-        x = x + att
-
-        # Feed-forward layer
-        og_x = x
-        x = self.ff_norm(x)
-        x = self.ff_proj(x)
-        x = self.act(x)
-        x = self.ff_out(x)
-        x = og_x + x
-
-        return x
-
-
-class Embedding(nn.Module):
-    def __init__(
-        self,
-        num_embeddings: int,
-        num_new_embeddings: int,
-        features: int,
-        initializer_range: float = 0.02,
-        new_embed_initializer_range: float = 0.02,
-    ):
-        super().__init__()
-        self.initializer_range = initializer_range
-        self.new_embed_initializer_range = new_embed_initializer_range
-
-        # Initialize embeddings
-        self.embedding = mx.random.normal(
-            (num_embeddings, features), scale=self.initializer_range
-        )
-        self.new_embedding = mx.random.normal(
-            (num_new_embeddings, features), scale=self.new_embed_initializer_range
-        )
-
-    def __call__(self, x: mx.array) -> mx.array:
-        return mx.concat([self.embedding, self.new_embedding], axis=0)[x]
-
-
-class Molmo(nn.Module):
-    def __init__(self, config: TextConfig):
-        super().__init__()
-        self.config = config
-
-        self.wte = Embedding(
-            config.embedding_size, config.additional_vocab_size, config.d_model
-        )
-
-        self.blocks = [MolmoBlock(config) for _ in range(config.n_layers)]
-
-        self.ln_f = nn.RMSNorm(config.d_model, eps=config.layer_norm_eps)
-
-        if not config.weight_tying:
-            self.ff_out = nn.Linear(config.d_model, config.vocab_size, bias=False)
-
-    def __call__(
-        self,
-        input_ids: mx.array,
-        inputs_embeds: Optional[mx.array] = None,
-        mask: Optional[mx.array] = None,
-        cache: Optional[KVCache] = None,
-    ) -> LanguageModelOutput:
-
-        if inputs_embeds is None:
-            h = self.wte(input_ids)
-        else:
-            h = inputs_embeds
-
-        if cache is None:
-            cache = [None] * self.config.n_layers
-
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-
-        for block, c in zip(self.blocks, cache):
-            h = block(h, mask, c)
-
-        h = self.ln_f(h)
-
-        if self.config.weight_tying:
-            logits = mx.matmul(h, self.wte.weight.T)
-        else:
-            logits = self.ff_out(h)
-
-        return LanguageModelOutput(logits=logits)
-
-
-class LanguageModel(nn.Module):
-    def __init__(self, config: TextConfig):
-        super().__init__()
-        self.config = config
-        self.model_type = config.model_type
-        if self.model_type != "molmo":
-            raise ValueError(
-                f"Model type {self.model_type} not supported. Currently only 'molmo' is supported"
-            )
-        self.model = Molmo(config)
-
-    def __call__(
-        self,
-        input_ids: mx.array,
-        inputs_embeds: Optional[mx.array] = None,
-        mask: Optional[mx.array] = None,
-        cache: Optional[KVCache] = None,
-    ) -> LanguageModelOutput:
-        outputs = self.model(input_ids, inputs_embeds, mask, cache)
-        return outputs
-
-    @staticmethod
-    def sanitize(weights):
-        # Remove unused precomputed rotary freqs
-        return {k: v for k, v in weights.items() if "rotary_emb.inv_freq" not in k}
-
-    @property
-    def layers(self):
-        return self.model.blocks
-
-    @property
-    def head_dim(self):
-        return self.config.d_model // self.config.n_heads
-
-    @property
-    def n_kv_heads(self):
-        return self.config.n_kv_heads

nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/molmo/molmo.py

@@ -1,133 +0,0 @@
-import glob
-import inspect
-import json
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Dict, List, Optional, Tuple, Union
-
-import mlx.core as mx
-import mlx.nn as nn
-import numpy as np
-from huggingface_hub import snapshot_download
-
-from .language import LanguageModel, TextConfig
-from .vision import VisionConfig, VisionModel
-
-
-@dataclass
-class ModelConfig:
-    text_config: TextConfig = field(default_factory=TextConfig)
-    vision_config: VisionConfig = field(default_factory=VisionConfig)
-    model_type: str = "molmo"
-    image_feature_dropout: float = 0.0
-    image_pooling_h: int = 2
-    image_pooling_w: int = 2
-    image_pooling_2d: str = "attention"
-    image_projector: str = "mlp"
-    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
-            }
-        )
-
-
-class Model(nn.Module):
-    def __init__(self, config: ModelConfig):
-        super().__init__()
-        self.config = config
-        self.language_model = LanguageModel(config.text_config)
-        self.vision_tower = VisionModel(config.vision_config)
-
-    @property
-    def layers(self):
-        return self.language_model.model.layers
-
-    def __call__(
-        self,
-        input_ids: mx.array,
-        pixel_values: mx.array,
-        mask: mx.array,
-        cache=None,
-        **kwargs,
-    ) -> Dict[str, Union[mx.array, List[Tuple[mx.array, mx.array]]]]:
-        if input_ids.ndim == 1:
-            input_ids = input_ids[None, :]
-
-        batch_size, seq_len = input_ids.shape
-
-        image_input_idx = kwargs.get("image_input_idx", None)
-        image_masks = kwargs.get("image_masks", None)
-
-        if pixel_values is not None:
-            assert (
-                image_masks is not None and image_input_idx is not None
-            ), "image_masks and image_input_idx must be provided when images are given"
-
-            dtype = self.vision_tower.image_vit.patch_embedding.weight.dtype
-            pixel_values = pixel_values.astype(dtype)
-
-            # Process images
-            if pixel_values.ndim == 3:
-                pixel_values = mx.expand_dims(pixel_values, 0)
-                image_masks = (
-                    mx.expand_dims(image_masks, 0) if image_masks is not None else None
-                )
-                image_input_idx = (
-                    mx.expand_dims(image_input_idx, 0)
-                    if image_input_idx is not None
-                    else None
-                )
-
-            image_features, cls_embed = self.vision_tower(pixel_values, image_masks)
-
-            # Insert image features into the input embeddings
-            num_image, num_patch = image_features.shape[1:3]
-
-            assert image_input_idx.shape == (
-                batch_size,
-                num_image,
-                num_patch,
-            ), f"image_input_idx.shape: {image_input_idx.shape}, expected: {(batch_size, num_image, num_patch)}"
-
-            # Insert image features into the input embeddings
-            image_features = image_features.reshape(
-                batch_size, num_image * num_patch, -1
-            )
-            image_input_idx = image_input_idx.reshape(batch_size, num_image * num_patch)
-
-            valid = np.where(image_input_idx >= 0)[0].tolist()
-            batch_idx = mx.arange(batch_size)
-            batch_idx = mx.tile(batch_idx[:, None], [1, image_features.shape[1]])
-
-            input_embeddings = self.language_model.model.wte(input_ids)
-            input_embeddings[
-                batch_idx[valid], image_input_idx[valid]
-            ] += image_features[valid]
-        else:
-            input_embeddings = None
-
-        # Forward pass through the language model
-        logits = self.language_model(
-            input_ids,
-            inputs_embeds=input_embeddings,
-            mask=mask,
-            cache=cache,
-        )
-
-        return logits
-
-    def sanitize(self, weights):
-        def transform_key(key):
-            if "model.transformer" in key:
-                key = key.replace("model.transformer", "language_model.model")
-            if "model.vision_backbone" in key:
-                key = key.replace("model.vision_backbone", "vision_tower")
-            return key
-
-        return {transform_key(k): v for k, v in weights.items()}