nexaai 1.0.10__cp310-cp310-macosx_14_0_universal2.whl → 1.0.11rc2__cp310-cp310-macosx_14_0_universal2.whl

nexaai/mlx_backend/embedding/interface.py

@@ -20,16 +20,11 @@ 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,
@@ -39,24 +34,13 @@ from ml import (
 # Import profiling module
 from profiling import ProfilingMixin, StopReason
 
-# 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
+# Import the model implementation
+from .modeling.nexa_jina_v2 import Model, ModelArgs
 
 
-class BaseMLXEmbedder(BaseEmbedder, ProfilingMixin, ABC):
+class Embedder(BaseEmbedder, ProfilingMixin):
     """
-    Abstract base embedder interface for MLX embedding models.
+    Embedder interface for MLX embedding models.
     API aligned with ml.py Embedder abstract base class.
     """
 
@@ -80,7 +64,7 @@ class BaseMLXEmbedder(BaseEmbedder, ProfilingMixin, ABC):
         
         self.model_path = model_path
         self.tokenizer_path = tokenizer_path
-        self.device = device if device is not None else "cpu"
+        self.device = device if device is not None else "cpu" # TODO: This device field is never used
         
         # Initialize model and tokenizer as None
         self.model = None
@@ -94,69 +78,6 @@ class BaseMLXEmbedder(BaseEmbedder, ProfilingMixin, ABC):
         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:
@@ -176,6 +97,10 @@ class JinaV2Embedder(BaseMLXEmbedder):
             print(f"Failed to load model: {e}")
             return False
 
+    def close(self) -> None:
+        """Close the model."""
+        self.destroy()
+
     def embed(
         self,
         texts: Sequence[str],
@@ -233,6 +158,22 @@ class JinaV2Embedder(BaseMLXEmbedder):
             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."""
         
@@ -340,267 +281,22 @@ class JinaV2Embedder(BaseMLXEmbedder):
         
         return embedding_list
 
-
-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.")
-        
-        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 config is None:
-            config = EmbeddingConfig()
-        
-        # 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
+    def _normalize_embedding(self, embedding: List[float], method: str) -> List[float]:
+        """Normalize embedding using specified method."""
+        if method == "none":
+            return embedding
         
-        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)
+        embedding_array = np.array(embedding)
         
-        # Check architectures field for JinaBertModel
-        architectures = config.get("architectures", [])
-        if "JinaBertModel" in architectures:
-            return "jina_v2"
+        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
         
-        # 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"
+        return embedding_array.tolist()
 
 
 # Factory function for creating embedder instances
@@ -608,10 +304,9 @@ def create_embedder(
     model_path: PathType,
     tokenizer_path: Optional[PathType] = None,
     device: Optional[str] = None,
-) -> MLXEmbedder:
-    """Create and return an MLXEmbedder instance that automatically routes to the appropriate implementation."""
+) -> Embedder:
+    """Create and return an Embedder instance."""
     if tokenizer_path is None:
         tokenizer_path = model_path
     
-    # Return the concrete MLXEmbedder which will handle routing internally
-    return MLXEmbedder(model_path, tokenizer_path, device)
+    return Embedder(model_path, tokenizer_path, device)

nexaai/mlx_backend/embedding/main.py

@@ -12,162 +12,71 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
-import sys
-import numpy as np
-from pathlib import Path
+from .interface import create_embedder, EmbeddingConfig
 
-# Add parent path for imports
-sys.path.insert(0, str(Path(__file__).parent.parent))
 
-# 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}")
+def test_embedding(model_path):
+    """Test embedding model functionality."""
     embedder = create_embedder(model_path=model_path)
-    print(f"✅ Created embedder type: {type(embedder).__name__}")
     
     # Load the model
-    print("\n📚 Loading embedding model...")
+    print("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.",
-        "Artificial intelligence is changing the world."
+        "Python is a programming language."
     ]
     
-    # 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}")
+    # Configure embedding
+    config = EmbeddingConfig(
+        batch_size=2,
+        normalize=True,
+        normalize_method="l2"
+    )
     
-    # Compute similarity matrix for normalized embeddings
-    print("\n" + "="*50)
-    print("SIMILARITY MATRIX (L2 Normalized)")
-    print("="*50)
+    print(f"\nGenerating embeddings for {len(test_texts)} texts...")
     
-    config = EmbeddingConfig(batch_size=len(test_texts), normalize=True, normalize_method="l2")
+    # Generate embeddings
     embeddings = embedder.embed(test_texts, config)
     
-    # Convert to numpy for easier computation
-    embeddings_np = np.array(embeddings)
-    similarity_matrix = np.dot(embeddings_np, embeddings_np.T)
-    
-    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()
+    # Display results
+    print("\nEmbedding Results:")
+    print("=" * 50)
     
-    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))
+    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}")
     
-    similarities.sort(reverse=True)
-    for sim, i, j in similarities[:3]:
-        print(f"  • Texts [{i}] and [{j}]: {sim:.4f}")
+    # 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}")
     
     # Cleanup
     embedder.close()
-    print("\n✅ Interface test completed successfully!")
+    print("\n✅ Embedding test completed!")
 
 
 if __name__ == "__main__":
     import argparse
-    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"
-    )
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", type=str, default="nexaml/jina-v2-fp16-mlx")
     args = parser.parse_args()
-    
-    test_embedding_interface(args.model_path, args.local)
+    test_embedding(args.model_path)