nexaai 1.0.4rc13__cp310-cp310-win_amd64.whl

nexaai/llm_impl/mlx_llm_impl.py

@@ -0,0 +1,249 @@
+from typing import Generator, Optional, Any
+
+from nexaai.common import ModelConfig, GenerationConfig, ChatMessage
+from nexaai.llm import LLM
+from nexaai.mlx_backend.llm.interface import LLM as MLXLLMInterface
+from nexaai.mlx_backend.ml import ModelConfig as MLXModelConfig, SamplerConfig as MLXSamplerConfig, GenerationConfig as MLXGenerationConfig, EmbeddingConfig
+
+
+class MLXLLMImpl(LLM):
+    def __init__(self, m_cfg: ModelConfig = ModelConfig()):
+        """Initialize MLX LLM implementation."""
+        super().__init__(m_cfg)
+        self._mlx_llm = None
+
+    @classmethod
+    def _load_from(cls,
+                   local_path: str,
+                   tokenizer_path: Optional[str] = None,
+                   m_cfg: ModelConfig = ModelConfig(),
+                   plugin_id: str = "mlx",
+                   device_id: Optional[str] = None
+        ) -> 'MLXLLMImpl':
+        """Load model from local path using MLX backend."""
+        try:
+            # MLX interface and configs are already imported
+            
+            # Convert our ModelConfig to MLX ModelConfig
+            mlx_config = MLXModelConfig()
+            mlx_config.n_ctx = m_cfg.n_ctx
+            mlx_config.n_threads = m_cfg.n_threads
+            mlx_config.n_threads_batch = m_cfg.n_threads_batch
+            mlx_config.n_batch = m_cfg.n_batch
+            mlx_config.n_ubatch = m_cfg.n_ubatch
+            mlx_config.n_seq_max = m_cfg.n_seq_max
+            mlx_config.chat_template_path = m_cfg.chat_template_path
+            mlx_config.chat_template_content = m_cfg.chat_template_content
+            
+            # Create instance and load MLX model
+            instance = cls(m_cfg)
+            instance._mlx_llm = MLXLLMInterface(
+                model_path=local_path,
+                tokenizer_path=tokenizer_path or local_path,
+                config=mlx_config,
+                device=device_id
+            )
+            
+            return instance
+        except Exception as e:
+            raise RuntimeError(f"Failed to load MLX LLM: {str(e)}")
+
+    def eject(self):
+        """Release the model from memory."""
+        if self._mlx_llm:
+            self._mlx_llm.destroy()
+            self._mlx_llm = None
+
+    def apply_chat_template(self, messages: list[ChatMessage]) -> str:
+        """Apply the chat template to messages."""
+        if not self._mlx_llm:
+            raise RuntimeError("MLX LLM not loaded")
+        
+        try:
+            # Convert to MLX ChatMessage format
+            mlx_messages = []
+            for msg in messages:
+                # Create a simple object with role and content attributes
+                class MLXChatMessage:
+                    def __init__(self, role, content):
+                        self.role = role
+                        self.content = content
+                mlx_messages.append(MLXChatMessage(msg["role"], msg["content"]))
+            
+            return self._mlx_llm.apply_chat_template(mlx_messages)
+        except Exception as e:
+            raise RuntimeError(f"Failed to apply chat template: {str(e)}")
+
+    def generate_stream(self, prompt: str, g_cfg: GenerationConfig = GenerationConfig()) -> Generator[str, None, None]:
+        """Generate text with streaming."""
+        if not self._mlx_llm:
+            raise RuntimeError("MLX LLM not loaded")
+        
+        try:
+            import queue
+            import threading
+            
+            # Convert GenerationConfig to MLX format
+            
+            mlx_gen_config = MLXGenerationConfig()
+            mlx_gen_config.max_tokens = g_cfg.max_tokens
+            mlx_gen_config.stop = g_cfg.stop_words
+            mlx_gen_config.image_paths = g_cfg.image_paths
+            mlx_gen_config.audio_paths = g_cfg.audio_paths
+            
+            if g_cfg.sampler_config:
+                mlx_sampler_config = MLXSamplerConfig()
+                mlx_sampler_config.temperature = g_cfg.sampler_config.temperature
+                mlx_sampler_config.top_p = g_cfg.sampler_config.top_p
+                mlx_sampler_config.top_k = g_cfg.sampler_config.top_k
+                mlx_sampler_config.repetition_penalty = g_cfg.sampler_config.repetition_penalty
+                mlx_sampler_config.presence_penalty = g_cfg.sampler_config.presence_penalty
+                mlx_sampler_config.frequency_penalty = g_cfg.sampler_config.frequency_penalty
+                mlx_sampler_config.seed = g_cfg.sampler_config.seed
+                mlx_sampler_config.grammar_path = g_cfg.sampler_config.grammar_path
+                mlx_sampler_config.grammar_string = g_cfg.sampler_config.grammar_string
+                mlx_gen_config.sampler_config = mlx_sampler_config
+            
+            # Create a queue for streaming tokens
+            token_queue = queue.Queue()
+            exception_container = [None]
+            self.reset_cancel()  # Reset cancel flag before generation
+            
+            def token_callback(token: str, user_data: Any = None) -> bool:
+                if self._cancel_event.is_set():
+                    token_queue.put(('end', None))
+                    return False
+                try:
+                    token_queue.put(('token', token))
+                    return True
+                except Exception as e:
+                    exception_container[0] = e
+                    return False
+            
+            # Run generation in a separate thread
+            def generate():
+                try:
+                    self._mlx_llm.generate_stream(prompt, mlx_gen_config, token_callback)
+                except Exception as e:
+                    exception_container[0] = e
+                finally:
+                    token_queue.put(('end', None))
+            
+            thread = threading.Thread(target=generate)
+            thread.start()
+            
+            # Yield tokens as they come from the queue
+            while True:
+                if exception_container[0]:
+                    raise exception_container[0]
+                    
+                try:
+                    msg_type, token = token_queue.get(timeout=0.1)
+                    if msg_type == 'end':
+                        break
+                    elif msg_type == 'token':
+                        yield token
+                except queue.Empty:
+                    if not thread.is_alive():
+                        break
+                    continue
+                    
+            thread.join()
+            
+            if exception_container[0]:
+                raise exception_container[0]
+            
+        except Exception as e:
+            raise RuntimeError(f"Failed to generate streaming text: {str(e)}")
+
+    def generate(self, prompt: str, g_cfg: GenerationConfig = GenerationConfig()) -> str:
+        """
+        Generate text without streaming.
+
+        Args:
+            prompt (str): The prompt to generate text from.
+            g_cfg (GenerationConfig): Generation configuration.
+
+        Returns:
+            str: The generated text.
+        """
+        if not self._mlx_llm:
+            raise RuntimeError("MLX LLM not loaded")
+        
+        try:
+            # Convert GenerationConfig to MLX format
+            
+            mlx_gen_config = MLXGenerationConfig()
+            mlx_gen_config.max_tokens = g_cfg.max_tokens
+            mlx_gen_config.stop = g_cfg.stop_words
+            mlx_gen_config.image_paths = g_cfg.image_paths
+            mlx_gen_config.audio_paths = g_cfg.audio_paths
+            
+            if g_cfg.sampler_config:
+                mlx_sampler_config = MLXSamplerConfig()
+                mlx_sampler_config.temperature = g_cfg.sampler_config.temperature
+                mlx_sampler_config.top_p = g_cfg.sampler_config.top_p
+                mlx_sampler_config.top_k = g_cfg.sampler_config.top_k
+                mlx_sampler_config.repetition_penalty = g_cfg.sampler_config.repetition_penalty
+                mlx_sampler_config.presence_penalty = g_cfg.sampler_config.presence_penalty
+                mlx_sampler_config.frequency_penalty = g_cfg.sampler_config.frequency_penalty
+                mlx_sampler_config.seed = g_cfg.sampler_config.seed
+                mlx_sampler_config.grammar_path = g_cfg.sampler_config.grammar_path
+                mlx_sampler_config.grammar_string = g_cfg.sampler_config.grammar_string
+                mlx_gen_config.sampler_config = mlx_sampler_config
+            
+            # Simple token callback that just continues
+            def token_callback(token: str, user_data: Any = None) -> bool:
+                return not self._cancel_event.is_set()
+            
+            # Use MLX streaming generation and return the full result
+            return self._mlx_llm.generate_stream(prompt, mlx_gen_config, token_callback)
+            
+        except Exception as e:
+            raise RuntimeError(f"Failed to generate text: {str(e)}")
+
+    def save_kv_cache(self, path: str):
+        """
+        Save the key-value cache to the file.
+
+        Args:
+            path (str): The path to the file.
+        """
+        if not self._mlx_llm:
+            raise RuntimeError("MLX LLM not loaded")
+        
+        try:
+            success = self._mlx_llm.save_kv_cache(path)
+            if not success:
+                raise RuntimeError("Failed to save KV cache")
+        except Exception as e:
+            raise RuntimeError(f"Failed to save KV cache: {str(e)}")
+
+    def load_kv_cache(self, path: str):
+        """
+        Load the key-value cache from the file.
+
+        Args:
+            path (str): The path to the file.
+        """
+        if not self._mlx_llm:
+            raise RuntimeError("MLX LLM not loaded")
+        
+        try:
+            success = self._mlx_llm.load_kv_cache(path)
+            if not success:
+                raise RuntimeError("Failed to load KV cache")
+        except Exception as e:
+            raise RuntimeError(f"Failed to load KV cache: {str(e)}")
+
+    def reset(self):
+        """
+        Reset the LLM model context and KV cache.
+        """
+        if not self._mlx_llm:
+            raise RuntimeError("MLX LLM not loaded")
+        
+        try:
+            self._mlx_llm.reset()
+        except Exception as e:
+            raise RuntimeError(f"Failed to reset MLX LLM: {str(e)}")

nexaai/llm_impl/pybind_llm_impl.py

@@ -0,0 +1,207 @@
+from typing import Generator, Optional
+import queue
+import threading
+
+from nexaai.common import ModelConfig, GenerationConfig, ChatMessage
+from nexaai.binds import llm_bind, common_bind
+from nexaai.runtime import _ensure_runtime
+from nexaai.llm import LLM
+
+
+class PyBindLLMImpl(LLM):
+    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
+
+    @classmethod
+    def _load_from(cls,
+                   local_path: str,
+                   tokenizer_path: Optional[str] = None,
+                   m_cfg: ModelConfig = ModelConfig(),
+                   plugin_id: str = "llama_cpp",
+                   device_id: Optional[str] = None
+        ) -> 'PyBindLLMImpl':
+        """Load model from local path."""
+        _ensure_runtime()
+
+        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
+        if m_cfg.n_gpu_layers is not None:
+            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
+
+        # Create handle : returns py::capsule with automatic cleanup
+        handle = llm_bind.ml_llm_create(
+            model_path=local_path,
+            tokenizer_path=tokenizer_path,
+            model_config=config,
+            plugin_id=plugin_id,
+            device_id=device_id
+        )
+        return cls(handle, m_cfg)
+
+    def eject(self):
+        """Release the model from memory."""
+        # py::capsule handles cleanup automatically
+        del self._handle
+        self._handle = None
+
+    def apply_chat_template(self, messages: list[ChatMessage]) -> str:
+        """Apply the chat template to messages."""
+        # Convert TypedDict to list of dicts for binding
+        message_dicts = [
+            {"role": m["role"], "content": m["content"]}
+            for m in messages
+        ]
+        return llm_bind.ml_llm_apply_chat_template(self._handle, message_dicts)
+
+    def generate_stream(self, prompt: str, g_cfg: GenerationConfig = GenerationConfig()) -> Generator[str, None, None]:
+        """Generate text with streaming."""
+        token_queue = queue.Queue()
+        exception_container = [None]
+        self.reset_cancel()  # Reset cancel flag before generation
+
+        def on_token(token: str, user_data) -> bool:
+            if self._cancel_event.is_set():
+                token_queue.put(('end', None))
+                return False  # Stop generation
+            try:
+                token_queue.put(('token', token))
+                return True  # Continue generation
+            except Exception as e:
+                exception_container[0] = e
+                return False  # Stop generation
+
+        config = self._convert_generation_config(g_cfg)
+
+        # Run generation in thread
+        def generate():
+            try:
+                llm_bind.ml_llm_generate(
+                    handle=self._handle,
+                    prompt=prompt,
+                    config=config,
+                    on_token=on_token,
+                    user_data=None
+                )
+            except Exception as e:
+                exception_container[0] = e
+            finally:
+                token_queue.put(('end', None))
+
+        thread = threading.Thread(target=generate)
+        thread.start()
+
+        # Yield tokens as they come
+        try:
+            while True:
+                msg_type, token = token_queue.get()
+                if msg_type == 'token':
+                    yield token
+                elif msg_type in ('error', 'end'):
+                    break
+        finally:
+            thread.join()
+
+        if exception_container[0]:
+            raise exception_container[0]
+
+    def generate(self, prompt: str, g_cfg: GenerationConfig = GenerationConfig()) -> str:
+        """
+        Generate text without streaming.
+
+        Args:
+            prompt (str): The prompt to generate text from. For chat models, this is the chat messages after chat template is applied.
+            g_cfg (GenerationConfig): Generation configuration.
+
+        Returns:
+            str: The generated text.
+        """
+        config = self._convert_generation_config(g_cfg)
+        result = llm_bind.ml_llm_generate(
+            handle=self._handle,
+            prompt=prompt,
+            config=config,
+            on_token=None,  # No callback for non-streaming
+            user_data=None
+        )
+        return result.get("text", "")
+
+    def save_kv_cache(self, path: str):
+        """
+        Save the key-value cache to the file.
+
+        Args:
+            path (str): The path to the file.
+        """
+        llm_bind.ml_llm_save_kv_cache(self._handle, path)
+
+    def load_kv_cache(self, path: str):
+        """
+        Load the key-value cache from the file.
+
+        Args:
+            path (str): The path to the file.
+        """
+        llm_bind.ml_llm_load_kv_cache(self._handle, path)
+
+    def reset(self):
+        """
+        Reset the LLM model context and KV cache. If not reset, the model will skip the number of evaluated tokens and treat tokens after those as the new incremental tokens.
+        If your past chat history changed, or you are starting a new chat, you should always reset the model before running generate.
+        """
+        llm_bind.ml_llm_reset(self._handle)
+
+    def _convert_generation_config(self, g_cfg: GenerationConfig):
+        """Convert GenerationConfig to binding format."""
+        config = common_bind.GenerationConfig()
+
+        # Set basic generation parameters
+        config.max_tokens = g_cfg.max_tokens
+
+        if g_cfg.stop_words:
+            config.stop = g_cfg.stop_words
+
+        if g_cfg.image_paths:
+            config.image_paths = g_cfg.image_paths
+
+        if g_cfg.audio_paths:
+            config.audio_paths = g_cfg.audio_paths
+
+        if g_cfg.sampler_config:
+            sampler = common_bind.SamplerConfig()
+            sampler.temperature = g_cfg.sampler_config.temperature
+            sampler.top_p = g_cfg.sampler_config.top_p
+            sampler.top_k = g_cfg.sampler_config.top_k
+            sampler.repetition_penalty = g_cfg.sampler_config.repetition_penalty
+            sampler.presence_penalty = g_cfg.sampler_config.presence_penalty
+            sampler.frequency_penalty = g_cfg.sampler_config.frequency_penalty
+            sampler.seed = g_cfg.sampler_config.seed
+
+            if g_cfg.sampler_config.grammar_path:
+                sampler.grammar_path = g_cfg.sampler_config.grammar_path
+
+            if g_cfg.sampler_config.grammar_string:
+                sampler.grammar_string = g_cfg.sampler_config.grammar_string
+
+            config.sampler_config = sampler
+
+        return config

nexaai/rerank.py

@@ -0,0 +1,51 @@
+from typing import List, Optional, Sequence
+from abc import abstractmethod
+from dataclasses import dataclass
+
+from nexaai.base import BaseModel
+
+
+@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: str = "llama_cpp",
+                   device_id: Optional[str] = None
+        ) -> 'Reranker':
+        """Load reranker model from local path, routing to appropriate implementation."""
+        if plugin_id == "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,91 @@
+# Note: This code is generated by Cursor, not tested yet.
+
+from typing import List, Optional, Sequence
+import os
+
+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: str = "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,42 @@
+from typing import List, Optional, Sequence
+
+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: str = "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")