nexaai 1.0.16rc13__cp310-cp310-macosx_15_0_universal2.whl

nexaai/rerank.py

@@ -0,0 +1,55 @@
+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,
+                   tokenizer_file: str = "tokenizer.json",
+                   plugin_id: Union[PluginID, str] = PluginID.LLAMA_CPP,
+                   device_id: Optional[str] = None
+        ) -> '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, tokenizer_file, plugin_id, device_id)
+        else:
+            from nexaai.rerank_impl.pybind_rerank_impl import PyBindRerankImpl
+            return PyBindRerankImpl._load_from(model_path, 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,92 @@
+# 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,
+                   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,
+                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,43 @@
+from typing import List, Optional, Sequence, Union
+
+from nexaai.common import PluginID
+from nexaai.rerank import Reranker, RerankConfig
+
+
+class PyBindRerankImpl(Reranker):
+    def __init__(self):
+        """Initialize PyBind Rerank implementation."""
+        super().__init__()
+        # TODO: Add PyBind-specific initialization
+
+    @classmethod
+    def _load_from(cls,
+                   model_path: str,
+                   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."""
+        # TODO: Implement PyBind reranker loading
+        instance = cls()
+        return instance
+
+    def eject(self):
+        """Destroy the model and free resources."""
+        # TODO: Implement PyBind reranker cleanup
+        pass
+
+    def load_model(self, model_path: str, extra_data: Optional[str] = None) -> bool:
+        """Load model from path."""
+        # TODO: Implement PyBind reranker model loading
+        raise NotImplementedError("PyBind reranker model loading not yet implemented")
+
+    def rerank(
+        self,
+        query: str,
+        documents: Sequence[str],
+        config: Optional[RerankConfig] = None,
+    ) -> List[float]:
+        """Rerank documents given a query."""
+        # TODO: Implement PyBind reranking
+        raise NotImplementedError("PyBind reranking not yet implemented")

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/tts.py

@@ -0,0 +1,74 @@
+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
+        ) -> '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/avatar_fetcher.py

@@ -0,0 +1,104 @@
+"""Utility for fetching avatar URLs from HuggingFace."""
+
+import logging
+from typing import Dict, Optional
+import httpx
+
+logger = logging.getLogger(__name__)
+
+
+def fetch_avatar_urls_from_hf_api(query: str, custom_endpoint: Optional[str] = None) -> Dict[str, str]:
+    """
+    Fetch avatar URLs from HuggingFace models-json endpoint.
+    
+    Args:
+        query: Search query to fetch models for
+        custom_endpoint: Optional custom HuggingFace endpoint
+        
+    Returns:
+        Dictionary mapping author names to avatar URLs
+    """
+    avatar_map = {}
+    try:
+        # Use the base URL from the configured endpoint
+        base_url = custom_endpoint if custom_endpoint else "https://huggingface.co"
+        
+        # Build the URL with query parameter
+        url = f"{base_url}/models-json?sort=trending&search={query}&withCount=true"
+        
+        # Make the HTTP request with a timeout
+        with httpx.Client(timeout=2.0) as client:
+            response = client.get(url)
+            
+            if response.status_code == 200:
+                data = response.json()
+                models = data.get("models", [])
+                
+                # Build a map of author names to avatar URLs
+                for model in models:
+                    author = model.get("author")
+                    author_data = model.get("authorData", {})
+                    avatar_url = author_data.get("avatarUrl")
+                    
+                    if author and avatar_url:
+                        # Handle relative URLs by prepending appropriate base URL
+                        if avatar_url.startswith("/"):
+                            avatar_url = f"{base_url}{avatar_url}"
+                        avatar_map[author] = avatar_url
+                
+                logger.debug(f"Fetched {len(avatar_map)} avatar URLs from HuggingFace API")
+            else:
+                logger.warning(f"Failed to fetch avatar URLs: HTTP {response.status_code}")
+                
+    except Exception as e:
+        logger.warning(f"Error fetching avatar URLs from HuggingFace API: {e}")
+        # Return empty map on error - we'll fall back to default behavior
+    
+    return avatar_map
+
+
+def get_avatar_url_for_repo(repo_id: str, search_query: Optional[str] = None, 
+                            custom_endpoint: Optional[str] = None) -> Optional[str]:
+    """
+    Get avatar URL for a repository ID.
+    
+    This method tries multiple strategies:
+    1. If search_query is provided, fetch from HuggingFace API with that query
+    2. Try fetching with the full repo_id as query
+    3. Try fetching with just the organization name as query
+    4. Fall back to CDN URL pattern
+    
+    Args:
+        repo_id: Repository ID in format "owner/repo"
+        search_query: Optional search query to use for fetching avatars
+        custom_endpoint: Optional custom HuggingFace endpoint
+        
+    Returns:
+        Avatar URL or None if not found
+    """
+    if "/" not in repo_id:
+        return None
+        
+    org_name = repo_id.split("/")[0]
+    
+    # Try with search query if provided
+    if search_query:
+        avatar_map = fetch_avatar_urls_from_hf_api(search_query, custom_endpoint)
+        avatar_url = avatar_map.get(org_name)
+        if avatar_url:
+            return avatar_url
+    
+    # Try with full repo_id
+    avatar_map = fetch_avatar_urls_from_hf_api(repo_id, custom_endpoint)
+    avatar_url = avatar_map.get(org_name)
+    if avatar_url:
+        return avatar_url
+    
+    # Try with just organization name
+    avatar_map = fetch_avatar_urls_from_hf_api(org_name, custom_endpoint)
+    avatar_url = avatar_map.get(org_name)
+    if avatar_url:
+        return avatar_url
+    
+    # Fallback to CDN URL pattern
+    return f"https://cdn-thumbnails.huggingface.co/social-thumbnails/{org_name}.png"

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)