nexaai 1.0.9__cp310-cp310-macosx_14_0_universal2.whl → 1.0.11rc1__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.9" 
+__version__ = "1.0.11-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 create_embedder
+from nexaai.mlx_backend.embedding.interface import Embedder as MLXEmbedderInterface
 from nexaai.mlx_backend.ml import ModelConfig as MLXModelConfig, SamplerConfig as MLXSamplerConfig, GenerationConfig as MLXGenerationConfig, EmbeddingConfig
 
 
@@ -27,12 +27,11 @@ class MLXEmbedderImpl(Embedder):
             MLXEmbedderImpl instance
         """
         try:
-            # Create instance
-            instance = cls()
+            # MLX interface is already imported
             
-            # 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(
+            # Create instance and load MLX embedder
+            instance = cls()
+            instance._mlx_embedder = MLXEmbedderInterface(
                 model_path=model_path,
                 tokenizer_path=tokenizer_file
             )

nexaai/mlx_backend/embedding/generate.py

@@ -23,46 +23,11 @@ from .modeling.nexa_jina_v2 import Model, ModelArgs
 from tokenizers import Tokenizer
 from huggingface_hub import snapshot_download
 
-# 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):
+def load_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 Jina V2 model {model_id}...")
+        print(f"📥 Downloading model {model_id}...")
         
         # Ensure modelfiles directory exists
         os.makedirs(f"{curr_dir}/modelfiles", exist_ok=True)
@@ -117,15 +82,15 @@ def load_jina_model(model_id):
     
     return model
 
-def load_jina_tokenizer():
-    """Load and configure the tokenizer for Jina V2."""
+def load_tokenizer():
+    """Load and configure the tokenizer."""
     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_jina_text(model, tokenizer, text):
-    """Encode a single text using Jina V2 and return its embedding."""
+def encode_text(model, tokenizer, text):
+    """Encode a single text and return its embedding."""
     # Tokenize the text
     encoding = tokenizer.encode(text)
     
@@ -148,186 +113,18 @@ def encode_jina_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."""
     
-    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!")
+    # 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()}")
 
 if __name__ == "__main__":
-    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)
+    model_id = "nexaml/jina-v2-fp16-mlx"
+    main(model_id)