nexaai 1.0.29__cp310-cp310-macosx_14_0_universal2.whl

nexaai/cv_impl/pybind_cv_impl.py

@@ -0,0 +1,124 @@
+from typing import Optional, Union, List
+import os
+
+from nexaai.common import PluginID, ModelConfig
+from nexaai.cv import CVModel, CVModelConfig, CVResults, CVResult, BoundingBox, CVCapabilities
+from nexaai.binds import cv_bind, common_bind
+from nexaai.runtime import _ensure_runtime
+
+
+class PyBindCVImpl(CVModel):
+    def __init__(self, handle: any, m_cfg: ModelConfig = ModelConfig()):
+        """Private constructor, should not be called directly."""
+        super().__init__()
+        self._handle = handle  # This is a py::capsule
+        self._model_config = None
+
+    @classmethod
+    def _load_from(cls,
+                   local_path: str,  # This is the local path after auto_download_model processing
+                   model_name: Optional[str] = None,
+                   m_cfg: CVModelConfig = CVModelConfig(CVCapabilities.OCR),
+                   plugin_id: Union[PluginID, str] = PluginID.LLAMA_CPP,
+                   device_id: Optional[str] = None,
+                   **kwargs
+                   ) -> 'PyBindCVImpl':
+        """Load CV model from configuration using PyBind backend."""
+        _ensure_runtime()
+
+        config = cv_bind.CVModelConfig()
+        config.capabilities = cv_bind.CVCapabilities(m_cfg.capabilities)
+        if m_cfg.det_model_path is not None:
+            config.det_model_path = m_cfg.det_model_path
+        else:
+            config.det_model_path = local_path
+
+        print("local_path: ", local_path)
+        print("m_cfg.rec_model_path: ", m_cfg.rec_model_path)
+        if m_cfg.rec_model_path is not None:
+            config.rec_model_path = m_cfg.rec_model_path
+        else:
+            config.rec_model_path = local_path
+        print("config.rec_model_path: ", config.rec_model_path)
+        
+        if m_cfg.char_dict_path is not None:
+            config.char_dict_path = m_cfg.char_dict_path
+
+        if m_cfg.model_path is not None:
+            config.model_path = m_cfg.model_path
+
+        if m_cfg.system_library_path is not None:
+            config.system_library_path = m_cfg.system_library_path
+
+        plugin_id_str = plugin_id.value if isinstance(
+            plugin_id, PluginID) else str(plugin_id)
+
+        model_name_to_use = model_name if model_name else local_path
+        handle = cv_bind.ml_cv_create(
+            model_name=model_name_to_use,
+            config=config,
+            plugin_id=plugin_id_str,
+            device_id=device_id,
+            license_id=None,
+            license_key=None
+        )
+
+        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 infer(self, input_image_path: str) -> CVResults:
+        """Perform inference on image."""
+        if self._handle is None:
+            raise RuntimeError("CV model not loaded. Call _load_from first.")
+
+        if not os.path.exists(input_image_path):
+            raise FileNotFoundError(
+                f"Input image not found: {input_image_path}")
+
+        try:
+            # Perform inference using the binding
+            result_dict = cv_bind.ml_cv_infer(
+                handle=self._handle,
+                input_image_path=input_image_path
+            )
+
+            # Convert result dictionary to CVResults
+            results = []
+            for result_data in result_dict["results"]:
+                # Create bounding box if present
+                bbox = None
+                if "bbox" in result_data and result_data["bbox"] is not None:
+                    bbox_data = result_data["bbox"]
+                    bbox = BoundingBox(
+                        x=bbox_data["x"],
+                        y=bbox_data["y"],
+                        width=bbox_data["width"],
+                        height=bbox_data["height"]
+                    )
+
+                # Create CV result
+                cv_result = CVResult(
+                    image_paths=result_data.get("image_paths"),
+                    image_count=result_data.get("image_count", 0),
+                    class_id=result_data.get("class_id", 0),
+                    confidence=result_data.get("confidence", 0.0),
+                    bbox=bbox,
+                    text=result_data.get("text"),
+                    embedding=result_data.get("embedding"),
+                    embedding_dim=result_data.get("embedding_dim", 0)
+                )
+                results.append(cv_result)
+
+            return CVResults(
+                results=results,
+                result_count=result_dict["result_count"]
+            )
+
+        except Exception as e:
+            raise RuntimeError(f"CV inference failed: {str(e)}")

nexaai/diarize.py

@@ -0,0 +1,80 @@
+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, ModelConfig
+
+
+@dataclass
+class DiarizeConfig:
+    """Configuration for speaker diarization."""
+
+    min_speakers: int = 0  # Minimum number of speakers (0 = auto-detect)
+    max_speakers: int = 0  # Maximum number of speakers (0 = no limit)
+
+
+@dataclass
+class SpeechSegment:
+    """Speech segment with speaker label and timestamps."""
+
+    start_time: float  # Segment start time in seconds
+    end_time: float  # Segment end time in seconds
+    speaker_label: str  # Speaker label (e.g., "SPEAKER_00")
+
+
+@dataclass
+class DiarizeResult:
+    """Result from speaker diarization."""
+
+    segments: Sequence[SpeechSegment]  # Array of speech segments
+    segment_count: int  # Number of segments
+    num_speakers: int  # Total unique speakers detected
+    duration: float  # Total audio duration in seconds
+
+
+class Diarize(BaseModel):
+    """Abstract base class for speaker diarization models."""
+
+    def __init__(self, m_cfg: ModelConfig = ModelConfig()):
+        """Initialize base Diarize class."""
+        self._m_cfg = m_cfg
+
+    @classmethod
+    def _load_from(
+        cls,
+        model_path: str,
+        model_name: Optional[str] = None,
+        m_cfg: ModelConfig = ModelConfig(),
+        plugin_id: Union[PluginID, str] = PluginID.LLAMA_CPP,
+        device_id: Optional[str] = None,
+        **kwargs
+    ) -> "Diarize":
+        """Load diarization model from local path using PyBind backend."""
+        from nexaai.diarize_impl.pybind_diarize_impl import PyBindDiarizeImpl
+
+        # There is no MLX implementation for diarization as of now.
+        return PyBindDiarizeImpl._load_from(
+            model_path, model_name, m_cfg, plugin_id, device_id
+        )
+
+    @abstractmethod
+    def infer(
+        self,
+        audio_path: str,
+        config: Optional[DiarizeConfig] = None,
+    ) -> DiarizeResult:
+        """
+        Perform speaker diarization on audio file.
+
+        Determines "who spoke when" in the audio recording, producing time-stamped segments
+        with speaker labels. Segments are time-ordered and non-overlapping.
+
+        Args:
+            audio_path: Path to audio file
+            config: Optional diarization configuration
+
+        Returns:
+            DiarizeResult with segments, speaker count, and duration
+        """
+        pass

nexaai/diarize_impl/__init__.py

@@ -0,0 +1 @@
+

nexaai/diarize_impl/pybind_diarize_impl.py

@@ -0,0 +1,125 @@
+from typing import Any, Optional, Union
+
+from nexaai.common import PluginID, ModelConfig
+from nexaai.diarize import Diarize, DiarizeConfig, DiarizeResult, SpeechSegment
+from nexaai.binds import diarize_bind, common_bind
+from nexaai.runtime import _ensure_runtime
+
+
+class PyBindDiarizeImpl(Diarize):
+    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,
+        m_cfg: ModelConfig = ModelConfig(),
+        plugin_id: Union[PluginID, str] = PluginID.LLAMA_CPP,
+        device_id: Optional[str] = None,
+    ) -> "PyBindDiarizeImpl":
+        """Load diarization 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 needed for diarization)
+        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 Diarize handle using the binding
+        handle = diarize_bind.ml_diarize_create(
+            model_path=model_path,
+            model_name=model_name,
+            model_config=config,
+            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 infer(
+        self,
+        audio_path: str,
+        config: Optional[DiarizeConfig] = None,
+    ) -> DiarizeResult:
+        """
+        Perform speaker diarization on audio file.
+
+        Determines "who spoke when" in the audio recording, producing time-stamped segments
+        with speaker labels. Segments are time-ordered and non-overlapping.
+
+        Args:
+            audio_path: Path to audio file
+            config: Optional diarization configuration
+
+        Returns:
+            DiarizeResult with segments, speaker count, and duration
+        """
+        if self._handle is None:
+            raise RuntimeError("Diarization model not loaded. Call _load_from first.")
+
+        # Convert DiarizeConfig to binding format if provided
+        diarize_config = None
+        if config:
+            diarize_config = diarize_bind.DiarizeConfig()
+            diarize_config.min_speakers = config.min_speakers
+            diarize_config.max_speakers = config.max_speakers
+
+        # Perform diarization using the binding
+        result_dict = diarize_bind.ml_diarize_infer(
+            handle=self._handle, audio_path=audio_path, config=diarize_config
+        )
+
+        # Convert result to DiarizeResult
+        segments = []
+        for segment_dict in result_dict.get("segments", []):
+            segments.append(
+                SpeechSegment(
+                    start_time=float(segment_dict["start_time"]),
+                    end_time=float(segment_dict["end_time"]),
+                    speaker_label=segment_dict["speaker_label"],
+                )
+            )
+
+        return DiarizeResult(
+            segments=segments,
+            segment_count=result_dict.get("segment_count", 0),
+            num_speakers=result_dict.get("num_speakers", 0),
+            duration=result_dict.get("duration", 0.0),
+        )

nexaai/embedder.py

@@ -0,0 +1,73 @@
+from typing import List, Union
+from dataclasses import dataclass
+from abc import abstractmethod
+import numpy as np
+
+from nexaai.base import BaseModel
+from nexaai.common import PluginID
+
+
+@dataclass
+class EmbeddingConfig:
+    batch_size: int = 32
+    normalize: bool = True
+    normalize_method: str = "l2"
+
+
+class Embedder(BaseModel):
+    def __init__(self):
+        """
+        Internal initializer
+        """
+        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, **kwargs):
+        """
+        Load an embedder from model files, routing to appropriate implementation.
+
+        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:
+            Embedder instance
+        """
+        # 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.embedder_impl.mlx_embedder_impl import MLXEmbedderImpl
+            return MLXEmbedderImpl._load_from(model_path, model_name, tokenizer_file, plugin_id)
+        else:
+            from nexaai.embedder_impl.pybind_embedder_impl import PyBindEmbedderImpl
+            return PyBindEmbedderImpl._load_from(model_path, model_name, tokenizer_file, plugin_id)
+
+    @abstractmethod
+    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)
+        """
+        pass
+
+    @abstractmethod
+    def get_embedding_dim(self) -> int:
+        """
+        Get the embedding dimension of the model
+
+        Returns:
+            The embedding dimension in int
+        """
+        pass

nexaai/embedder_impl/mlx_embedder_impl.py

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

nexaai/embedder_impl/pybind_embedder_impl.py

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