nexaai 1.0.11rc2__cp310-cp310-macosx_13_0_x86_64.whl → 1.0.12rc1__cp310-cp310-macosx_13_0_x86_64.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.11-rc2" 
+__version__ = "1.0.12-rc1" 

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/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)