nexaai 1.0.4rc10__py3-none-macosx_11_0_arm64.whl

nexaai/mlx_backend/vlm/interface.py

@@ -0,0 +1,406 @@
+from __future__ import annotations
+
+import json
+import os
+import time
+from typing import Any, List, Optional, Sequence, Tuple, Union
+import mlx.core as mx
+import codecs
+from dataclasses import dataclass
+
+# Import configs and callback types from ml.py for API alignment
+from ml import (
+    VLM as BaseVLM,
+    SamplerConfig,
+    GenerationConfig,
+    ChatMessage,
+    EmbeddingConfig,
+    TokenCallback,
+    Path,
+    Tool,  # Add Path alias for type hints
+)
+
+# Import profiling module
+from profiling import ProfilingMixin, ProfilingData, StopReason
+
+# Import from the actual mlx_vlm structure
+from .generate import generate, stream_generate, load
+from .modeling.prompt_utils import apply_chat_template
+
+# --------------------------------------------------------------------------------------
+# Updated GenerationResult to match the new structure
+# --------------------------------------------------------------------------------------
+
+@dataclass
+class GenerationResult:
+    text: str = ""
+    token: Optional[int] = None
+    logprobs: Optional[List[float]] = None
+    prompt_tokens: int = 0
+    generation_tokens: int = 0
+    total_tokens: int = 0
+    prompt_tps: float = 0.0
+    generation_tps: float = 0.0
+    peak_memory: float = 0.0
+# --------------------------------------------------------------------------------------
+# VLM (Vision-Language Model)
+# --------------------------------------------------------------------------------------
+
+class VLM(ProfilingMixin):
+    """
+    Vision-Language Models for mlx-vlm
+    API aligned with ml.py VLM abstract base class.
+    """
+
+    def __init__(
+        self,
+        model_path: Path,
+        mmproj_path: Path,
+        context_length: int,
+        device: Optional[str] = None,
+    ) -> None:
+        # Initialize profiling mixin
+        ProfilingMixin.__init__(self)
+
+        # Check if model_path is a file, if so use its parent directory
+        if os.path.isfile(model_path):
+            model_path = os.path.dirname(model_path)
+
+        self.model_path = model_path
+        self.mmproj_path = mmproj_path
+        self.context_length = context_length
+        self.device = device
+
+        self.model, self.processor = load(str(model_path))
+
+        # Init deafutl sampler config with defualt.
+        self.sampler_config = SamplerConfig()
+
+    def destroy(self) -> None:
+        """Destroy the model and free resources."""
+        self.model = None
+        self.processor = None
+
+    def reset(self) -> None:
+        """Reset the model state."""
+        self._reset_cache()
+
+    def _reset_cache(self) -> None:
+        """Reset the KV cache."""
+        # If the model has a cache, reset it
+        if hasattr(self.model, "cache"):
+            self.model.cache = None
+
+    # Tokenization
+    def encode(self, text: str) -> List[int]:
+        """Encode text to token IDs."""
+        return self.processor.encode(text)
+
+    def decode(self, token_ids: Sequence[int]) -> str:
+        """Decode token IDs to text."""
+        return self.processor.decode(token_ids)
+
+    # Sampler
+    def set_sampler(self, config: SamplerConfig) -> None:
+        """Set sampler configuration."""
+        self.sampler_config = config
+
+    def reset_sampler(self) -> None:
+        """Reset sampler to default configuration."""
+        self.sampler_config = None
+
+    # Generation
+    def generate(
+        self,
+        prompt: str,
+        config: Optional[GenerationConfig] = None,
+    ) -> GenerationResult:
+        """Generate text from prompt."""
+        # Start profiling
+        self._start_profiling()
+
+        gen_kwargs = {}
+        if config is not None:
+            gen_kwargs = config.__dict__.copy()
+            # Remove image_paths and audio_paths from config as they'll be handled separately
+            gen_kwargs.pop('image_paths', None)
+            gen_kwargs.pop('audio_paths', None)
+        if self.sampler_config is not None:
+            gen_kwargs.update(self.sampler_config.__dict__)
+
+        # Get image and audio paths from config
+        image_paths = config.image_paths if config else None
+        audio_paths = config.audio_paths if config else None
+
+        # Convert paths to strings for generate function
+        image_list = [str(path) for path in image_paths] if image_paths else None
+        audio_list = [str(path) for path in audio_paths] if audio_paths else None
+
+        # End prompt processing, start decode
+        self._prompt_end()
+        self._decode_start()
+
+        try:
+            # Start timing for generation
+            generation_start_time = time.perf_counter()
+
+            text, stats = generate(
+                self.model,
+                self.processor,
+                prompt,
+                image=image_list,
+                audio=audio_list,
+                **gen_kwargs,
+            )
+
+            # End timing for generation
+            generation_end_time = time.perf_counter()
+
+            # Calculate average time per token and estimate TTFT
+            generated_tokens = stats.get("output_tokens", 0)
+            if generated_tokens > 0:
+                total_generation_time = generation_end_time - generation_start_time
+                avg_time_per_token = total_generation_time / generated_tokens
+                # TTFT = prompt processing time + first token generation time
+                # This provides a more accurate estimate than the previous approximation
+                estimated_ttft = (self._profiling_context.prompt_end_time - self._profiling_context.prompt_start_time) + avg_time_per_token
+                # Update the profiling context with estimated TTFT
+                self._profiling_context.first_token_time = self._profiling_context.prompt_start_time + estimated_ttft
+                self._profiling_context.ttft_recorded = True
+            else:
+                # If no tokens generated, use total generation time as TTFT
+                self._record_ttft()
+
+            # Update profiling data
+            prompt_tokens = stats.get("input_tokens", 0)
+            self._update_prompt_tokens(prompt_tokens)
+            self._update_generated_tokens(generated_tokens)
+            self._set_stop_reason(StopReason.ML_STOP_REASON_COMPLETED)
+            self._decode_end()
+            self._end_profiling()
+
+            return GenerationResult(
+                text=text,
+                prompt_tokens=prompt_tokens,
+                generation_tokens=generated_tokens,
+                total_tokens=stats.get("total_tokens", 0),
+                prompt_tps=stats.get("prompt_tps", 0.0),
+                generation_tps=stats.get("generation_tps", 0.0),
+                peak_memory=stats.get("peak_memory", 0.0),
+            )
+        except Exception as e:
+            self._set_stop_reason(StopReason.ML_STOP_REASON_UNKNOWN)
+            self._decode_end()
+            self._end_profiling()
+            raise RuntimeError(f"Generation error: {str(e)}")
+
+    def generate_stream(
+        self,
+        prompt: str,
+        config: Optional[GenerationConfig],
+        on_token: Optional[TokenCallback],
+    ) -> GenerationResult:
+        """Generate text with streaming callback. Unified method for both text and multimodal generation."""
+        # Start profiling
+        self._start_profiling()
+
+        gen_kwargs = {}
+        if config is not None:
+            gen_kwargs = config.__dict__.copy()
+            # Remove image_paths and audio_paths from config as they'll be handled separately
+            gen_kwargs.pop('image_paths', None)
+            gen_kwargs.pop('audio_paths', None)
+        if self.sampler_config is not None:
+            gen_kwargs.update(self.sampler_config.__dict__)
+
+        # Get image and audio paths from config
+        image_paths = config.image_paths if config else None
+        audio_paths = config.audio_paths if config else None
+
+        # Convert paths to strings for stream_generate function
+        image_list = [str(path) for path in image_paths] if image_paths else None
+        audio_list = [str(path) for path in audio_paths] if audio_paths else None
+
+        # End prompt processing, start decode
+        self._prompt_end()
+        self._decode_start()
+
+        text = ""
+        last_result = None
+        first_token = True
+
+        try:
+            for result in stream_generate(
+                self.model,
+                self.processor,
+                prompt,
+                image=image_list,
+                audio=audio_list,
+                **gen_kwargs,
+            ):
+                                # Record TTFT on first token
+                if first_token:
+                    self._record_ttft()
+                    first_token = False
+
+                # Call the token callback if provided
+                if on_token is not None:
+                    if not on_token(result.text):
+                        self._set_stop_reason(StopReason.ML_STOP_REASON_USER)
+                        break
+                text += result.text
+                last_result = result
+
+            # Set stop reason if not user stop
+            if self._profiling_context.stop_reason != StopReason.ML_STOP_REASON_USER:
+                self._set_stop_reason(StopReason.ML_STOP_REASON_EOS)
+
+            # Update profiling data
+            if last_result:
+                self._update_prompt_tokens(last_result.prompt_tokens)
+                self._update_generated_tokens(last_result.generation_tokens)
+
+            self._decode_end()
+            self._end_profiling()
+
+            return GenerationResult(
+                text=text,
+                token=last_result.token if last_result else None,
+                logprobs=last_result.logprobs if last_result else None,
+                prompt_tokens=last_result.prompt_tokens if last_result else 0,
+                generation_tokens=last_result.generation_tokens if last_result else 0,
+                total_tokens=(last_result.prompt_tokens + last_result.generation_tokens) if last_result else 0,
+                prompt_tps=last_result.prompt_tps if last_result else 0.0,
+                generation_tps=last_result.generation_tps if last_result else 0.0,
+                peak_memory=last_result.peak_memory if last_result else 0.0,
+            )
+        except Exception as e:
+            self._set_stop_reason(StopReason.ML_STOP_REASON_UNKNOWN)
+            self._decode_end()
+            self._end_profiling()
+            raise RuntimeError(f"Streaming generation error: {str(e)}")
+
+    # Legacy multimodal methods - kept for backward compatibility but delegate to unified method
+    def generate_multimodal(
+        self,
+        prompt: str,
+        image_paths: Optional[Sequence[Path]] = None,
+        audio_paths: Optional[Sequence[Path]] = None,
+        config: Optional[GenerationConfig] = None,
+    ) -> str:
+        """Generate text from prompt with multiple images and audio."""
+        # Create config with media paths if not provided
+        if config is None:
+            config = GenerationConfig()
+
+        # Update config with provided paths
+        if image_paths is not None:
+            config.image_paths = image_paths
+        if audio_paths is not None:
+            config.audio_paths = audio_paths
+
+        # Delegate to unified generate method and extract text
+        result = self.generate(prompt, config)
+        return result.text
+
+    def generate_stream_multimodal(
+        self,
+        prompt: str,
+        image_paths: Optional[Sequence[Path]] = None,
+        audio_paths: Optional[Sequence[Path]] = None,
+        config: Optional[GenerationConfig] = None,
+        on_token: Optional[TokenCallback] = None,
+    ) -> str:
+        """Generate text from prompt with multiple images and audio using streaming callback."""
+        # Create config with media paths if not provided
+        if config is None:
+            config = GenerationConfig()
+
+        # Update config with provided paths
+        if image_paths is not None:
+            config.image_paths = image_paths
+        if audio_paths is not None:
+            config.audio_paths = audio_paths
+
+        # Delegate to unified generate_stream method and extract text
+        result = self.generate_stream(prompt, config, on_token)
+        return result.text
+
+    def get_chat_template(self, template_name: str) -> str:
+        """Get chat template by name."""
+        # This is a stub; actual implementation depends on processor internals
+        if hasattr(self.processor, "get_chat_template"):
+            return self.processor.get_chat_template(template_name)
+        return ""
+
+    def apply_chat_template(self, messages: Sequence[ChatMessage], tools: Optional[str] = None, enable_thinking: bool = True) -> str:
+        """Apply chat template to messages with optional tools support."""
+        if hasattr(self.processor, "apply_chat_template"):
+            # Convert ChatMessage objects to dictionaries for the processor
+            messages_dict = [{"role": msg.role, "content": msg.content} for msg in messages]
+            
+            parsed_tools = None
+            if tools is not None and tools.strip():
+                parsed_tools = json.loads(tools)
+            
+            result = apply_chat_template(self.processor, self.model.config, messages_dict, add_generation_prompt=True, enable_thinking=enable_thinking, tools=parsed_tools)
+            return result
+        # Fallback: join messages
+        return "\n".join([f"{m.role}: {m.content}" for m in messages])
+
+    def apply_chat_template_with_media(self, messages: Sequence[ChatMessage], num_images: int = 0, num_audios: int = 0, tools: Optional[str] = None, enable_thinking: bool = True) -> str:
+        """Apply chat template to messages with proper image/audio token insertion and optional tools support."""
+        # Convert ChatMessage objects to dictionaries for the processor
+        messages_dict = [{"role": msg.role, "content": msg.content} for msg in messages]
+
+        parsed_tools = None
+        if tools is not None and tools.strip():
+            parsed_tools = json.loads(tools)
+
+        # Use the same logic as generate.py
+        return apply_chat_template(
+            self.processor,
+            self.model.config,
+            messages_dict,
+            num_images=num_images,
+            num_audios=num_audios,
+            enable_thinking=enable_thinking,
+            tools=parsed_tools
+        )
+
+    # Embeddings
+    def embed(
+        self,
+        texts: Sequence[str],
+        config: Optional[EmbeddingConfig] = None,
+    ) -> List[List[float]]:
+        """Generate embeddings for texts with profiling."""
+        # Start profiling
+        self._start_profiling()
+
+        try:
+            # If processor/model supports embeddings, use it; otherwise, stub
+            if hasattr(self.model, "embed"):
+                embed_kwargs = config.__dict__ if config else {}
+
+                # End prompt processing, start decode
+                self._prompt_end()
+                self._decode_start()
+
+                result = self.model.embed(texts, **embed_kwargs)
+
+                # End timing and finalize profiling data
+                self._update_generated_tokens(0)  # No generation in embedding
+                self._set_stop_reason(StopReason.ML_STOP_REASON_COMPLETED)
+                self._decode_end()
+                self._end_profiling()
+
+                return result
+            else:
+                raise NotImplementedError("Embedding not supported for this model.")
+
+        except Exception as e:
+            self._set_stop_reason(StopReason.ML_STOP_REASON_UNKNOWN)
+            self._decode_end()
+            self._end_profiling()
+            raise RuntimeError(f"Error generating embeddings: {str(e)}")

nexaai/mlx_backend/vlm/main.py

@@ -0,0 +1,157 @@
+from .interface import VLM
+from ml import GenerationConfig, SamplerConfig, ChatMessage
+import re
+import os
+import codecs
+
+def parse_media_from_input(user_input):
+    """Parse quoted media files from user input and return prompt and media paths"""
+    # Find all quoted strings (both single and double quotes)
+    quoted_pattern = r'["\']([^"\']*)["\']'
+    quoted_matches = re.findall(quoted_pattern, user_input)
+    
+    # Remove quoted strings from the input to get the actual prompt
+    prompt = re.sub(quoted_pattern, '', user_input).strip()
+    
+    # Separate image and audio files based on extensions
+    image_extensions = {'.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff', '.webp'}
+    audio_extensions = {'.mp3', '.wav', '.flac', '.aac', '.ogg', '.m4a'}
+    
+    image_paths = []
+    audio_paths = []
+    
+    for quoted_file in quoted_matches:
+        if quoted_file:  # Skip empty quotes
+            # Expand user path if it starts with ~
+            if quoted_file.startswith('~'):
+                quoted_file = os.path.expanduser(quoted_file)
+            
+            # Check if file exists
+            if not os.path.exists(quoted_file):
+                print(f"Warning: File '{quoted_file}' not found")
+                continue
+                
+            file_ext = os.path.splitext(quoted_file.lower())[1]
+            if file_ext in image_extensions:
+                image_paths.append(quoted_file)
+            elif file_ext in audio_extensions:
+                audio_paths.append(quoted_file)
+    
+    return prompt, image_paths if image_paths else None, audio_paths if audio_paths else None
+
+def test_vlm_generate_stream(model_path):
+    # Specify the checkpoint
+    context_length = 2048
+
+    # Load the corresponding model and VLM instance
+    vlm = VLM(
+        model_path=model_path,
+        mmproj_path=None,  # Not needed for this model
+        context_length=context_length,
+        device=None
+    )
+
+    # Configure sampler
+    sampler_config = SamplerConfig(
+        temperature=0.7,
+        top_p=0.9
+    )
+    vlm.set_sampler(sampler_config)
+
+    # Chat history using ChatMessage objects (following ml.py API)
+    chat = []
+
+    print("Multi-round VLM conversation started. Type 'quit' or 'exit' to end.")
+    print("Include images/audios in quotes, e.g.: 'describe \"image1.jpg\" \"image2.png\"'")
+    print("You can also use single quotes: 'describe '/path/to/image.jpg''")
+    print("=" * 50)
+
+    def on_token(text_chunk, user_data):
+        """Token callback for streaming"""
+        print(text_chunk, end="", flush=True)
+        if user_data is not None:
+            user_data["response"] += text_chunk
+        return True
+
+    while True:
+        # Get user input
+        user_input = input("\nUser: ").strip()
+
+        # Check for exit commands
+        if user_input.lower() in ["quit", "exit", "q"]:
+            print("Goodbye!")
+            break
+
+        if not user_input:
+            continue
+
+        # Parse media files and prompt from user input
+        prompt_text, image_paths, audio_paths = parse_media_from_input(user_input)
+        
+        print(f"image_paths: {image_paths}")
+        print(f"audio_paths: {audio_paths}")
+        
+        # If no text prompt after parsing, use the original input
+        if not prompt_text.strip():
+            prompt_text = user_input
+            image_paths = None
+            audio_paths = None
+
+        # Add user message to chat history using ChatMessage (following ml.py API)
+        chat.append(ChatMessage(role="user", content=prompt_text))
+
+        # Calculate number of images and audios for chat template
+        num_images = len(image_paths) if image_paths else 0
+        num_audios = len(audio_paths) if audio_paths else 0
+
+        # Apply chat template with image/audio token insertion
+        try:
+            formatted_prompt = vlm.apply_chat_template_with_media(chat, num_images=num_images, num_audios=num_audios)
+        except (NotImplementedError, AttributeError):
+            # Fallback to manual formatting if chat template is not implemented
+            formatted_prompt = ""
+            for msg in chat:
+                formatted_prompt += f"{msg.role}: {msg.content}\n"
+            formatted_prompt += "Assistant: "
+
+        # Generation config with media paths
+        generation_config = GenerationConfig(
+            max_tokens=512,
+            sampler_config=sampler_config,
+            image_paths=image_paths,
+            audio_paths=audio_paths
+        )
+
+        # Generate response
+        print("Assistant: ", end="", flush=True)
+        
+        try:
+            # Use streaming generation with callback - single method handles all cases
+            user_data = {"response": ""}
+            
+            # Always use the unified generate_stream method
+            response = vlm.generate_stream(
+                prompt=formatted_prompt,
+                config=generation_config,
+                on_token=on_token,
+                user_data=user_data
+            )
+            
+            print()  # New line after streaming
+            
+            # Add assistant response to chat history using ChatMessage
+            chat.append(ChatMessage(role="assistant", content=user_data["response"]))
+            
+        except Exception as e:
+            print(f"Error generating response: {e}")
+            print()
+
+    # Clean up
+    vlm.destroy()
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", type=str, default="mlx-community/gemma-3-4b-it-8bit")
+    args = parser.parse_args()
+    test_vlm_generate_stream(args.model_path)

nexaai/mlx_backend/vlm/modeling/convert.py

@@ -0,0 +1,68 @@
+# Copyright © 2023-2024 Apple Inc.
+
+import argparse
+
+from .utils import MODEL_CONVERSION_DTYPES, convert
+
+
+def configure_parser() -> argparse.ArgumentParser:
+    """
+    Configures and returns the argument parser for the script.
+
+    Returns:
+        argparse.ArgumentParser: Configured argument parser.
+    """
+    parser = argparse.ArgumentParser(
+        description="Convert Hugging Face model to MLX format"
+    )
+
+    parser.add_argument("--hf-path", type=str, help="Path to the Hugging Face model.")
+    parser.add_argument(
+        "--mlx-path", type=str, default="mlx_model", help="Path to save the MLX model."
+    )
+    parser.add_argument(
+        "-q", "--quantize", help="Generate a quantized model.", action="store_true"
+    )
+    parser.add_argument(
+        "--q-group-size", help="Group size for quantization.", type=int, default=64
+    )
+    parser.add_argument(
+        "--q-bits", help="Bits per weight for quantization.", type=int, default=4
+    )
+    parser.add_argument(
+        "--dtype",
+        help="Type to save the parameter. Defaults to config.json's `torch_dtype` or the current model weights dtype",
+        type=str,
+        choices=MODEL_CONVERSION_DTYPES,
+        default=None,
+    )
+    parser.add_argument(
+        "--upload-repo",
+        help="The Hugging Face repo to upload the model to.",
+        type=str,
+        default=None,
+    )
+    parser.add_argument(
+        "-d",
+        "--dequantize",
+        help="Dequantize a quantized model.",
+        action="store_true",
+        default=False,
+    )
+    parser.add_argument(
+        "--skip-vision",
+        help="Skip vision module quantization.",
+        action="store_true",
+        default=False,
+    )
+    return parser
+
+
+def main():
+    parser = configure_parser()
+    args = parser.parse_args()
+    convert(**vars(args))
+
+
+if __name__ == "__main__":
+    main()

nexaai/mlx_backend/vlm/modeling/models/aya_vision/__init__.py

@@ -0,0 +1,8 @@
+from .aya_vision import (
+    LanguageModel,
+    Model,
+    ModelConfig,
+    TextConfig,
+    VisionConfig,
+    VisionModel,
+)