nexaai 1.0.29__cp310-cp310-macosx_14_0_universal2.whl

nexaai/rerank.py

@@ -0,0 +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

nexaai/rerank_impl/mlx_rerank_impl.py

@@ -0,0 +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)}")
+
+

nexaai/rerank_impl/pybind_rerank_impl.py

@@ -0,0 +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 []

nexaai/runtime.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+import atexit
+import threading
+from typing import Optional, Any
+
+from nexaai.binds import common_bind
+
+_init_lock        = threading.Lock()
+_runtime_alive    = False          # global flag
+
+def _ensure_runtime() -> None:
+    """Initialise the runtime exactly once (thread‑safe, lazy)."""
+    global _runtime_alive
+    if not _runtime_alive:
+        with _init_lock:
+            if not _runtime_alive:          # double‑checked locking
+                common_bind.ml_init()
+                _runtime_alive = True
+                atexit.register(_shutdown_runtime)
+
+def _shutdown_runtime() -> None:
+    """Tear the runtime down; idempotent and registered with atexit."""
+    global _runtime_alive
+    if _runtime_alive:
+        common_bind.ml_deinit()
+        _runtime_alive = False
+
+# Public helper so advanced users can reclaim memory on demand
+shutdown = _shutdown_runtime
+
+def is_initialized() -> bool:
+    """Check if the runtime has been initialized."""
+    return _runtime_alive
+
+# ----------------------------------------------------------------------
+# Single public class
+# ----------------------------------------------------------------------
+class Session:
+    """
+    Model session **and** runtime guard in one object.
+
+        sess = myrt.Session("foo.mdl")
+        out  = sess.run(inputs)
+        sess.close()           # optional (model only)
+
+    The global runtime is initialised lazily when the first Session
+    is created and stays alive until:
+    • the interpreter exits, or
+    • `myrt.shutdown()` is called.
+    """
+
+    # ---- construction -------------------------------------------------
+    def __init__(self, model_path: str) -> None:
+        _ensure_runtime()
+
+    # safety net – make GC close the model
+    def __del__(self) -> None:
+        try:
+            self.close()
+        except Exception:
+            pass
+
+    # allow `with Session(...) as s:` syntax
+    def __enter__(self) -> "Session":
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        self.close()

nexaai/runtime_error.py

@@ -0,0 +1,24 @@
+"""Runtime errors for Nexa SDK operations."""
+
+
+class NexaRuntimeError(Exception):
+    """Base class for Nexa runtime errors."""
+    
+    def __init__(self, message: str, error_code: int = None):
+        self.error_code = error_code
+        super().__init__(message)
+
+
+class ContextLengthExceededError(NexaRuntimeError):
+    """Raised when the input context length exceeds the model's maximum."""
+    
+    def __init__(self, message: str = "Input context length exceeded model's maximum", error_code: int = None):
+        super().__init__(message, error_code)
+
+
+class GenerationError(NexaRuntimeError):
+    """Raised when generation fails."""
+    
+    def __init__(self, message: str = "Generation failed", error_code: int = None):
+        super().__init__(message, error_code)
+

nexaai/tts.py

@@ -0,0 +1,75 @@
+from typing import List, Optional, Union
+from abc import abstractmethod
+from dataclasses import dataclass
+
+from nexaai.base import BaseModel
+from nexaai.common import PluginID
+
+
+@dataclass
+class TTSConfig:
+    """Configuration for TTS."""
+    voice: str = "default"
+    speed: float = 1.0
+    seed: int = -1  # –1 for random
+    sample_rate: int = 22050
+
+
+@dataclass
+class TTSSamplerConfig:
+    """Configuration for TTS sampling."""
+    temperature: float = 1.0
+    noise_scale: float = 0.667
+    length_scale: float = 1.0
+
+
+@dataclass
+class TTSResult:
+    """Result from TTS processing."""
+    audio_path: str  # Path where the synthesized audio is saved
+    duration_seconds: float
+    sample_rate: int
+    channels: int
+    num_samples: int
+
+
+class TTS(BaseModel):
+    """Abstract base class for Text-to-Speech models."""
+
+    def __init__(self):
+        """Initialize base TTS class."""
+        pass
+
+    @classmethod
+    def _load_from(cls,
+                   model_path: str,
+                   vocoder_path: str,
+                   plugin_id: Union[PluginID, str] = PluginID.LLAMA_CPP,
+                   device_id: Optional[str] = None,
+                   **kwargs
+        ) -> 'TTS':
+        """Load TTS 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.tts_impl.mlx_tts_impl import MLXTTSImpl
+            return MLXTTSImpl._load_from(model_path, vocoder_path, plugin_id, device_id)
+        else:
+            from nexaai.tts_impl.pybind_tts_impl import PyBindTTSImpl
+            return PyBindTTSImpl._load_from(model_path, vocoder_path, plugin_id, device_id)
+
+    @abstractmethod
+    def synthesize(
+        self,
+        text: str,
+        config: Optional[TTSConfig] = None,
+        output_path: Optional[str] = None,
+    ) -> TTSResult:
+        """Synthesize speech from text and save to filesystem."""
+        pass
+
+    @abstractmethod
+    def list_available_voices(self) -> List[str]:
+        """List available voices."""
+        pass

nexaai/tts_impl/mlx_tts_impl.py

@@ -0,0 +1,94 @@
+# Note: This code is generated by Cursor, not tested yet.
+
+from typing import List, Optional, Union
+import os
+
+from nexaai.common import PluginID
+from nexaai.tts import TTS, TTSConfig, TTSResult
+from nexaai.mlx_backend.tts.interface import MlxTts as MLXTTSInterface
+
+
+class MLXTTSImpl(TTS):
+    def __init__(self):
+        """Initialize MLX TTS implementation."""
+        super().__init__()
+        self._mlx_tts = None
+
+    @classmethod
+    def _load_from(cls,
+                   model_path: str,
+                   vocoder_path: str,
+                   plugin_id: Union[PluginID, str] = PluginID.MLX,
+                   device_id: Optional[str] = None
+        ) -> 'MLXTTSImpl':
+        """Load TTS model from local path using MLX backend."""
+        try:
+            # MLX TTS interface is already imported
+            
+            # Create instance and load MLX TTS
+            instance = cls()
+            instance._mlx_tts = MLXTTSInterface(
+                model_path=model_path,
+                vocoder_path=vocoder_path,
+                device=device_id
+            )
+            
+            return instance
+        except Exception as e:
+            raise RuntimeError(f"Failed to load MLX TTS: {str(e)}")
+
+    def eject(self):
+        """Destroy the model and free resources."""
+        if self._mlx_tts:
+            self._mlx_tts.destroy()
+            self._mlx_tts = None
+
+    def synthesize(
+        self,
+        text: str,
+        config: Optional[TTSConfig] = None,
+        output_path: Optional[str] = None,
+    ) -> TTSResult:
+        """Synthesize speech from text and save to filesystem."""
+        if not self._mlx_tts:
+            raise RuntimeError("MLX TTS not loaded")
+        
+        try:
+            # Convert our config to MLX format if provided
+            mlx_config = None
+            if config:
+                from nexaai.mlx_backend.ml import TTSConfig as MLXTTSConfig
+                
+                mlx_config = MLXTTSConfig(
+                    voice=config.voice,
+                    speed=config.speed,
+                    seed=config.seed,
+                    sample_rate=config.sample_rate
+                )
+            
+            # Use MLX TTS synthesis
+            result = self._mlx_tts.synthesize(text, mlx_config, output_path)
+            
+            # Convert MLX result to our format
+            return TTSResult(
+                audio_path=result.audio_path,
+                duration_seconds=result.duration_seconds,
+                sample_rate=result.sample_rate,
+                channels=result.channels,
+                num_samples=result.num_samples
+            )
+            
+        except Exception as e:
+            raise RuntimeError(f"Failed to synthesize speech: {str(e)}")
+
+    def list_available_voices(self) -> List[str]:
+        """List available voices."""
+        if not self._mlx_tts:
+            raise RuntimeError("MLX TTS not loaded")
+        
+        try:
+            return self._mlx_tts.list_available_voices()
+        except Exception as e:
+            raise RuntimeError(f"Failed to list available voices: {str(e)}")
+
+

nexaai/tts_impl/pybind_tts_impl.py

@@ -0,0 +1,43 @@
+from typing import List, Optional, Union
+
+from nexaai.common import PluginID
+from nexaai.tts import TTS, TTSConfig, TTSResult
+
+
+class PyBindTTSImpl(TTS):
+    def __init__(self):
+        """Initialize PyBind TTS implementation."""
+        super().__init__()
+        # TODO: Add PyBind-specific initialization
+
+    @classmethod
+    def _load_from(cls,
+                   model_path: str,
+                   vocoder_path: str,
+                   plugin_id: Union[PluginID, str] = PluginID.LLAMA_CPP,
+                   device_id: Optional[str] = None
+        ) -> 'PyBindTTSImpl':
+        """Load TTS model from local path using PyBind backend."""
+        # TODO: Implement PyBind TTS loading
+        instance = cls()
+        return instance
+
+    def eject(self):
+        """Destroy the model and free resources."""
+        # TODO: Implement PyBind TTS cleanup
+        pass
+
+    def synthesize(
+        self,
+        text: str,
+        config: Optional[TTSConfig] = None,
+        output_path: Optional[str] = None,
+    ) -> TTSResult:
+        """Synthesize speech from text and save to filesystem."""
+        # TODO: Implement PyBind TTS synthesis
+        raise NotImplementedError("PyBind TTS synthesis not yet implemented")
+
+    def list_available_voices(self) -> List[str]:
+        """List available voices."""
+        # TODO: Implement PyBind TTS voice listing
+        raise NotImplementedError("PyBind TTS voice listing not yet implemented")

nexaai/utils/decode.py

@@ -0,0 +1,18 @@
+"""
+Utility functions for text decoding operations.
+"""
+
+
+def safe_decode(data):
+    """
+    Safely decode bytes or text, handling UTF-8 errors.
+    
+    Args:
+        data: Input data that can be bytes or text
+        
+    Returns:
+        str: Decoded string with errors replaced if any
+    """
+    if isinstance(data, bytes):
+        return data.decode('utf-8', errors='replace')
+    return str(data)