nexaai 1.0.21rc5__cp313-cp313-win_arm64.whl → 1.0.21rc16__cp313-cp313-win_arm64.whl

nexaai/embedder_impl/mlx_embedder_impl.py

@@ -1,118 +1,118 @@
-from typing import List, Union
-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.ml import ModelConfig as MLXModelConfig, SamplerConfig as MLXSamplerConfig, GenerationConfig as MLXGenerationConfig, EmbeddingConfig
-
-
-class MLXEmbedderImpl(Embedder):
-    def __init__(self):
-        """Initialize MLX Embedder implementation."""
-        super().__init__()
-        self._mlx_embedder = None
-
-    @classmethod
-    def _load_from(cls, model_path: str, model_name: str = None, tokenizer_file: str = "tokenizer.json", plugin_id: Union[PluginID, str] = PluginID.MLX):
-        """
-        Load an embedder from model files using MLX backend.
-
-        Args:
-            model_path: Path to the model file
-            model_name: Name of the model
-            tokenizer_file: Path to the tokenizer file (default: "tokenizer.json")
-            plugin_id: Plugin ID to use for the model (default: PluginID.MLX)
-
-        Returns:
-            MLXEmbedderImpl instance
-        """
-        try:
-            # Create instance
-            instance = cls()
-            
-            # 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,
-                # model_name=model_name, # FIXME: For MLX Embedder, model_name is not used
-                tokenizer_path=tokenizer_file
-            )
-            
-            # Load the model
-            success = instance._mlx_embedder.load_model(model_path)
-            if not success:
-                raise RuntimeError("Failed to load MLX embedder model")
-            
-            return instance
-        except Exception as e:
-            raise RuntimeError(f"Failed to load MLX Embedder: {str(e)}")
-
-    def eject(self):
-        """
-        Clean up resources and destroy the embedder
-        """
-        if self._mlx_embedder:
-            self._mlx_embedder.destroy()
-            self._mlx_embedder = None
-
-    def generate(self, texts: Union[List[str], str] = None, config: EmbeddingConfig = EmbeddingConfig(), input_ids: Union[List[int], List[List[int]]] = None) -> np.ndarray:
-        """
-        Generate embeddings for the given texts or input_ids.
-
-        Args:
-            texts: List of strings or single string to embed
-            input_ids: Pre-tokenized input as:
-                      - Single sequence: list of integers [1, 2, 3, 4]
-                      - Multiple sequences: list of lists [[1, 2, 3], [4, 5, 6]]
-            config: Configuration for embedding generation
-
-        Returns:
-            numpy array of embeddings with shape (num_sequences, embedding_dim)
-        """
-        if not self._mlx_embedder:
-            raise RuntimeError("MLX Embedder not loaded")
-        
-        if texts is None and input_ids is None:
-            raise ValueError("Either texts or input_ids must be provided")
-        
-        # MLX embedder currently only supports text input, not pre-tokenized input_ids
-        if input_ids is not None:
-            raise NotImplementedError("MLX embedder does not support input_ids, only text input")
-        
-        try:
-            # Convert single string to list if needed
-            if isinstance(texts, str):
-                texts = [texts]
-            
-            # MLX config classes are already imported
-            
-            # Convert our config to MLX config
-            mlx_config = EmbeddingConfig()
-            mlx_config.batch_size = config.batch_size
-            mlx_config.normalize = config.normalize
-            mlx_config.normalize_method = config.normalize_method
-            
-            # Generate embeddings using MLX
-            embeddings = self._mlx_embedder.embed(texts, mlx_config)
-            
-            # Convert to numpy array
-            return np.array(embeddings, dtype=np.float32)
-            
-        except Exception as e:
-            raise RuntimeError(f"Failed to generate embeddings: {str(e)}")
-
-    def get_embedding_dim(self) -> int:
-        """
-        Get the embedding dimension of the model
-
-        Returns:
-            The embedding dimension in int
-        """
-        if not self._mlx_embedder:
-            raise RuntimeError("MLX Embedder not loaded")
-        
-        try:
-            return self._mlx_embedder.embedding_dim()
-        except Exception as e:
-            raise RuntimeError(f"Failed to get embedding dimension: {str(e)}")
+from typing import List, Union
+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.ml import ModelConfig as MLXModelConfig, SamplerConfig as MLXSamplerConfig, GenerationConfig as MLXGenerationConfig, EmbeddingConfig
+
+
+class MLXEmbedderImpl(Embedder):
+    def __init__(self):
+        """Initialize MLX Embedder implementation."""
+        super().__init__()
+        self._mlx_embedder = None
+
+    @classmethod
+    def _load_from(cls, model_path: str, model_name: str = None, tokenizer_file: str = "tokenizer.json", plugin_id: Union[PluginID, str] = PluginID.MLX):
+        """
+        Load an embedder from model files using MLX backend.
+
+        Args:
+            model_path: Path to the model file
+            model_name: Name of the model
+            tokenizer_file: Path to the tokenizer file (default: "tokenizer.json")
+            plugin_id: Plugin ID to use for the model (default: PluginID.MLX)
+
+        Returns:
+            MLXEmbedderImpl instance
+        """
+        try:
+            # Create instance
+            instance = cls()
+            
+            # 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,
+                # model_name=model_name, # FIXME: For MLX Embedder, model_name is not used
+                tokenizer_path=tokenizer_file
+            )
+            
+            # Load the model
+            success = instance._mlx_embedder.load_model(model_path)
+            if not success:
+                raise RuntimeError("Failed to load MLX embedder model")
+            
+            return instance
+        except Exception as e:
+            raise RuntimeError(f"Failed to load MLX Embedder: {str(e)}")
+
+    def eject(self):
+        """
+        Clean up resources and destroy the embedder
+        """
+        if self._mlx_embedder:
+            self._mlx_embedder.destroy()
+            self._mlx_embedder = None
+
+    def generate(self, texts: Union[List[str], str] = None, config: EmbeddingConfig = EmbeddingConfig(), input_ids: Union[List[int], List[List[int]]] = None) -> np.ndarray:
+        """
+        Generate embeddings for the given texts or input_ids.
+
+        Args:
+            texts: List of strings or single string to embed
+            input_ids: Pre-tokenized input as:
+                      - Single sequence: list of integers [1, 2, 3, 4]
+                      - Multiple sequences: list of lists [[1, 2, 3], [4, 5, 6]]
+            config: Configuration for embedding generation
+
+        Returns:
+            numpy array of embeddings with shape (num_sequences, embedding_dim)
+        """
+        if not self._mlx_embedder:
+            raise RuntimeError("MLX Embedder not loaded")
+        
+        if texts is None and input_ids is None:
+            raise ValueError("Either texts or input_ids must be provided")
+        
+        # MLX embedder currently only supports text input, not pre-tokenized input_ids
+        if input_ids is not None:
+            raise NotImplementedError("MLX embedder does not support input_ids, only text input")
+        
+        try:
+            # Convert single string to list if needed
+            if isinstance(texts, str):
+                texts = [texts]
+            
+            # MLX config classes are already imported
+            
+            # Convert our config to MLX config
+            mlx_config = EmbeddingConfig()
+            mlx_config.batch_size = config.batch_size
+            mlx_config.normalize = config.normalize
+            mlx_config.normalize_method = config.normalize_method
+            
+            # Generate embeddings using MLX
+            embeddings = self._mlx_embedder.embed(texts, mlx_config)
+            
+            # Convert to numpy array
+            return np.array(embeddings, dtype=np.float32)
+            
+        except Exception as e:
+            raise RuntimeError(f"Failed to generate embeddings: {str(e)}")
+
+    def get_embedding_dim(self) -> int:
+        """
+        Get the embedding dimension of the model
+
+        Returns:
+            The embedding dimension in int
+        """
+        if not self._mlx_embedder:
+            raise RuntimeError("MLX Embedder not loaded")
+        
+        try:
+            return self._mlx_embedder.embedding_dim()
+        except Exception as e:
+            raise RuntimeError(f"Failed to get embedding dimension: {str(e)}")

nexaai/embedder_impl/pybind_embedder_impl.py

@@ -1,96 +1,96 @@
-from typing import List, Union
-import numpy as np
-
-from nexaai.common import PluginID
-from nexaai.embedder import Embedder, EmbeddingConfig
-from nexaai.binds import embedder_bind
-from nexaai.runtime import _ensure_runtime
-
-
-class PyBindEmbedderImpl(Embedder):
-    def __init__(self, _handle_ptr):
-        """
-        Internal initializer
-        """
-        super().__init__()
-        self._handle = _handle_ptr
-
-    @classmethod
-    def _load_from(cls, model_path: str, model_name: str = None, tokenizer_file: str = "tokenizer.json", plugin_id: Union[PluginID, str] = PluginID.LLAMA_CPP):
-        """
-        Load an embedder from model files
-
-        Args:
-            model_path: Path to the model file
-            model_name: Name of the model
-            tokenizer_file: Path to the tokenizer file (default: "tokenizer.json")
-            plugin_id: Plugin ID to use for the model (default: PluginID.LLAMA_CPP)
-
-        Returns:
-            PyBindEmbedderImpl instance
-        """
-        _ensure_runtime()
-        # Convert enum to string for C++ binding
-        plugin_id_str = plugin_id.value if isinstance(plugin_id, PluginID) else plugin_id
-        # New parameter order: model_path, plugin_id, tokenizer_path (optional)
-        handle = embedder_bind.ml_embedder_create(model_path, model_name, plugin_id_str, tokenizer_file)
-        return cls(handle)
-
-    def eject(self):
-        """
-        Clean up resources and destroy the embedder
-        """
-        # Destructor of the handle will unload the model correctly
-        del self._handle
-        self._handle = None
-
-    def generate(self, texts: Union[List[str], str] = None, config: EmbeddingConfig = EmbeddingConfig(), input_ids: Union[List[int], List[List[int]]] = None) -> np.ndarray:
-        """
-        Generate embeddings for the given texts or input_ids.
-
-        Args:
-            texts: List of strings or single string to embed
-            input_ids: Pre-tokenized input as:
-                      - Single sequence: list of integers [1, 2, 3, 4]
-                      - Multiple sequences: list of lists [[1, 2, 3], [4, 5, 6]]
-            config: Configuration for embedding generation
-
-        Returns:
-            numpy array of embeddings with shape (num_sequences, embedding_dim)
-        """
-        if texts is None and input_ids is None:
-            raise ValueError("Either texts or input_ids must be provided")
-
-        # Create bind config
-        bind_config = embedder_bind.EmbeddingConfig()
-        bind_config.batch_size = config.batch_size
-        bind_config.normalize = config.normalize
-        bind_config.normalize_method = config.normalize_method
-
-        # Convert single string to list if needed
-        if isinstance(texts, str):
-            texts = [texts]
-
-        # Convert input_ids to 2D format if needed
-        processed_input_ids = None
-        if input_ids is not None:
-            if len(input_ids) > 0 and isinstance(input_ids[0], int):
-                # Single sequence: convert [1, 2, 3] to [[1, 2, 3]]
-                processed_input_ids = [input_ids]
-            else:
-                # Multiple sequences: already in correct format [[1, 2], [3, 4]]
-                processed_input_ids = input_ids
-
-        # Pass both parameters, let the ABI handle validation
-        embeddings = embedder_bind.ml_embedder_embed(self._handle, bind_config, texts, processed_input_ids)
-
-        return embeddings
-
-    def get_embedding_dim(self) -> int:
-        """
-        Get the embedding dimension of the model
-
-        Returns:
-            The embedding dimension in int
-        """
-        return embedder_bind.ml_embedder_embedding_dim(self._handle)
+from typing import List, Union
+import numpy as np
+
+from nexaai.common import PluginID
+from nexaai.embedder import Embedder, EmbeddingConfig
+from nexaai.binds import embedder_bind
+from nexaai.runtime import _ensure_runtime
+
+
+class PyBindEmbedderImpl(Embedder):
+    def __init__(self, _handle_ptr):
+        """
+        Internal initializer
+        """
+        super().__init__()
+        self._handle = _handle_ptr
+
+    @classmethod
+    def _load_from(cls, model_path: str, model_name: str = None, tokenizer_file: str = "tokenizer.json", plugin_id: Union[PluginID, str] = PluginID.LLAMA_CPP):
+        """
+        Load an embedder from model files
+
+        Args:
+            model_path: Path to the model file
+            model_name: Name of the model
+            tokenizer_file: Path to the tokenizer file (default: "tokenizer.json")
+            plugin_id: Plugin ID to use for the model (default: PluginID.LLAMA_CPP)
+
+        Returns:
+            PyBindEmbedderImpl instance
+        """
+        _ensure_runtime()
+        # Convert enum to string for C++ binding
+        plugin_id_str = plugin_id.value if isinstance(plugin_id, PluginID) else plugin_id
+        # New parameter order: model_path, plugin_id, tokenizer_path (optional)
+        handle = embedder_bind.ml_embedder_create(model_path, model_name, plugin_id_str, tokenizer_file)
+        return cls(handle)
+
+    def eject(self):
+        """
+        Clean up resources and destroy the embedder
+        """
+        # Destructor of the handle will unload the model correctly
+        del self._handle
+        self._handle = None
+
+    def generate(self, texts: Union[List[str], str] = None, config: EmbeddingConfig = EmbeddingConfig(), input_ids: Union[List[int], List[List[int]]] = None) -> np.ndarray:
+        """
+        Generate embeddings for the given texts or input_ids.
+
+        Args:
+            texts: List of strings or single string to embed
+            input_ids: Pre-tokenized input as:
+                      - Single sequence: list of integers [1, 2, 3, 4]
+                      - Multiple sequences: list of lists [[1, 2, 3], [4, 5, 6]]
+            config: Configuration for embedding generation
+
+        Returns:
+            numpy array of embeddings with shape (num_sequences, embedding_dim)
+        """
+        if texts is None and input_ids is None:
+            raise ValueError("Either texts or input_ids must be provided")
+
+        # Create bind config
+        bind_config = embedder_bind.EmbeddingConfig()
+        bind_config.batch_size = config.batch_size
+        bind_config.normalize = config.normalize
+        bind_config.normalize_method = config.normalize_method
+
+        # Convert single string to list if needed
+        if isinstance(texts, str):
+            texts = [texts]
+
+        # Convert input_ids to 2D format if needed
+        processed_input_ids = None
+        if input_ids is not None:
+            if len(input_ids) > 0 and isinstance(input_ids[0], int):
+                # Single sequence: convert [1, 2, 3] to [[1, 2, 3]]
+                processed_input_ids = [input_ids]
+            else:
+                # Multiple sequences: already in correct format [[1, 2], [3, 4]]
+                processed_input_ids = input_ids
+
+        # Pass both parameters, let the ABI handle validation
+        embeddings = embedder_bind.ml_embedder_embed(self._handle, bind_config, texts, processed_input_ids)
+
+        return embeddings
+
+    def get_embedding_dim(self) -> int:
+        """
+        Get the embedding dimension of the model
+
+        Returns:
+            The embedding dimension in int
+        """
+        return embedder_bind.ml_embedder_embedding_dim(self._handle)