nexaai 1.0.4rc16__cp310-cp310-macosx_14_0_universal2.whl → 1.0.6__cp310-cp310-macosx_14_0_universal2.whl

nexaai/_version.py

@@ -1,4 +1,4 @@
 # This file is generated by CMake from _version.py.in
 # Do not modify this file manually - it will be overwritten
 
-__version__ = "1.0.4-rc16" 
+__version__ = "1.0.6" 

nexaai/common.py

@@ -59,10 +59,45 @@ class ModelConfig:
 
 @dataclass(frozen=True) # Read-only
 class ProfilingData:
-    start_time: int
-    end_time: int
-    prompt_start_time: int = None
-    prompt_end_time: int = None
-    decode_start_time: int = None
-    decode_ent_time: int = None
-    first_token_time: int = None
+    """Profiling data structure for LLM/VLM performance metrics."""
+    ttft: int = 0             # Time to first token (us)
+    prompt_time: int = 0      # Prompt processing time (us)
+    decode_time: int = 0      # Token generation time (us)
+    prompt_tokens: int = 0    # Number of prompt tokens
+    generated_tokens: int = 0  # Number of generated tokens
+    audio_duration: int = 0   # Audio duration (us)
+    prefill_speed: float = 0.0  # Prefill speed (tokens/sec)
+    decoding_speed: float = 0.0 # Decoding speed (tokens/sec)
+    real_time_factor: float = 0.0 # Real-Time Factor (RTF)
+    stop_reason: str = ""     # Stop reason: "eos", "length", "user", "stop_sequence"
+    
+    @classmethod
+    def from_dict(cls, data: dict) -> "ProfilingData":
+        """Create ProfilingData from dictionary."""
+        return cls(
+            ttft=data.get("ttft", 0),
+            prompt_time=data.get("prompt_time", 0),
+            decode_time=data.get("decode_time", 0),
+            prompt_tokens=data.get("prompt_tokens", 0),
+            generated_tokens=data.get("generated_tokens", 0),
+            audio_duration=data.get("audio_duration", 0),
+            prefill_speed=data.get("prefill_speed", 0.0),
+            decoding_speed=data.get("decoding_speed", 0.0),
+            real_time_factor=data.get("real_time_factor", 0.0),
+            stop_reason=data.get("stop_reason", "")
+        )
+    
+    def to_dict(self) -> dict:
+        """Convert to dictionary."""
+        return {
+            "ttft": self.ttft,
+            "prompt_time": self.prompt_time,
+            "decode_time": self.decode_time,
+            "prompt_tokens": self.prompt_tokens,
+            "generated_tokens": self.generated_tokens,
+            "audio_duration": self.audio_duration,
+            "prefill_speed": self.prefill_speed,
+            "decoding_speed": self.decoding_speed,
+            "real_time_factor": self.real_time_factor,
+            "stop_reason": self.stop_reason
+        }

nexaai/embedder_impl/mlx_embedder_impl.py

@@ -3,7 +3,7 @@ import numpy as np
 
 from nexaai.common import PluginID
 from nexaai.embedder import Embedder, EmbeddingConfig
-from nexaai.mlx_backend.embedding.interface import Embedder as MLXEmbedderInterface
+from nexaai.mlx_backend.embedding.interface import create_embedder
 from nexaai.mlx_backend.ml import ModelConfig as MLXModelConfig, SamplerConfig as MLXSamplerConfig, GenerationConfig as MLXGenerationConfig, EmbeddingConfig
 
 
@@ -27,11 +27,12 @@ class MLXEmbedderImpl(Embedder):
             MLXEmbedderImpl instance
         """
         try:
-            # MLX interface is already imported
-            
-            # Create instance and load MLX embedder
+            # Create instance
             instance = cls()
-            instance._mlx_embedder = MLXEmbedderInterface(
+            
+            # Use the factory function to create the appropriate embedder based on model type
+            # This will automatically detect if it's JinaV2 or generic model and route correctly
+            instance._mlx_embedder = create_embedder(
                 model_path=model_path,
                 tokenizer_path=tokenizer_file
             )

nexaai/llm.py

@@ -4,7 +4,7 @@ import queue
 import threading
 
 from nexaai.common import ModelConfig, GenerationConfig, ChatMessage, PluginID
-from nexaai.base import BaseModel
+from nexaai.base import BaseModel, ProfilingData
 
 class LLM(BaseModel):
     def __init__(self, m_cfg: ModelConfig = ModelConfig()):
@@ -63,6 +63,10 @@ class LLM(BaseModel):
         """
         pass
 
+    def get_profiling_data(self) -> Optional[ProfilingData]:
+        """Get profiling data from the last generation."""  
+        pass
+
     @abstractmethod
     def save_kv_cache(self, path: str):
         """

nexaai/llm_impl/mlx_llm_impl.py

@@ -1,5 +1,6 @@
 from typing import Generator, Optional, Any, Sequence, Union
 
+from nexaai.base import ProfilingData
 from nexaai.common import ModelConfig, GenerationConfig, ChatMessage, PluginID
 from nexaai.llm import LLM
 from nexaai.mlx_backend.llm.interface import LLM as MLXLLMInterface
@@ -215,6 +216,12 @@ class MLXLLMImpl(LLM):
         except Exception as e:
             raise RuntimeError(f"Failed to generate text: {str(e)}")
 
+    def get_profiling_data(self) -> Optional[ProfilingData]:
+        """Get profiling data from the last generation."""
+        if not self._mlx_llm:
+            raise RuntimeError("MLX LLM not loaded")
+        return self._mlx_llm.get_profiling_data()
+
     def save_kv_cache(self, path: str):
         """
         Save the key-value cache to the file.

nexaai/llm_impl/pybind_llm_impl.py

@@ -2,6 +2,7 @@ from typing import Generator, Optional, Union
 import queue
 import threading
 
+from nexaai.base import ProfilingData
 from nexaai.common import ModelConfig, GenerationConfig, ChatMessage, PluginID
 from nexaai.binds import llm_bind, common_bind
 from nexaai.runtime import _ensure_runtime
@@ -13,6 +14,7 @@ class PyBindLLMImpl(LLM):
         """Private constructor, should not be called directly."""
         super().__init__(m_cfg)
         self._handle = handle  # This is a py::capsule
+        self._profiling_data = None
 
     @classmethod
     def _load_from(cls,
@@ -97,13 +99,14 @@ class PyBindLLMImpl(LLM):
         # Run generation in thread
         def generate():
             try:
-                llm_bind.ml_llm_generate(
+                result = llm_bind.ml_llm_generate(
                     handle=self._handle,
                     prompt=prompt,
                     config=config,
                     on_token=on_token,
                     user_data=None
                 )
+                self._profiling_data = ProfilingData.from_dict(result.get("profile_data", {}))
             except Exception as e:
                 exception_container[0] = e
             finally:
@@ -145,8 +148,14 @@ class PyBindLLMImpl(LLM):
             on_token=None,  # No callback for non-streaming
             user_data=None
         )
+
+        self._profiling_data = ProfilingData.from_dict(result.get("profile_data", {}))
         return result.get("text", "")
 
+    def get_profiling_data(self) -> Optional[ProfilingData]:
+        """Get profiling data."""
+        return self._profiling_data
+
     def save_kv_cache(self, path: str):
         """
         Save the key-value cache to the file.

nexaai/mlx_backend/embedding/generate.py

@@ -23,11 +23,46 @@ from .modeling.nexa_jina_v2 import Model, ModelArgs
 from tokenizers import Tokenizer
 from huggingface_hub import snapshot_download
 
-def load_model(model_id):
+# Try to import mlx_embeddings for general embedding support
+try:
+    import mlx_embeddings
+    MLX_EMBEDDINGS_AVAILABLE = True
+except ImportError:
+    MLX_EMBEDDINGS_AVAILABLE = False
+    # Suppress warning during import to avoid interfering with C++ tests
+    # The warning will be shown when actually trying to use mlx_embeddings functionality
+    pass
+
+def detect_model_type(model_path):
+    """Detect if the model is Jina V2 or generic mlx_embeddings model."""
+    config_path = os.path.join(model_path, "config.json") if os.path.isdir(model_path) else f"{model_path}/config.json"
+    
+    if not os.path.exists(config_path):
+        # Try default modelfiles directory
+        config_path = f"{curr_dir}/modelfiles/config.json"
+        if not os.path.exists(config_path):
+            return "generic"
+    
+    try:
+        with open(config_path, "r") as f:
+            config = json.load(f)
+        
+        # Check if it's a Jina V2 model
+        architectures = config.get("architectures", [])
+        if "JinaBertModel" in architectures:
+            return "jina_v2"
+        
+        return "generic"
+    except Exception:
+        return "generic"
+
+# ========== Jina V2 Direct Implementation ==========
+
+def load_jina_model(model_id):
     """Initialize and load the Jina V2 model with FP16 weights."""
     # Load configuration from config.json
     if not os.path.exists(f"{curr_dir}/modelfiles/config.json"):
-        print(f"📥 Downloading model {model_id}...")
+        print(f"📥 Downloading Jina V2 model {model_id}...")
         
         # Ensure modelfiles directory exists
         os.makedirs(f"{curr_dir}/modelfiles", exist_ok=True)
@@ -82,15 +117,15 @@ def load_model(model_id):
     
     return model
 
-def load_tokenizer():
-    """Load and configure the tokenizer."""
+def load_jina_tokenizer():
+    """Load and configure the tokenizer for Jina V2."""
     tokenizer = Tokenizer.from_file(f"{curr_dir}/modelfiles/tokenizer.json")
     tokenizer.enable_padding(pad_id=0, pad_token="[PAD]")
     tokenizer.enable_truncation(max_length=512)
     return tokenizer
 
-def encode_text(model, tokenizer, text):
-    """Encode a single text and return its embedding."""
+def encode_jina_text(model, tokenizer, text):
+    """Encode a single text using Jina V2 and return its embedding."""
     # Tokenize the text
     encoding = tokenizer.encode(text)
     
@@ -113,18 +148,186 @@ def encode_text(model, tokenizer, text):
     
     return embeddings
 
+# ========== MLX Embeddings Direct Implementation ==========
+
+def load_mlx_embeddings_model(model_id):
+    """Load model using mlx_embeddings package."""
+    if not MLX_EMBEDDINGS_AVAILABLE:
+        print("Warning: mlx_embeddings not available. Please install it to use general embedding models.")
+        raise ImportError("mlx_embeddings package is not available. Please install it first.")
+    
+    # Download model if needed
+    model_path = f"{curr_dir}/modelfiles"
+    
+    if not os.path.exists(f"{model_path}/config.json"):
+        print(f"📥 Downloading model {model_id}...")
+        os.makedirs(model_path, exist_ok=True)
+        
+        try:
+            snapshot_download(
+                repo_id=model_id,
+                local_dir=model_path,
+                resume_download=True,
+                local_dir_use_symlinks=False
+            )
+            print("✅ Model download completed!")
+        except Exception as e:
+            print(f"❌ Failed to download model: {e}")
+            raise
+    
+    # Load model and tokenizer using mlx_embeddings
+    model, tokenizer = mlx_embeddings.load(model_path)
+    return model, tokenizer
+
+def encode_mlx_embeddings_text(model, tokenizer, texts, model_path=None):
+    """Generate embeddings using mlx_embeddings."""
+    if isinstance(texts, str):
+        texts = [texts]
+    
+    # Check if this is a Gemma3TextModel by checking config
+    # WORKAROUND: Gemma3TextModel has a bug where it expects 'inputs' as positional arg
+    # but mlx_embeddings.generate passes 'input_ids' as keyword arg
+    # See: https://github.com/ml-explore/mlx-examples/issues/... (bug report pending)
+    is_gemma = False
+    if model_path:
+        config_path = os.path.join(model_path, "config.json") if os.path.isdir(model_path) else f"{model_path}/config.json"
+    else:
+        config_path = f"{curr_dir}/modelfiles/config.json"
+    
+    if os.path.exists(config_path):
+        try:
+            with open(config_path, "r") as f:
+                config = json.load(f)
+                architectures = config.get("architectures", [])
+                is_gemma = "Gemma3TextModel" in architectures
+        except Exception:
+            pass
+    
+    if is_gemma:
+        # HARDCODED WORKAROUND for Gemma3TextModel bug
+        # Use direct tokenization and model call instead of mlx_embeddings.generate
+        # This avoids the bug where generate passes 'input_ids' as keyword arg
+        # but Gemma3TextModel.__call__ expects 'inputs' as positional arg
+        
+        # Tokenize using batch_encode_plus for Gemma models
+        encoded_input = tokenizer.batch_encode_plus(
+            texts,
+            padding=True,
+            truncation=True,
+            return_tensors='mlx',
+            max_length=512
+        )
+        
+        # Get input tensors
+        input_ids = encoded_input['input_ids']
+        attention_mask = encoded_input.get('attention_mask', None)
+        
+        # Call model with positional input_ids and keyword attention_mask
+        # This matches Gemma3TextModel's expected signature:
+        # def __call__(self, inputs: mx.array, attention_mask: Optional[mx.array] = None)
+        output = model(input_ids, attention_mask=attention_mask)
+        
+        # Get the normalized embeddings
+        return output.text_embeds
+    else:
+        # Normal path for non-Gemma models
+        # Use standard mlx_embeddings.generate approach
+        output = mlx_embeddings.generate(
+            model,
+            tokenizer,
+            texts=texts,
+            max_length=512,
+            padding=True,
+            truncation=True
+        )
+        
+        return output.text_embeds
+
 def main(model_id):
     """Main function to handle user input and generate embeddings."""
     
-    # Load model and tokenizer
-    model = load_model(model_id)
-    tokenizer = load_tokenizer()
-    user_input = "Hello, how are you?"
-    embedding = encode_text(model, tokenizer, user_input)
-    print(f"Embedding shape: {embedding.shape}")
-    print(f"Embedding sample values: {embedding.flatten()[:5].tolist()}")
-    print(f"Embedding min: {embedding.min()}, Max: {embedding.max()}, Mean: {embedding.mean()}, Std: {embedding.std()}")
+    print(f"🔍 Loading model: {model_id}")
+    
+    # Detect model type
+    model_type = detect_model_type(f"{curr_dir}/modelfiles")
+    
+    # First try to download/check if model exists
+    if not os.path.exists(f"{curr_dir}/modelfiles/config.json"):
+        # Download the model first to detect its type
+        print(f"Model not found locally. Downloading...")
+        os.makedirs(f"{curr_dir}/modelfiles", exist_ok=True)
+        try:
+            snapshot_download(
+                repo_id=model_id, 
+                local_dir=f"{curr_dir}/modelfiles",
+                resume_download=True,
+                local_dir_use_symlinks=False
+            )
+            print("✅ Model download completed!")
+            # Re-detect model type after download
+            model_type = detect_model_type(f"{curr_dir}/modelfiles")
+        except Exception as e:
+            print(f"❌ Failed to download model: {e}")
+            raise
+    
+    print(f"📦 Detected model type: {model_type}")
+    
+    # Test texts
+    test_texts = [
+        "Hello, how are you?",
+        "What is machine learning?",
+        "The weather is nice today."
+    ]
+    
+    if model_type == "jina_v2":
+        print("Using Jina V2 direct implementation")
+        
+        # Load Jina V2 model
+        model = load_jina_model(model_id)
+        tokenizer = load_jina_tokenizer()
+        
+        print("\nGenerating embeddings for test texts:")
+        for text in test_texts:
+            embedding = encode_jina_text(model, tokenizer, text)
+            print(f"\nText: '{text}'")
+            print(f"  Embedding shape: {embedding.shape}")
+            print(f"  Sample values (first 5): {embedding.flatten()[:5].tolist()}")
+            print(f"  Stats - Min: {embedding.min():.4f}, Max: {embedding.max():.4f}, Mean: {embedding.mean():.4f}")
+    
+    else:
+        print("Using mlx_embeddings direct implementation")
+        
+        if not MLX_EMBEDDINGS_AVAILABLE:
+            print("❌ mlx_embeddings is not installed. Please install it to use generic models.")
+            return
+        
+        # Load generic model using mlx_embeddings
+        model, tokenizer = load_mlx_embeddings_model(model_id)
+        
+        print("\nGenerating embeddings for test texts:")
+        # Pass model_path to handle Gemma workaround if needed
+        embeddings = encode_mlx_embeddings_text(model, tokenizer, test_texts, model_path=f"{curr_dir}/modelfiles")
+        
+        for i, text in enumerate(test_texts):
+            embedding = embeddings[i]
+            print(f"\nText: '{text}'")
+            print(f"  Embedding shape: {embedding.shape}")
+            print(f"  Sample values (first 5): {embedding[:5].tolist()}")
+            
+            # Calculate stats
+            emb_array = mx.array(embedding) if not isinstance(embedding, mx.array) else embedding
+            print(f"  Stats - Min: {emb_array.min():.4f}, Max: {emb_array.max():.4f}, Mean: {emb_array.mean():.4f}")
+    
+    print("\n✅ Direct embedding generation completed!")
 
 if __name__ == "__main__":
-    model_id = "nexaml/jina-v2-fp16-mlx"
-    main(model_id)
+    import argparse
+    parser = argparse.ArgumentParser(description="Generate embeddings using direct implementation")
+    parser.add_argument(
+        "--model_id", 
+        type=str, 
+        default="nexaml/jina-v2-fp16-mlx",
+        help="Model ID from Hugging Face Hub (e.g., 'nexaml/jina-v2-fp16-mlx' or 'mlx-community/embeddinggemma-300m-bf16')"
+    )
+    args = parser.parse_args()
+    main(args.model_id)