nexaai 1.0.29__cp310-cp310-macosx_14_0_universal2.whl

nexaai/__init__.py

@@ -0,0 +1,99 @@
+"""
+NexaAI Python bindings for NexaSDK C-lib backend.
+"""
+
+import sys
+import os
+
+# Add mlx_backend to Python path as individual module (only if it exists)
+_current_dir = os.path.dirname(os.path.abspath(__file__))
+_mlx_backend_path = os.path.join(_current_dir, "mlx_backend")
+# Only add to path if the directory exists (it won't on Windows)
+if os.path.exists(_mlx_backend_path) and _mlx_backend_path not in sys.path:
+    sys.path.insert(0, _mlx_backend_path)
+
+try:
+    from ._version import __version__
+except ImportError:
+    # Fallback for development or when version file hasn't been generated yet
+    __version__ = "0.0.1"
+
+# Import common configuration classes first (no external dependencies)
+from .common import ModelConfig, GenerationConfig, ChatMessage, SamplerConfig, PluginID
+
+# Import logging functionality
+from .log import set_logger, get_error_message
+
+# Import runtime errors
+from .runtime_error import (
+    NexaRuntimeError,
+    ContextLengthExceededError,
+    GenerationError,
+)
+
+# Create alias for PluginID to be accessible as plugin_id
+plugin_id = PluginID
+
+# Import new feature classes (no external dependencies in base classes)
+from .llm import LLM
+from .embedder import Embedder, EmbeddingConfig
+from .vlm import VLM
+from .asr import ASR, ASRConfig, ASRResult
+from .cv import CVModel, CVModelConfig, CVResult, CVResults, CVCapabilities, BoundingBox
+from .rerank import Reranker, RerankConfig
+from .image_gen import (
+    ImageGen,
+    ImageGenerationConfig,
+    ImageSamplerConfig,
+    SchedulerConfig,
+    Image,
+)
+from .tts import TTS, TTSConfig, TTSSamplerConfig, TTSResult
+from .diarize import Diarize, DiarizeConfig, DiarizeResult, SpeechSegment
+
+# Build __all__ list dynamically
+__all__ = [
+    "__version__",
+    # Common configurations (always available)
+    "ModelConfig",
+    "GenerationConfig",
+    "ChatMessage",
+    "SamplerConfig",
+    "EmbeddingConfig",
+    "PluginID",
+    "plugin_id",
+    # Logging functionality
+    "set_logger",
+    "get_error_message",
+    # Runtime errors
+    "NexaRuntimeError",
+    "ContextLengthExceededError",
+    "GenerationError",
+    "LLM",
+    "Embedder",
+    "VLM",
+    "ASR",
+    "CVModel",
+    "Reranker",
+    "ImageGen",
+    "TTS",
+    "Diarize",
+    "ASRConfig",
+    "ASRResult",
+    "CVModelConfig",
+    "CVResult",
+    "CVResults",
+    "CVCapabilities",
+    "BoundingBox",
+    "RerankConfig",
+    "ImageGenerationConfig",
+    "ImageSamplerConfig",
+    "SchedulerConfig",
+    "Image",
+    "TTSConfig",
+    "TTSSamplerConfig",
+    "TTSResult",
+    "DiarizeConfig",
+    "DiarizeResult",
+    "SpeechSegment",
+]

nexaai/_version.py

@@ -0,0 +1,4 @@
+# This file is generated by CMake from _version.py.in
+# Do not modify this file manually - it will be overwritten
+
+__version__ = "1.0.29" 

nexaai/asr.py

@@ -0,0 +1,68 @@
+from typing import List, Optional, Sequence, Tuple, Union
+from abc import abstractmethod
+from dataclasses import dataclass
+
+from nexaai.base import BaseModel
+from nexaai.common import PluginID, ModelConfig
+
+
+@dataclass
+class ASRConfig:
+    """Configuration for ASR."""
+    timestamps: str = "none"  # "none" | "segment" | "word"
+    beam_size: int = 5
+    stream: bool = False
+
+
+@dataclass
+class ASRResult:
+    """Result from ASR processing."""
+    transcript: str
+    confidence_scores: Sequence[float]
+    timestamps: Sequence[Tuple[float, float]]
+
+
+class ASR(BaseModel):
+    """Abstract base class for Automatic Speech Recognition models."""
+
+    def __init__(self, m_cfg: ModelConfig = ModelConfig()):
+        """Initialize base ASR class."""
+        self._m_cfg = m_cfg
+
+    @classmethod
+    def _load_from(cls,
+                   model_path: str,
+                   model_name: Optional[str] = None,
+                   tokenizer_path: Optional[str] = None,
+                   language: Optional[str] = None,
+                   m_cfg: ModelConfig = ModelConfig(),
+                   plugin_id: Union[PluginID, str] = PluginID.LLAMA_CPP,
+                   device_id: Optional[str] = None,
+                   **kwargs
+        ) -> 'ASR':
+        """Load ASR 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.asr_impl.mlx_asr_impl import MLXASRImpl
+            return MLXASRImpl._load_from(model_path, model_name, tokenizer_path, language, m_cfg, plugin_id, device_id)
+        else:
+            from nexaai.asr_impl.pybind_asr_impl import PyBindASRImpl
+            return PyBindASRImpl._load_from(model_path, model_name, tokenizer_path, language, m_cfg, plugin_id, device_id)
+
+
+    @abstractmethod
+    def transcribe(
+        self,
+        audio_path: str,
+        language: Optional[str] = None,
+        config: Optional[ASRConfig] = None,
+    ) -> ASRResult:
+        """Transcribe audio file to text."""
+        pass
+
+    @abstractmethod
+    def list_supported_languages(self) -> List[str]:
+        """List supported languages."""
+        pass

nexaai/asr_impl/mlx_asr_impl.py

@@ -0,0 +1,93 @@
+# Note: This code is generated by Cursor, not tested yet.
+
+from typing import List, Optional, Union
+
+from nexaai.common import PluginID, ModelConfig
+from nexaai.asr import ASR, ASRConfig, ASRResult
+from nexaai.mlx_backend.asr.interface import MlxAsr as MLXASRInterface
+from nexaai.mlx_backend.ml import ModelConfig as MLXModelConfig, SamplerConfig as MLXSamplerConfig, GenerationConfig as MLXGenerationConfig, EmbeddingConfig
+
+
+class MLXASRImpl(ASR):
+    def __init__(self, m_cfg: ModelConfig = ModelConfig()):
+        """Initialize MLX ASR implementation."""
+        super().__init__(m_cfg)
+        self._mlx_asr = None
+
+    @classmethod
+    def _load_from(cls,
+                   model_path: str,
+                   model_name: Optional[str] = None,
+                   tokenizer_path: Optional[str] = None,
+                   language: Optional[str] = None,
+                   m_cfg: 'ModelConfig' = None,
+                   plugin_id: Union[PluginID, str] = PluginID.MLX,
+                   device_id: Optional[str] = None
+        ) -> 'MLXASRImpl':
+        """Load ASR model from local path using MLX backend."""
+        try:
+            
+            # Create instance and load MLX ASR
+            instance = cls(m_cfg)
+            instance._mlx_asr = MLXASRInterface(
+                model_path=model_path,
+                tokenizer_path=tokenizer_path,
+                language=language,
+                device=device_id
+            )
+            
+            return instance
+        except Exception as e:
+            raise RuntimeError(f"Failed to load MLX ASR: {str(e)}")
+
+    def eject(self):
+        """Destroy the model and free resources."""
+        if self._mlx_asr:
+            self._mlx_asr.destroy()
+            self._mlx_asr = None
+            
+    def transcribe(
+        self,
+        audio_path: str,
+        language: Optional[str] = None,
+        config: Optional[ASRConfig] = None,
+    ) -> ASRResult:
+        """Transcribe audio file to text."""
+        if not self._mlx_asr:
+            raise RuntimeError("MLX ASR not loaded")
+        
+        try:
+            # Convert our config to MLX format if provided
+            mlx_config = None
+            if config:
+                from nexaai.mlx_backend.ml import ASRConfig as MLXASRConfig
+                
+                mlx_config = MLXASRConfig()
+                mlx_config.timestamps = config.timestamps
+                mlx_config.beam_size = config.beam_size
+                mlx_config.stream = config.stream
+            
+            # Use MLX ASR transcription
+            result = self._mlx_asr.transcribe(audio_path, language, mlx_config)
+            
+            # Convert MLX result to our format
+            return ASRResult(
+                transcript=result.transcript,
+                confidence_scores=result.confidence_scores,
+                timestamps=result.timestamps
+            )
+            
+        except Exception as e:
+            raise RuntimeError(f"Failed to transcribe audio: {str(e)}")
+
+    def list_supported_languages(self) -> List[str]:
+        """List supported languages."""
+        if not self._mlx_asr:
+            raise RuntimeError("MLX ASR not loaded")
+        
+        try:
+            return self._mlx_asr.list_supported_languages()
+        except Exception as e:
+            raise RuntimeError(f"Failed to list supported languages: {str(e)}")
+
+

nexaai/asr_impl/pybind_asr_impl.py

@@ -0,0 +1,127 @@
+from typing import List, Optional, Union
+
+from nexaai.common import PluginID, ModelConfig
+from nexaai.asr import ASR, ASRConfig, ASRResult
+from nexaai.binds import asr_bind, common_bind
+from nexaai.runtime import _ensure_runtime
+
+
+class PyBindASRImpl(ASR):
+    def __init__(self, handle: any, m_cfg: ModelConfig = ModelConfig()):
+        """Private constructor, should not be called directly."""
+        super().__init__(m_cfg)
+        self._handle = handle  # This is a py::capsule
+        self._model_config = None
+
+    @classmethod
+    def _load_from(cls,
+                   model_path: str,
+                   model_name: Optional[str] = None,
+                   tokenizer_path: Optional[str] = None,
+                   language: Optional[str] = None,
+                   m_cfg: ModelConfig = ModelConfig(),
+                   plugin_id: Union[PluginID, str] = PluginID.LLAMA_CPP,
+                   device_id: Optional[str] = None
+        ) -> 'PyBindASRImpl':
+        """Load ASR model from local path using PyBind backend."""
+        _ensure_runtime()
+        
+        # Create model config
+        config = common_bind.ModelConfig()
+        
+        config.n_ctx = m_cfg.n_ctx
+        if m_cfg.n_threads is not None:
+            config.n_threads = m_cfg.n_threads
+        if m_cfg.n_threads_batch is not None:
+            config.n_threads_batch = m_cfg.n_threads_batch
+        if m_cfg.n_batch is not None:
+            config.n_batch = m_cfg.n_batch
+        if m_cfg.n_ubatch is not None:
+            config.n_ubatch = m_cfg.n_ubatch
+        if m_cfg.n_seq_max is not None:
+            config.n_seq_max = m_cfg.n_seq_max
+        config.n_gpu_layers = m_cfg.n_gpu_layers
+        
+        # handle chat template strings
+        if m_cfg.chat_template_path:
+            config.chat_template_path = m_cfg.chat_template_path
+        
+        if m_cfg.chat_template_content:
+            config.chat_template_content = m_cfg.chat_template_content
+        
+        # Convert plugin_id to string
+        plugin_id_str = plugin_id.value if isinstance(plugin_id, PluginID) else str(plugin_id)
+        
+        # Create ASR handle using the binding
+        handle = asr_bind.ml_asr_create(
+            model_path=model_path,
+            model_name=model_name,
+            tokenizer_path=tokenizer_path,
+            model_config=config,
+            language=language,
+            plugin_id=plugin_id_str,
+            device_id=device_id,
+            license_id=None,  # Optional
+            license_key=None  # Optional
+        )
+        
+        return cls(handle, m_cfg)
+
+    def eject(self):
+        """Release the model from memory."""
+        # py::capsule handles cleanup automatically
+        if hasattr(self, '_handle') and self._handle is not None:
+            del self._handle
+            self._handle = None
+
+    def transcribe(
+        self,
+        audio_path: str,
+        language: Optional[str] = None,
+        config: Optional[ASRConfig] = None,
+    ) -> ASRResult:
+        """Transcribe audio file to text."""
+        if self._handle is None:
+            raise RuntimeError("ASR model not loaded. Call _load_from first.")
+        
+        # Convert ASRConfig to binding format if provided
+        asr_config = None
+        if config:
+            asr_config = asr_bind.ASRConfig()
+            asr_config.timestamps = config.timestamps
+            asr_config.beam_size = config.beam_size
+            asr_config.stream = config.stream
+        
+        # Perform transcription using the binding
+        result_dict = asr_bind.ml_asr_transcribe(
+            handle=self._handle,
+            audio_path=audio_path,
+            language=language,
+            config=asr_config
+        )
+        
+        # Convert result to ASRResult
+        transcript = result_dict.get("transcript", "")
+        confidence_scores = result_dict.get("confidence_scores")
+        timestamps = result_dict.get("timestamps")
+        
+        # Convert timestamps to the expected format
+        timestamp_pairs = []
+        if timestamps:
+            for start, end in timestamps:
+                timestamp_pairs.append((float(start), float(end)))
+        
+        return ASRResult(
+            transcript=transcript,
+            confidence_scores=confidence_scores or [],
+            timestamps=timestamp_pairs
+        )
+
+    def list_supported_languages(self) -> List[str]:
+        """List supported languages."""
+        if self._handle is None:
+            raise RuntimeError("ASR model not loaded. Call _load_from first.")
+        
+        # Get supported languages using the binding
+        languages = asr_bind.ml_asr_list_supported_languages(handle=self._handle)
+        return languages

nexaai/base.py

@@ -0,0 +1,39 @@
+from abc import ABC, abstractmethod
+from nexaai.common import ProfilingData
+from nexaai.utils.model_manager import auto_download_model
+
+class BaseModel(ABC):
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.eject()
+
+    def __del__(self):
+        self.eject()
+
+    @classmethod
+    @auto_download_model
+    def from_(cls, name_or_path: str, **kwargs) -> "BaseModel":
+        """
+        initialize model from (1) HF (2) if not found, then from local path
+        """
+
+        return cls._load_from(name_or_path, **kwargs)
+
+    @classmethod
+    @abstractmethod
+    def _load_from(cls, name_or_path: str, **kwargs) -> "BaseModel":
+        """
+        Model-specific loading logic. Must be implemented by each model type.
+        Called after model is available locally.
+        """  
+        pass
+
+    @abstractmethod
+    def eject(self):
+        pass
+
+    def get_profiling_data(self) -> ProfilingData:
+        pass

nexaai/binds/__init__.py

@@ -0,0 +1,7 @@
+from .common_bind import *
+from .llm_bind import *
+from .embedder_bind import *
+from .vlm_bind import *
+from .rerank_bind import *
+from .asr_bind import *
+from .diarize_bind import *