nexaai 1.0.11rc2__cp310-cp310-macosx_14_0_universal2.whl → 1.0.12rc1__cp310-cp310-macosx_14_0_universal2.whl

nexaai/mlx_backend/embedding/interface.py

@@ -20,11 +20,16 @@ import mlx.core as mx
 import numpy as np
 from pathlib import Path
 from typing import Any, List, Optional, Sequence
+from abc import ABC, abstractmethod
 
 # Import necessary modules 
 from tokenizers import Tokenizer
 
 # Import from ml.py for API alignment
+import sys
+from pathlib import Path as PathLib
+sys.path.insert(0, str(PathLib(__file__).parent.parent))
+
 from ml import (
     Embedder as BaseEmbedder,
     EmbeddingConfig,
@@ -34,13 +39,24 @@ from ml import (
 # Import profiling module
 from profiling import ProfilingMixin, StopReason
 
-# Import the model implementation
-from .modeling.nexa_jina_v2 import Model, ModelArgs
+# Import the model implementation for Jina
+try:
+    from .modeling.nexa_jina_v2 import Model, ModelArgs
+except ImportError:
+    # Fallback for when module is run directly
+    from modeling.nexa_jina_v2 import Model, ModelArgs
+
+# Import mlx_embeddings for general embedding support
+try:
+    import mlx_embeddings
+    MLX_EMBEDDINGS_AVAILABLE = True
+except ImportError:
+    MLX_EMBEDDINGS_AVAILABLE = False
 
 
-class Embedder(BaseEmbedder, ProfilingMixin):
+class BaseMLXEmbedder(BaseEmbedder, ProfilingMixin, ABC):
     """
-    Embedder interface for MLX embedding models.
+    Abstract base embedder interface for MLX embedding models.
     API aligned with ml.py Embedder abstract base class.
     """
 
@@ -64,7 +80,7 @@ class Embedder(BaseEmbedder, ProfilingMixin):
         
         self.model_path = model_path
         self.tokenizer_path = tokenizer_path
-        self.device = device if device is not None else "cpu" # TODO: This device field is never used
+        self.device = device if device is not None else "cpu"
         
         # Initialize model and tokenizer as None
         self.model = None
@@ -78,6 +94,69 @@ class Embedder(BaseEmbedder, ProfilingMixin):
         self.config = None
         self.reset_profiling()
 
+    @abstractmethod
+    def load_model(self, model_path: PathType) -> bool:
+        """Load model from path."""
+        pass
+
+    def close(self) -> None:
+        """Close the model."""
+        self.destroy()
+
+    @abstractmethod
+    def embed(
+        self,
+        texts: Sequence[str],
+        config: Optional[EmbeddingConfig] = None,
+        clear_cache: bool = True,
+    ) -> List[List[float]]:
+        """Generate embeddings for texts."""
+        pass
+
+    @abstractmethod
+    def embedding_dim(self) -> int:
+        """Get embedding dimension."""
+        pass
+
+    def set_lora(self, lora_id: int) -> None:
+        """Set active LoRA adapter. (Disabled for embedding models)"""
+        raise NotImplementedError("LoRA is not supported for embedding models")
+
+    def add_lora(self, lora_path: PathType) -> int:
+        """Add LoRA adapter and return its ID. (Disabled for embedding models)"""
+        raise NotImplementedError("LoRA is not supported for embedding models")
+
+    def remove_lora(self, lora_id: int) -> None:
+        """Remove LoRA adapter. (Disabled for embedding models)"""
+        raise NotImplementedError("LoRA is not supported for embedding models")
+
+    def list_loras(self) -> List[int]:
+        """List available LoRA adapters. (Disabled for embedding models)"""
+        raise NotImplementedError("LoRA is not supported for embedding models")
+
+    def _normalize_embedding(self, embedding: List[float], method: str) -> List[float]:
+        """Normalize embedding using specified method."""
+        if method == "none":
+            return embedding
+        
+        embedding_array = np.array(embedding)
+        
+        if method == "l2":
+            norm = np.linalg.norm(embedding_array)
+            if norm > 0:
+                embedding_array = embedding_array / norm
+        elif method == "mean":
+            mean_val = np.mean(embedding_array)
+            embedding_array = embedding_array - mean_val
+        
+        return embedding_array.tolist()
+
+
+class JinaV2Embedder(BaseMLXEmbedder):
+    """
+    Embedder implementation specifically for Jina V2 models.
+    """
+
     def load_model(self, model_path: PathType) -> bool:
         """Load model from path."""
         try:
@@ -97,10 +176,6 @@ class Embedder(BaseEmbedder, ProfilingMixin):
             print(f"Failed to load model: {e}")
             return False
 
-    def close(self) -> None:
-        """Close the model."""
-        self.destroy()
-
     def embed(
         self,
         texts: Sequence[str],
@@ -158,22 +233,6 @@ class Embedder(BaseEmbedder, ProfilingMixin):
             return 768  # Default dimension for Jina v2
         return self.config.hidden_size
 
-    def set_lora(self, lora_id: int) -> None:
-        """Set active LoRA adapter. (Disabled for embedding models)"""
-        raise NotImplementedError("LoRA is not supported for embedding models")
-
-    def add_lora(self, lora_path: PathType) -> int:
-        """Add LoRA adapter and return its ID. (Disabled for embedding models)"""
-        raise NotImplementedError("LoRA is not supported for embedding models")
-
-    def remove_lora(self, lora_id: int) -> None:
-        """Remove LoRA adapter. (Disabled for embedding models)"""
-        raise NotImplementedError("LoRA is not supported for embedding models")
-
-    def list_loras(self) -> List[int]:
-        """List available LoRA adapters. (Disabled for embedding models)"""
-        raise NotImplementedError("LoRA is not supported for embedding models")
-    
     def _load_jina_model(self, model_dir: str) -> Model:
         """Initialize and load the Jina V2 model with FP16 weights."""
         
@@ -281,22 +340,267 @@ class Embedder(BaseEmbedder, ProfilingMixin):
         
         return embedding_list
 
-    def _normalize_embedding(self, embedding: List[float], method: str) -> List[float]:
-        """Normalize embedding using specified method."""
-        if method == "none":
-            return embedding
+
+class MlxEmbeddingEmbedder(BaseMLXEmbedder):
+    """
+    Embedder implementation using mlx_embeddings package for general embedding models.
+    """
+
+    def load_model(self, model_path: PathType) -> bool:
+        """Load model from path using mlx_embeddings."""
+        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.")
         
-        embedding_array = np.array(embedding)
+        try:
+            # Use the provided model_path or fall back to instance path
+            if model_path:
+                if os.path.isfile(model_path):
+                    model_path = os.path.dirname(model_path)
+                self.model_path = model_path
+            
+            # Load model and tokenizer using mlx_embeddings
+            self.model, self.tokenizer = mlx_embeddings.load(self.model_path)
+            
+            # Load config to get dimensions
+            config_path = os.path.join(self.model_path, "config.json")
+            if os.path.exists(config_path):
+                with open(config_path, "r") as f:
+                    self.config = json.load(f)
+            
+            return True
+        except Exception as e:
+            print(f"Failed to load model: {e}")
+            return False
+
+    def embed(
+        self,
+        texts: Sequence[str],
+        config: Optional[EmbeddingConfig] = None,
+        clear_cache: bool = True,
+    ) -> List[List[float]]:
+        """Generate embeddings for texts using mlx_embeddings."""
+        if self.model is None or self.tokenizer is None:
+            raise RuntimeError("Model not loaded. Call load_model() first.")
         
-        if method == "l2":
-            norm = np.linalg.norm(embedding_array)
-            if norm > 0:
-                embedding_array = embedding_array / norm
-        elif method == "mean":
-            mean_val = np.mean(embedding_array)
-            embedding_array = embedding_array - mean_val
+        if config is None:
+            config = EmbeddingConfig()
         
-        return embedding_array.tolist()
+        # Start profiling
+        self._start_profiling()
+        
+        try:
+            # Calculate total tokens for profiling
+            if hasattr(self.tokenizer, 'encode'):
+                total_tokens = sum(len(self.tokenizer.encode(text)) for text in texts)
+            else:
+                # For tokenizers that don't have simple encode method
+                total_tokens = len(texts) * 50  # Rough estimate
+            
+            self._update_prompt_tokens(total_tokens)
+            
+            # End prompt processing, start decode
+            self._prompt_end()
+            self._decode_start()
+            
+            # Check if this is a Gemma3TextModel
+            # 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 self.config and "architectures" in self.config:
+                architectures = self.config.get("architectures", [])
+                is_gemma = "Gemma3TextModel" in architectures
+            
+            if is_gemma:
+                # HARDCODED WORKAROUND for Gemma3TextModel bug
+                # Use direct tokenization and model call instead of mlx_embeddings.generate
+                max_length = config.max_length if hasattr(config, 'max_length') else 512
+                
+                # Tokenize using batch_encode_plus
+                encoded_input = self.tokenizer.batch_encode_plus(
+                    list(texts),
+                    padding=True,
+                    truncation=True,
+                    return_tensors='mlx',
+                    max_length=max_length
+                )
+                
+                # 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
+                output = self.model(input_ids, attention_mask=attention_mask)
+                
+                # Extract embeddings
+                embeddings_tensor = output.text_embeds
+            else:
+                # Normal path for non-Gemma models
+                # Generate embeddings using mlx_embeddings standard approach
+                output = mlx_embeddings.generate(
+                    self.model,
+                    self.tokenizer,
+                    texts=list(texts),
+                    max_length=config.max_length if hasattr(config, 'max_length') else 512,
+                    padding=True,
+                    truncation=True
+                )
+                
+                # Extract embeddings
+                embeddings_tensor = output.text_embeds
+            
+            # Convert to list format
+            embeddings = []
+            for i in range(embeddings_tensor.shape[0]):
+                embedding = embeddings_tensor[i].tolist()
+                
+                # Apply normalization if requested
+                if config.normalize:
+                    embedding = self._normalize_embedding(embedding, config.normalize_method)
+                
+                embeddings.append(embedding)
+            
+            if clear_cache:
+                mx.clear_cache()
+            
+            # End timing and finalize profiling data
+            self._update_generated_tokens(0)  # No generation in embedding
+            self._set_stop_reason(StopReason.ML_STOP_REASON_COMPLETED)
+            self._decode_end()
+            self._end_profiling()
+            
+            return embeddings
+            
+        except Exception as e:
+            self._set_stop_reason(StopReason.ML_STOP_REASON_UNKNOWN)
+            self._decode_end()
+            self._end_profiling()
+            raise RuntimeError(f"Error generating embeddings: {str(e)}")
+
+    def embedding_dim(self) -> int:
+        """Get embedding dimension."""
+        if self.config is None:
+            return 768  # Default dimension
+        
+        # Try different config keys that might contain the dimension
+        if "hidden_size" in self.config:
+            return self.config["hidden_size"]
+        elif "d_model" in self.config:
+            return self.config["d_model"]
+        elif "dim" in self.config:
+            return self.config["dim"]
+        else:
+            return 768  # Fallback default
+
+
+class MLXEmbedder(BaseMLXEmbedder):
+    """
+    Concrete embedder class that routes to the appropriate implementation.
+    This class can be instantiated directly (for C++ compatibility) and will
+    automatically delegate to JinaV2Embedder or MlxEmbeddingEmbedder based on model type.
+    """
+    
+    def __init__(
+        self,
+        model_path: PathType,
+        tokenizer_path: PathType,
+        device: Optional[str] = None,
+    ) -> None:
+        """Initialize the Embedder model."""
+        super().__init__(model_path, tokenizer_path, device)
+        self._impl = None  # Will hold the actual implementation
+    
+    def _get_implementation(self) -> BaseMLXEmbedder:
+        """Get or create the appropriate implementation based on model type."""
+        if self._impl is None:
+            # Detect model type and create appropriate implementation
+            model_type = _detect_model_type(self.model_path)
+            
+            if model_type == "jina_v2":
+                self._impl = JinaV2Embedder(self.model_path, self.tokenizer_path, self.device)
+            else:
+                self._impl = MlxEmbeddingEmbedder(self.model_path, self.tokenizer_path, self.device)
+            
+            # Copy over any existing state
+            if self.model is not None:
+                self._impl.model = self.model
+            if self.tokenizer is not None:
+                self._impl.tokenizer = self.tokenizer
+            if self.config is not None:
+                self._impl.config = self.config
+        
+        return self._impl
+    
+    def load_model(self, model_path: PathType) -> bool:
+        """Load model from path."""
+        # Get the appropriate implementation and delegate
+        impl = self._get_implementation()
+        result = impl.load_model(model_path)
+        
+        # Sync state back
+        self.model = impl.model
+        self.tokenizer = impl.tokenizer
+        self.config = impl.config
+        
+        return result
+    
+    def embed(
+        self,
+        texts: Sequence[str],
+        config: Optional[EmbeddingConfig] = None,
+        clear_cache: bool = True,
+    ) -> List[List[float]]:
+        """Generate embeddings for texts."""
+        # Get the appropriate implementation and delegate
+        impl = self._get_implementation()
+        return impl.embed(texts, config, clear_cache)
+    
+    def embedding_dim(self) -> int:
+        """Get embedding dimension."""
+        # Get the appropriate implementation and delegate
+        impl = self._get_implementation()
+        return impl.embedding_dim()
+    
+    def destroy(self) -> None:
+        """Destroy the model and free resources."""
+        super().destroy()
+        if self._impl is not None:
+            self._impl.destroy()
+            self._impl = None
+
+
+# Backward compatibility alias
+Embedder = MLXEmbedder
+
+
+def _detect_model_type(model_path: PathType) -> str:
+    """Detect the model type from config.json."""
+    if os.path.isfile(model_path):
+        model_path = os.path.dirname(model_path)
+    
+    config_path = os.path.join(model_path, "config.json")
+    
+    if not os.path.exists(config_path):
+        # If no config.json, assume it's a generic model
+        return "generic"
+    
+    try:
+        with open(config_path, "r") as f:
+            config = json.load(f)
+        
+        # Check architectures field for JinaBertModel
+        architectures = config.get("architectures", [])
+        if "JinaBertModel" in architectures:
+            return "jina_v2"
+        
+        # Default to generic mlx_embeddings for other models
+        return "generic"
+    
+    except Exception as e:
+        print(f"Warning: Could not parse config.json: {e}")
+        return "generic"
 
 
 # Factory function for creating embedder instances
@@ -304,9 +608,10 @@ def create_embedder(
     model_path: PathType,
     tokenizer_path: Optional[PathType] = None,
     device: Optional[str] = None,
-) -> Embedder:
-    """Create and return an Embedder instance."""
+) -> MLXEmbedder:
+    """Create and return an MLXEmbedder instance that automatically routes to the appropriate implementation."""
     if tokenizer_path is None:
         tokenizer_path = model_path
     
-    return Embedder(model_path, tokenizer_path, device)
+    # Return the concrete MLXEmbedder which will handle routing internally
+    return MLXEmbedder(model_path, tokenizer_path, device)

nexaai/mlx_backend/embedding/main.py

@@ -12,71 +12,162 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .interface import create_embedder, EmbeddingConfig
+import os
+import sys
+import numpy as np
+from pathlib import Path
 
+# Add parent path for imports
+sys.path.insert(0, str(Path(__file__).parent.parent))
 
-def test_embedding(model_path):
-    """Test embedding model functionality."""
+# Import from interface (uses the factory pattern with routing)
+from .interface import create_embedder
+from .interface import EmbeddingConfig
+from huggingface_hub import snapshot_download
+
+
+def download_model_if_needed(model_id, local_dir):
+    """Download model from Hugging Face Hub if not present locally."""
+    if not os.path.exists(os.path.join(local_dir, "config.json")):
+        print(f"📥 Model not found locally. Downloading {model_id}...")
+        os.makedirs(local_dir, exist_ok=True)
+        try:
+            snapshot_download(
+                repo_id=model_id,
+                local_dir=local_dir,
+                resume_download=True,
+                local_dir_use_symlinks=False
+            )
+            print("✅ Model download completed!")
+        except Exception as e:
+            print(f"❌ Failed to download model: {e}")
+            raise
+
+
+def test_embedding_interface(model_path, is_local=False):
+    """Test embedding model functionality using the interface."""
+    
+    print("=" * 70)
+    print("TESTING EMBEDDING MODEL VIA INTERFACE")
+    print("=" * 70)
+    
+    # Handle model path - download if it's a HF model ID
+    if not is_local and "/" in model_path:
+        # It's a HuggingFace model ID
+        local_dir = f"./modelfiles/{model_path.replace('/', '_')}"
+        download_model_if_needed(model_path, local_dir)
+        model_path = local_dir
+    
+    # Create embedder using factory function (will auto-detect model type)
+    print(f"\n🔍 Creating embedder for: {model_path}")
     embedder = create_embedder(model_path=model_path)
+    print(f"✅ Created embedder type: {type(embedder).__name__}")
     
     # Load the model
-    print("Loading embedding model...")
+    print("\n📚 Loading embedding model...")
     success = embedder.load_model(model_path)
     
     if not success:
-        print("Failed to load model!")
+        print("❌ Failed to load model!")
         return
     
     print("✅ Model loaded successfully!")
-    print(f"Embedding dimension: {embedder.embedding_dim()}")
+    print(f"📏 Embedding dimension: {embedder.embedding_dim()}")
     
     # Test texts
     test_texts = [
         "Hello, how are you?",
         "What is machine learning?",
         "The weather is nice today.",
-        "Python is a programming language."
+        "Python is a programming language.",
+        "Artificial intelligence is changing the world."
     ]
     
-    # Configure embedding
-    config = EmbeddingConfig(
-        batch_size=2,
-        normalize=True,
-        normalize_method="l2"
-    )
+    # Configure embedding with different settings
+    configs = [
+        EmbeddingConfig(batch_size=2, normalize=True, normalize_method="l2"),
+        EmbeddingConfig(batch_size=3, normalize=False),
+    ]
+    
+    for config_idx, config in enumerate(configs):
+        print(f"\n{'='*50}")
+        print(f"TEST {config_idx + 1}: Config - Batch: {config.batch_size}, "
+              f"Normalize: {config.normalize}, Method: {config.normalize_method}")
+        print('='*50)
+        
+        # Generate embeddings
+        embeddings = embedder.embed(test_texts, config)
+        
+        # Display results
+        print(f"\n📊 Generated {len(embeddings)} embeddings")
+        
+        for i, (text, embedding) in enumerate(zip(test_texts[:3], embeddings[:3])):
+            print(f"\n  Text {i+1}: '{text}'")
+            print(f"    Dimension: {len(embedding)}")
+            print(f"    First 5 values: {[f'{v:.4f}' for v in embedding[:5]]}")
+            
+            # Calculate magnitude
+            magnitude = np.linalg.norm(embedding)
+            print(f"    Magnitude: {magnitude:.6f}")
     
-    print(f"\nGenerating embeddings for {len(test_texts)} texts...")
+    # Compute similarity matrix for normalized embeddings
+    print("\n" + "="*50)
+    print("SIMILARITY MATRIX (L2 Normalized)")
+    print("="*50)
     
-    # Generate embeddings
+    config = EmbeddingConfig(batch_size=len(test_texts), normalize=True, normalize_method="l2")
     embeddings = embedder.embed(test_texts, config)
     
-    # Display results
-    print("\nEmbedding Results:")
-    print("=" * 50)
+    # Convert to numpy for easier computation
+    embeddings_np = np.array(embeddings)
+    similarity_matrix = np.dot(embeddings_np, embeddings_np.T)
     
-    for i, (text, embedding) in enumerate(zip(test_texts, embeddings)):
-        print(f"\nText {i+1}: '{text}'")
-        print(f"Embedding shape: {len(embedding)}")
-        print(f"First 5 values: {embedding[:5]}")
-        
-        # Calculate magnitude for normalized embeddings
-        magnitude = sum(x*x for x in embedding) ** 0.5
-        print(f"Magnitude: {magnitude:.6f}")
+    print("\nTexts:")
+    for i, text in enumerate(test_texts):
+        print(f"  [{i}] {text[:30]}...")
+    
+    print("\nSimilarity Matrix:")
+    print("     ", end="")
+    for i in range(len(test_texts)):
+        print(f"  [{i}] ", end="")
+    print()
     
-    # Test similarity between first two embeddings
-    if len(embeddings) >= 2:
-        emb1, emb2 = embeddings[0], embeddings[1]
-        similarity = sum(a*b for a, b in zip(emb1, emb2))
-        print(f"\nCosine similarity between text 1 and 2: {similarity:.6f}")
+    for i in range(len(test_texts)):
+        print(f"  [{i}]", end="")
+        for j in range(len(test_texts)):
+            print(f" {similarity_matrix[i, j]:5.2f}", end="")
+        print()
+    
+    # Find most similar pairs
+    print("\n🔍 Most Similar Pairs (excluding self-similarity):")
+    similarities = []
+    for i in range(len(test_texts)):
+        for j in range(i+1, len(test_texts)):
+            similarities.append((similarity_matrix[i, j], i, j))
+    
+    similarities.sort(reverse=True)
+    for sim, i, j in similarities[:3]:
+        print(f"  • Texts [{i}] and [{j}]: {sim:.4f}")
     
     # Cleanup
     embedder.close()
-    print("\n✅ Embedding test completed!")
+    print("\n✅ Interface test completed successfully!")
 
 
 if __name__ == "__main__":
     import argparse
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model_path", type=str, default="nexaml/jina-v2-fp16-mlx")
+    parser = argparse.ArgumentParser(description="Test embedding models via interface")
+    parser.add_argument(
+        "--model_path", 
+        type=str, 
+        default="nexaml/jina-v2-fp16-mlx",
+        help="Model path (local) or HuggingFace model ID"
+    )
+    parser.add_argument(
+        "--local",
+        action="store_true",
+        help="Indicate if model_path is a local directory"
+    )
     args = parser.parse_args()
-    test_embedding(args.model_path)
+    
+    test_embedding_interface(args.model_path, args.local)