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

nexaai/rerank.py

@@ -1,57 +1,57 @@
-from typing import List, Optional, Sequence, Union
-from abc import abstractmethod
-from dataclasses import dataclass
-
-from nexaai.base import BaseModel
-from nexaai.common import PluginID
-
-
-@dataclass
-class RerankConfig:
-    """Configuration for reranking."""
-    batch_size: int = 1
-    normalize: bool = True
-    normalize_method: str = "softmax"  # "softmax" | "min-max" | "none"
-
-
-class Reranker(BaseModel):
-    """Abstract base class for reranker models."""
-
-    def __init__(self):
-        """Initialize base Reranker class."""
-        pass
-
-    @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,
-                   device_id: Optional[str] = None,
-                   **kwargs
-        ) -> 'Reranker':
-        """Load reranker model from local path, routing to appropriate implementation."""
-        # Check plugin_id value for routing - handle both enum and string
-        plugin_value = plugin_id.value if isinstance(plugin_id, PluginID) else plugin_id
-        
-        if plugin_value == "mlx":
-            from nexaai.rerank_impl.mlx_rerank_impl import MLXRerankImpl
-            return MLXRerankImpl._load_from(model_path, model_name, tokenizer_file, plugin_id, device_id)
-        else:
-            from nexaai.rerank_impl.pybind_rerank_impl import PyBindRerankImpl
-            return PyBindRerankImpl._load_from(model_path, model_name, tokenizer_file, plugin_id, device_id)
-
-    @abstractmethod
-    def load_model(self, model_path: str, extra_data: Optional[str] = None) -> bool:
-        """Load model from path."""
-        pass
-
-    @abstractmethod
-    def rerank(
-        self,
-        query: str,
-        documents: Sequence[str],
-        config: Optional[RerankConfig] = None,
-    ) -> List[float]:
-        """Rerank documents given a query."""
-        pass
+from typing import List, Optional, Sequence, Union
+from abc import abstractmethod
+from dataclasses import dataclass
+
+from nexaai.base import BaseModel
+from nexaai.common import PluginID
+
+
+@dataclass
+class RerankConfig:
+    """Configuration for reranking."""
+    batch_size: int = 1
+    normalize: bool = True
+    normalize_method: str = "softmax"  # "softmax" | "min-max" | "none"
+
+
+class Reranker(BaseModel):
+    """Abstract base class for reranker models."""
+
+    def __init__(self):
+        """Initialize base Reranker class."""
+        pass
+
+    @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,
+                   device_id: Optional[str] = None,
+                   **kwargs
+        ) -> 'Reranker':
+        """Load reranker model from local path, routing to appropriate implementation."""
+        # Check plugin_id value for routing - handle both enum and string
+        plugin_value = plugin_id.value if isinstance(plugin_id, PluginID) else plugin_id
+        
+        if plugin_value == "mlx":
+            from nexaai.rerank_impl.mlx_rerank_impl import MLXRerankImpl
+            return MLXRerankImpl._load_from(model_path, model_name, tokenizer_file, plugin_id, device_id)
+        else:
+            from nexaai.rerank_impl.pybind_rerank_impl import PyBindRerankImpl
+            return PyBindRerankImpl._load_from(model_path, model_name, tokenizer_file, plugin_id, device_id)
+
+    @abstractmethod
+    def load_model(self, model_path: str, extra_data: Optional[str] = None) -> bool:
+        """Load model from path."""
+        pass
+
+    @abstractmethod
+    def rerank(
+        self,
+        query: str,
+        documents: Sequence[str],
+        config: Optional[RerankConfig] = None,
+    ) -> List[float]:
+        """Rerank documents given a query."""
+        pass

nexaai/rerank_impl/mlx_rerank_impl.py

@@ -1,94 +1,94 @@
-# Note: This code is generated by Cursor, not tested yet.
-
-from typing import List, Optional, Sequence, Union
-import os
-
-from nexaai.common import PluginID
-from nexaai.rerank import Reranker, RerankConfig
-from nexaai.mlx_backend.rerank.interface import Reranker as MLXRerankInterface, create_reranker
-
-
-class MLXRerankImpl(Reranker):
-    def __init__(self):
-        """Initialize MLX Rerank implementation."""
-        super().__init__()
-        self._mlx_reranker = 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,
-                   device_id: Optional[str] = None
-        ) -> 'MLXRerankImpl':
-        """Load reranker model from local path using MLX backend."""
-        try:
-            # MLX Rerank interfaces are already imported
-            
-            # Create instance and load MLX reranker
-            instance = cls()
-            instance._mlx_reranker = create_reranker(
-                model_path=model_path,
-                # model_name=model_name, # FIXME: For MLX Reranker, model_name is not used
-                tokenizer_path=tokenizer_file,
-                device=device_id
-            )
-            
-            # Load the model
-            success = instance._mlx_reranker.load_model(model_path)
-            if not success:
-                raise RuntimeError("Failed to load MLX reranker model")
-            
-            return instance
-        except Exception as e:
-            raise RuntimeError(f"Failed to load MLX Reranker: {str(e)}")
-
-    def eject(self):
-        """Destroy the model and free resources."""
-        if self._mlx_reranker:
-            self._mlx_reranker.destroy()
-            self._mlx_reranker = None
-
-    def load_model(self, model_path: str, extra_data: Optional[str] = None) -> bool:
-        """Load model from path."""
-        if not self._mlx_reranker:
-            raise RuntimeError("MLX Reranker not initialized")
-        
-        try:
-            return self._mlx_reranker.load_model(model_path, extra_data)
-        except Exception as e:
-            raise RuntimeError(f"Failed to load reranker model: {str(e)}")
-
-    def rerank(
-        self,
-        query: str,
-        documents: Sequence[str],
-        config: Optional[RerankConfig] = None,
-    ) -> List[float]:
-        """Rerank documents given a query."""
-        if not self._mlx_reranker:
-            raise RuntimeError("MLX Reranker not loaded")
-        
-        try:
-            # Convert our config to MLX format if provided
-            mlx_config = None
-            if config:
-                from nexaai.mlx_backend.rerank.interface import RerankConfig as MLXRerankConfig
-                
-                mlx_config = MLXRerankConfig(
-                    batch_size=config.batch_size,
-                    normalize=config.normalize,
-                    normalize_method=config.normalize_method
-                )
-            
-            # Use MLX reranking
-            scores = self._mlx_reranker.rerank(query, documents, mlx_config)
-            
-            # Convert mx.array to Python list of floats
-            return scores.tolist()
-            
-        except Exception as e:
-            raise RuntimeError(f"Failed to rerank documents: {str(e)}")
-
-
+# Note: This code is generated by Cursor, not tested yet.
+
+from typing import List, Optional, Sequence, Union
+import os
+
+from nexaai.common import PluginID
+from nexaai.rerank import Reranker, RerankConfig
+from nexaai.mlx_backend.rerank.interface import Reranker as MLXRerankInterface, create_reranker
+
+
+class MLXRerankImpl(Reranker):
+    def __init__(self):
+        """Initialize MLX Rerank implementation."""
+        super().__init__()
+        self._mlx_reranker = 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,
+                   device_id: Optional[str] = None
+        ) -> 'MLXRerankImpl':
+        """Load reranker model from local path using MLX backend."""
+        try:
+            # MLX Rerank interfaces are already imported
+            
+            # Create instance and load MLX reranker
+            instance = cls()
+            instance._mlx_reranker = create_reranker(
+                model_path=model_path,
+                # model_name=model_name, # FIXME: For MLX Reranker, model_name is not used
+                tokenizer_path=tokenizer_file,
+                device=device_id
+            )
+            
+            # Load the model
+            success = instance._mlx_reranker.load_model(model_path)
+            if not success:
+                raise RuntimeError("Failed to load MLX reranker model")
+            
+            return instance
+        except Exception as e:
+            raise RuntimeError(f"Failed to load MLX Reranker: {str(e)}")
+
+    def eject(self):
+        """Destroy the model and free resources."""
+        if self._mlx_reranker:
+            self._mlx_reranker.destroy()
+            self._mlx_reranker = None
+
+    def load_model(self, model_path: str, extra_data: Optional[str] = None) -> bool:
+        """Load model from path."""
+        if not self._mlx_reranker:
+            raise RuntimeError("MLX Reranker not initialized")
+        
+        try:
+            return self._mlx_reranker.load_model(model_path, extra_data)
+        except Exception as e:
+            raise RuntimeError(f"Failed to load reranker model: {str(e)}")
+
+    def rerank(
+        self,
+        query: str,
+        documents: Sequence[str],
+        config: Optional[RerankConfig] = None,
+    ) -> List[float]:
+        """Rerank documents given a query."""
+        if not self._mlx_reranker:
+            raise RuntimeError("MLX Reranker not loaded")
+        
+        try:
+            # Convert our config to MLX format if provided
+            mlx_config = None
+            if config:
+                from nexaai.mlx_backend.rerank.interface import RerankConfig as MLXRerankConfig
+                
+                mlx_config = MLXRerankConfig(
+                    batch_size=config.batch_size,
+                    normalize=config.normalize,
+                    normalize_method=config.normalize_method
+                )
+            
+            # Use MLX reranking
+            scores = self._mlx_reranker.rerank(query, documents, mlx_config)
+            
+            # Convert mx.array to Python list of floats
+            return scores.tolist()
+            
+        except Exception as e:
+            raise RuntimeError(f"Failed to rerank documents: {str(e)}")
+
+

nexaai/rerank_impl/pybind_rerank_impl.py

@@ -1,136 +1,136 @@
-from typing import List, Optional, Sequence, Union
-import numpy as np
-
-from nexaai.common import PluginID
-from nexaai.rerank import Reranker, RerankConfig
-from nexaai.binds import rerank_bind, common_bind
-from nexaai.runtime import _ensure_runtime
-
-
-class PyBindRerankImpl(Reranker):
-    def __init__(self, _handle_ptr):
-        """
-        Internal initializer
-        
-        Args:
-            _handle_ptr: Capsule handle to the C++ reranker object
-        """
-        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,
-                   device_id: Optional[str] = None
-        ) -> 'PyBindRerankImpl':
-        """
-        Load reranker model from local path using PyBind backend.
-        
-        Args:
-            model_path: Path to the model file
-            model_name: Name of the model (optional)
-            tokenizer_file: Path to the tokenizer file (default: "tokenizer.json")
-            plugin_id: Plugin ID to use for the model (default: PluginID.LLAMA_CPP)
-            device_id: Device ID to use for the model (optional)
-            
-        Returns:
-            PyBindRerankImpl instance
-        """
-        _ensure_runtime()
-        
-        # Convert enum to string for C++ binding
-        plugin_id_str = plugin_id.value if isinstance(plugin_id, PluginID) else plugin_id
-        
-        # Create model config
-        model_config = common_bind.ModelConfig()
-        
-        # Create reranker handle with new API signature
-        handle = rerank_bind.ml_reranker_create(
-            model_path,
-            model_name,
-            tokenizer_file,
-            model_config,
-            plugin_id_str,
-            device_id
-        )
-        
-        return cls(handle)
-
-    def eject(self):
-        """
-        Clean up resources and destroy the reranker
-        """
-        # Destructor of the handle will unload the model correctly
-        if hasattr(self, '_handle') and self._handle is not None:
-            del self._handle
-            self._handle = None
-
-    def load_model(self, model_path: str, extra_data: Optional[str] = None) -> bool:
-        """
-        Load model from path.
-        
-        Note: This method is not typically used directly. Use _load_from instead.
-        
-        Args:
-            model_path: Path to the model file
-            extra_data: Additional data (unused)
-            
-        Returns:
-            True if successful
-        """
-        # This method is part of the BaseModel interface but typically not used
-        # directly for PyBind implementations since _load_from handles creation
-        raise NotImplementedError("Use _load_from class method to load models")
-
-    def rerank(
-        self,
-        query: str,
-        documents: Sequence[str],
-        config: Optional[RerankConfig] = None,
-    ) -> List[float]:
-        """
-        Rerank documents given a query.
-        
-        Args:
-            query: Query text as UTF-8 string
-            documents: List of document texts to rerank
-            config: Optional reranking configuration
-            
-        Returns:
-            List of ranking scores (one per document)
-        """
-        if self._handle is None:
-            raise RuntimeError("Reranker handle is None. Model may have been ejected.")
-        
-        # Use default config if not provided
-        if config is None:
-            config = RerankConfig()
-        
-        # Create bind config
-        bind_config = rerank_bind.RerankConfig()
-        bind_config.batch_size = config.batch_size
-        bind_config.normalize = config.normalize
-        bind_config.normalize_method = config.normalize_method
-        
-        # Convert documents to list if needed
-        documents_list = list(documents)
-        
-        # Call the binding which returns a dict with scores and profile_data
-        result = rerank_bind.ml_reranker_rerank(
-            self._handle,
-            query,
-            documents_list,
-            bind_config
-        )
-        
-        # Extract scores from result dict
-        scores_array = result.get("scores", np.array([]))
-        
-        # Convert numpy array to list of floats
-        if isinstance(scores_array, np.ndarray):
-            return scores_array.tolist()
-        else:
-            return []
+from typing import List, Optional, Sequence, Union
+import numpy as np
+
+from nexaai.common import PluginID
+from nexaai.rerank import Reranker, RerankConfig
+from nexaai.binds import rerank_bind, common_bind
+from nexaai.runtime import _ensure_runtime
+
+
+class PyBindRerankImpl(Reranker):
+    def __init__(self, _handle_ptr):
+        """
+        Internal initializer
+        
+        Args:
+            _handle_ptr: Capsule handle to the C++ reranker object
+        """
+        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,
+                   device_id: Optional[str] = None
+        ) -> 'PyBindRerankImpl':
+        """
+        Load reranker model from local path using PyBind backend.
+        
+        Args:
+            model_path: Path to the model file
+            model_name: Name of the model (optional)
+            tokenizer_file: Path to the tokenizer file (default: "tokenizer.json")
+            plugin_id: Plugin ID to use for the model (default: PluginID.LLAMA_CPP)
+            device_id: Device ID to use for the model (optional)
+            
+        Returns:
+            PyBindRerankImpl instance
+        """
+        _ensure_runtime()
+        
+        # Convert enum to string for C++ binding
+        plugin_id_str = plugin_id.value if isinstance(plugin_id, PluginID) else plugin_id
+        
+        # Create model config
+        model_config = common_bind.ModelConfig()
+        
+        # Create reranker handle with new API signature
+        handle = rerank_bind.ml_reranker_create(
+            model_path,
+            model_name,
+            tokenizer_file,
+            model_config,
+            plugin_id_str,
+            device_id
+        )
+        
+        return cls(handle)
+
+    def eject(self):
+        """
+        Clean up resources and destroy the reranker
+        """
+        # Destructor of the handle will unload the model correctly
+        if hasattr(self, '_handle') and self._handle is not None:
+            del self._handle
+            self._handle = None
+
+    def load_model(self, model_path: str, extra_data: Optional[str] = None) -> bool:
+        """
+        Load model from path.
+        
+        Note: This method is not typically used directly. Use _load_from instead.
+        
+        Args:
+            model_path: Path to the model file
+            extra_data: Additional data (unused)
+            
+        Returns:
+            True if successful
+        """
+        # This method is part of the BaseModel interface but typically not used
+        # directly for PyBind implementations since _load_from handles creation
+        raise NotImplementedError("Use _load_from class method to load models")
+
+    def rerank(
+        self,
+        query: str,
+        documents: Sequence[str],
+        config: Optional[RerankConfig] = None,
+    ) -> List[float]:
+        """
+        Rerank documents given a query.
+        
+        Args:
+            query: Query text as UTF-8 string
+            documents: List of document texts to rerank
+            config: Optional reranking configuration
+            
+        Returns:
+            List of ranking scores (one per document)
+        """
+        if self._handle is None:
+            raise RuntimeError("Reranker handle is None. Model may have been ejected.")
+        
+        # Use default config if not provided
+        if config is None:
+            config = RerankConfig()
+        
+        # Create bind config
+        bind_config = rerank_bind.RerankConfig()
+        bind_config.batch_size = config.batch_size
+        bind_config.normalize = config.normalize
+        bind_config.normalize_method = config.normalize_method
+        
+        # Convert documents to list if needed
+        documents_list = list(documents)
+        
+        # Call the binding which returns a dict with scores and profile_data
+        result = rerank_bind.ml_reranker_rerank(
+            self._handle,
+            query,
+            documents_list,
+            bind_config
+        )
+        
+        # Extract scores from result dict
+        scores_array = result.get("scores", np.array([]))
+        
+        # Convert numpy array to list of floats
+        if isinstance(scores_array, np.ndarray):
+            return scores_array.tolist()
+        else:
+            return []