@modular-prompt/driver 0.12.0 → 0.13.1

package/src/mlx-ml/python/backends/mlx_lm.py

@@ -0,0 +1,202 @@
+from __future__ import annotations
+
+import json
+import os
+import sys
+from typing import Any, Iterator
+
+from mlx_lm import load as mlx_lm_load
+from mlx_lm import stream_generate as mlx_lm_stream_generate
+from mlx_lm.models.cache import make_prompt_cache, save_prompt_cache, load_prompt_cache, trim_prompt_cache
+from mlx_lm.sample_utils import make_sampler
+
+from backends.base import ModelBackend
+from utils.token_utils import is_eod_token
+
+
+class MlxLmBackend(ModelBackend):
+    """`mlx_lm` backend for text-only models."""
+
+    def __init__(self) -> None:
+        self.model: Any | None = None
+        self.tokenizer: Any | None = None
+
+    def load(self, model_name: str) -> None:
+        self.model, self.tokenizer = mlx_lm_load(model_name)
+
+    def get_tokenizer(self) -> Any:
+        return self.tokenizer
+
+    def stream_generate(
+        self,
+        prompt: str | list[int],
+        options: dict,
+        images: list | None = None,
+        prompt_cache: list | None = None,
+    ) -> Iterator[Any]:
+        if self.model is None or self.tokenizer is None:
+            raise RuntimeError("Model is not loaded")
+
+        final_options = {"max_tokens": 1000, **options}
+        temperature = final_options.pop("temperature", 1.0)
+        top_p = final_options.pop("top_p", 0.0)
+        top_k = final_options.pop("top_k", 0)
+        final_options["sampler"] = make_sampler(
+            temp=temperature,
+            top_p=top_p,
+            top_k=top_k,
+        )
+
+        if prompt_cache is not None:
+            final_options["prompt_cache"] = prompt_cache
+
+        for response in mlx_lm_stream_generate(
+            self.model,
+            self.tokenizer,
+            prompt,
+            **final_options,
+        ):
+            if is_eod_token(response, self.tokenizer):
+                break
+            yield response
+
+    # get_cache_offset is inherited from ModelBackend base class
+
+    def _tokenize_prompt(self, prompt: str) -> list[int]:
+        """Tokenize a prompt string using the same logic as stream_generate."""
+        add_special = self.tokenizer.bos_token is None or not prompt.startswith(
+            self.tokenizer.bos_token
+        )
+        return self.tokenizer.encode(prompt, add_special_tokens=add_special)
+
+    @staticmethod
+    def _write_cache_meta(
+        cache_path: str,
+        token_count: int,
+        prefix_offsets: list[int] | None = None,
+        prefix_hashes: list[str] | None = None,
+    ) -> None:
+        meta_path = cache_path + '.meta.json'
+        try:
+            meta: dict[str, Any] = {"token_count": token_count}
+            if prefix_offsets and prefix_hashes:
+                meta["prefix_offsets"] = prefix_offsets
+                meta["prefix_hashes"] = prefix_hashes
+            with open(meta_path, 'w') as f:
+                json.dump(meta, f)
+        except Exception as e:
+            sys.stderr.write(f"Failed to write cache meta: {e}\n")
+
+    @staticmethod
+    def _read_cache_meta(cache_path: str) -> int | None:
+        meta_path = cache_path + '.meta.json'
+        try:
+            with open(meta_path) as f:
+                meta = json.load(f)
+                count = meta.get('token_count')
+                return int(count) if count is not None else None
+        except (FileNotFoundError, json.JSONDecodeError, ValueError, TypeError):
+            return None
+
+    def cache_prefill(
+        self,
+        cache_path: str,
+        prompt: str,
+        base_cache_path: str | None = None,
+        trim_to_tokens: int | None = None,
+        prefix_offsets: list[int] | None = None,
+        prefix_hashes: list[str] | None = None,
+    ) -> dict:
+        if self.model is None or self.tokenizer is None:
+            raise RuntimeError("Model is not loaded")
+
+        full_tokens = self._tokenize_prompt(prompt)
+        token_count = len(full_tokens)
+        effective_prompt: str | list[int] = prompt
+
+        if base_cache_path is not None:
+            try:
+                prompt_cache = load_prompt_cache(base_cache_path)
+
+                if trim_to_tokens is not None:
+                    current_offset = self.get_cache_offset(prompt_cache)
+                    if current_offset > trim_to_tokens:
+                        trim_count = current_offset - trim_to_tokens
+                        trim_prompt_cache(prompt_cache, trim_count)
+                        if os.getenv('MLX_DEBUG'):
+                            sys.stderr.write(
+                                f"Trimmed base cache: {current_offset} → {trim_to_tokens} tokens\n"
+                            )
+                        cache_offset = trim_to_tokens
+                    else:
+                        cache_offset = current_offset
+                else:
+                    meta_offset = self._read_cache_meta(base_cache_path)
+                    if meta_offset is None:
+                        # Legacy cache without meta file - create fresh cache
+                        sys.stderr.write(
+                            f"WARNING: Cache file exists but no .meta.json found at {base_cache_path}. "
+                            "Creating fresh cache (may be from old implementation).\n"
+                        )
+                        cache_offset = 0
+                        prompt_cache = make_prompt_cache(self.model)
+                    else:
+                        cache_offset = meta_offset
+
+                if cache_offset > 0 and prompt_cache is not None:
+                    if cache_offset < token_count:
+                        effective_prompt = full_tokens[cache_offset:]
+                        if os.getenv('MLX_DEBUG'):
+                            sys.stderr.write(
+                                f"Incremental prefill from: {base_cache_path} "
+                                f"(skip {cache_offset}/{token_count} tokens)\n"
+                            )
+                    else:
+                        if os.getenv('MLX_DEBUG'):
+                            sys.stderr.write(
+                                f"Base cache covers entire prompt "
+                                f"({cache_offset} >= {token_count}), saving as-is\n"
+                            )
+                        save_prompt_cache(cache_path, prompt_cache)
+                        self._write_cache_meta(cache_path, token_count, prefix_offsets, prefix_hashes)
+                        return {"cache_path": cache_path, "token_count": token_count}
+                else:
+                    if os.getenv('MLX_DEBUG'):
+                        sys.stderr.write(f"Incremental prefill from: {base_cache_path}\n")
+            except Exception as e:
+                sys.stderr.write(f"Base cache load failed, creating fresh: {e}\n")
+                prompt_cache = make_prompt_cache(self.model)
+        else:
+            prompt_cache = make_prompt_cache(self.model)
+
+        if os.getenv('MLX_DEBUG'):
+            sys.stderr.write(f"Prefill prompt: {token_count} tokens\n")
+
+        for _ in mlx_lm_stream_generate(
+            self.model, self.tokenizer, effective_prompt,
+            prompt_cache=prompt_cache, max_tokens=1,
+        ):
+            break
+
+        save_prompt_cache(cache_path, prompt_cache)
+        self._write_cache_meta(cache_path, token_count, prefix_offsets, prefix_hashes)
+        if os.getenv('MLX_DEBUG'):
+            sys.stderr.write(f"Cache created: {cache_path} ({token_count} tokens)\n")
+        return {"cache_path": cache_path, "token_count": token_count}
+
+    def load_cache_from_file(self, cache_path: str) -> list | None:
+        try:
+            return load_prompt_cache(cache_path)
+        except FileNotFoundError:
+            sys.stderr.write(f"Cache file not found: {cache_path}\n")
+            return None
+        except Exception as e:
+            sys.stderr.write(f"Failed to load cache: {e}\n")
+            return None
+
+    def supports_vision(self) -> bool:
+        return False
+
+    @property
+    def model_kind(self) -> str:
+        return "lm"

package/src/mlx-ml/python/backends/mlx_vlm.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+import sys
+from typing import Any, Iterator
+
+from mlx_vlm import load as mlx_vlm_load
+from mlx_vlm import stream_generate as mlx_vlm_stream_generate
+
+from backends.base import ModelBackend
+from utils.vlm_utils import load_and_resize_images
+
+
+class MlxVlmBackend(ModelBackend):
+    """`mlx_vlm` backend for vision-language models."""
+
+    def __init__(self) -> None:
+        self.model: Any | None = None
+        self.processor: Any | None = None
+        self.drafter: Any | None = None
+        self.drafter_kind: str | None = None
+        self.draft_block_size: int | None = None
+
+    def load(self, model_name: str) -> None:
+        self.model, self.processor = mlx_vlm_load(model_name)
+
+    def load_drafter(self, drafter_model: str) -> None:
+        from mlx_vlm.speculative.drafters import load_drafter
+        self.drafter, self.drafter_kind = load_drafter(drafter_model)
+        sys.stderr.write(f"Drafter loaded: {drafter_model} (kind={self.drafter_kind})\n")
+
+    def has_drafter(self) -> bool:
+        return self.drafter is not None
+
+    def get_tokenizer(self) -> Any:
+        return self.processor
+
+    def stream_generate(
+        self, prompt: str | list[int], options: dict, images: list | None = None,
+        prompt_cache: list | None = None,
+    ) -> Iterator[Any]:
+        if self.model is None or self.processor is None:
+            raise RuntimeError("Model is not loaded")
+
+        final_options = dict(options)
+        temperature = final_options.pop("temperature", 1.0)
+        max_tokens = final_options.pop("max_tokens", 1000)
+        top_p = final_options.pop("top_p", 0.0)
+        top_k = final_options.pop("top_k", 0)
+
+        processed_images = None
+        if images:
+            max_image_size = final_options.pop("max_image_size", 768)
+            processed_images = load_and_resize_images(images, max_image_size)
+
+        draft_kwargs = {}
+        if self.drafter:
+            draft_kwargs["draft_model"] = self.drafter
+            draft_kwargs["draft_kind"] = self.drafter_kind
+            if self.draft_block_size is not None:
+                draft_kwargs["draft_block_size"] = self.draft_block_size
+
+        # Generate and collect tokens
+        token_count = 0
+        for result in mlx_vlm_stream_generate(
+            self.model,
+            self.processor,
+            prompt,
+            image=processed_images,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            **draft_kwargs,
+        ):
+            token_count += 1
+            yield result
+
+        # Output speculative decoding stats if drafter was used
+        if self.drafter and hasattr(self.drafter, 'accept_lens'):
+            accept_lens = self.drafter.accept_lens
+            if accept_lens:
+                avg_accepted = sum(accept_lens) / len(accept_lens)
+                # Show stats unless MLX_NO_STATS environment variable is set
+                import os
+                if not os.getenv('MLX_NO_STATS'):
+                    sys.stderr.write(f"\n[Speculative Decoding Stats]\n")
+                    sys.stderr.write(f"  Rounds: {len(accept_lens)}\n")
+                    sys.stderr.write(f"  Average accepted tokens/round: {avg_accepted:.2f}\n")
+                    sys.stderr.write(f"  Total tokens generated: {token_count}\n")
+                    sys.stderr.write(f"  Speedup factor: {avg_accepted:.2f}x (theoretical)\n")
+                # Clear for next generation
+                self.drafter.accept_lens = []
+
+    def supports_vision(self) -> bool:
+        return True
+
+    @property
+    def model_kind(self) -> str:
+        return "vlm"

package/src/mlx-ml/python/handlers/__init__.py

@@ -0,0 +1,6 @@
+from handlers.cache import handle_cache_prefill
+from handlers.capabilities import handle_capabilities
+from handlers.chat import handle_chat
+from handlers.completion import handle_completion
+from handlers.format_test import handle_format_test
+from handlers.tokenize import handle_tokenize

package/src/mlx-ml/python/handlers/cache.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+import json
+import sys
+
+from backends.base import ModelBackend
+from utils.prompt_builder import generate_merged_prompt, supports_chat_template
+
+
+def handle_cache_prefill(
+    backend: ModelBackend,
+    capabilities: dict,
+    cache_path: str,
+    messages: list,
+    base_cache_path: str | None = None,
+    trim_to_tokens: int | None = None,
+    prefix_offsets: list[int] | None = None,
+    prefix_hashes: list[str] | None = None,
+    tools: list | None = None,
+    reasoning_effort: str | None = None,
+) -> None:
+    tokenizer = backend.get_tokenizer()
+
+    extra_kwargs = {}
+    if tools is not None:
+        extra_kwargs["tools"] = tools
+    if reasoning_effort is not None:
+        extra_kwargs["reasoning_effort"] = reasoning_effort
+    if supports_chat_template(tokenizer):
+        try:
+            prompt = tokenizer.apply_chat_template(
+                messages,
+                add_generation_prompt=False,
+                tokenize=False,
+                **extra_kwargs,
+            )
+        except TypeError:
+            try:
+                fallback_kwargs = {}
+                if tools is not None:
+                    fallback_kwargs["tools"] = tools
+                prompt = tokenizer.apply_chat_template(
+                    messages,
+                    add_generation_prompt=False,
+                    tokenize=False,
+                    **fallback_kwargs,
+                )
+            except TypeError:
+                try:
+                    prompt = tokenizer.apply_chat_template(
+                        messages,
+                        add_generation_prompt=False,
+                        tokenize=False,
+                    )
+                except Exception:
+                    prompt = generate_merged_prompt(messages, capabilities)
+                    sys.stderr.write(
+                        "--- cache_prefill: fallback to generate_merged_prompt\n"
+                    )
+        except Exception:
+            prompt = generate_merged_prompt(messages, capabilities)
+            sys.stderr.write(
+                "--- cache_prefill: fallback to generate_merged_prompt\n"
+            )
+    else:
+        prompt = generate_merged_prompt(messages, capabilities)
+
+    # Only show debug output if MLX_DEBUG environment variable is set
+    import os
+    if os.getenv('MLX_DEBUG'):
+        sys.stderr.write(f"--- cache_prefill {cache_path}\n")
+    result = backend.cache_prefill(
+        cache_path, prompt, base_cache_path,
+        trim_to_tokens=trim_to_tokens,
+        prefix_offsets=prefix_offsets,
+        prefix_hashes=prefix_hashes,
+    )
+    if prefix_offsets and prefix_hashes:
+        result["prefix_offsets"] = prefix_offsets
+        result["prefix_hashes"] = prefix_hashes
+    print(json.dumps(result), end="\0", flush=True)

package/src/mlx-ml/python/handlers/capabilities.py

@@ -0,0 +1,6 @@
+import json
+
+
+def handle_capabilities(capabilities: dict) -> None:
+    """capabilities API の処理。JSON出力してnull文字で終端"""
+    print(json.dumps(capabilities), end="\0", flush=True)

package/src/mlx-ml/python/handlers/chat.py

@@ -0,0 +1,221 @@
+from __future__ import annotations
+
+import json
+import re
+import sys
+
+from backends.base import ModelBackend
+from mlx_lm.models.cache import trim_prompt_cache
+from utils.prompt_builder import generate_merged_prompt, supports_chat_template
+
+
+def _read_cache_token_count(cache_path: str) -> int | None:
+    """Read token count from the sidecar .meta.json file."""
+    meta_path = cache_path + '.meta.json'
+    try:
+        with open(meta_path) as f:
+            meta = json.load(f)
+            count = meta.get('token_count')
+            return int(count) if count is not None else None
+    except (FileNotFoundError, json.JSONDecodeError, ValueError, TypeError):
+        return None
+
+
+def _stream_to_stdout(
+    backend: ModelBackend,
+    prompt: str | list[int],
+    options: dict,
+    images: list | None = None,
+    primer: str | None = None,
+    prompt_cache: list | None = None,
+) -> None:
+    if primer is not None:
+        print(primer, end="", flush=True)
+
+    last_response = None
+    for response in backend.stream_generate(prompt, options, images, prompt_cache=prompt_cache):
+        print(response.text.replace("\0", "").replace("\x1e", ""), end="", flush=True)
+        last_response = response
+
+    meta: dict = {}
+    if last_response is not None:
+        if hasattr(last_response, "prompt_tokens"):
+            meta["prompt_tokens"] = last_response.prompt_tokens
+        if hasattr(last_response, "generation_tokens"):
+            meta["generation_tokens"] = last_response.generation_tokens
+
+    if meta:
+        print(f"\x1e__META__:{json.dumps(meta)}", end="\0", flush=True)
+    else:
+        print("", end="\0", flush=True)
+
+
+def handle_chat(
+    backend: ModelBackend,
+    capabilities: dict,
+    messages: list,
+    primer: str | None = None,
+    options: dict | None = None,
+    tools: list | None = None,
+    images: list | None = None,
+    max_image_size: int = 768,
+    reasoning_effort: str | None = None,
+    cache_path: str | None = None,
+    cache_trim_tokens: int | None = None,
+) -> None:
+    """chat API の処理"""
+    if options is None:
+        options = {}
+
+    tokenizer = backend.get_tokenizer()
+
+    if backend.supports_vision():
+        add_generation_prompt = True
+        fmt_messages = list(messages)
+        if primer is not None:
+            fmt_messages.append({"role": "assistant", "content": primer})
+            add_generation_prompt = False
+
+        try:
+            prompt = tokenizer.apply_chat_template(
+                fmt_messages,
+                tools=tools,
+                add_generation_prompt=add_generation_prompt,
+                tokenize=False,
+            )
+        except TypeError:
+            prompt = tokenizer.apply_chat_template(
+                fmt_messages,
+                add_generation_prompt=add_generation_prompt,
+                tokenize=False,
+            )
+
+        if primer is not None:
+            prompt = primer.join(prompt.split(primer)[0:-1]) + primer
+
+        display_prompt = re.sub(r'(<\|image_pad\|>)+', '<|image_pad|>...', prompt)
+        sys.stderr.write(f"--- vlm prompt (images: {len(images) if images else 0}, max_size: {max_image_size})\n{display_prompt}\n")
+
+        final_options = dict(options)
+        final_options["max_image_size"] = max_image_size
+        _stream_to_stdout(
+            backend,
+            prompt,
+            final_options,
+            images=images,
+            primer=primer,
+        )
+        return
+
+    prompt_cache = backend.load_cache_from_file(cache_path) if cache_path else None
+    cache_tokens = 0
+    if prompt_cache is not None:
+        if cache_trim_tokens is not None:
+            current_offset = backend.get_cache_offset(prompt_cache)
+            if current_offset > cache_trim_tokens:
+                trim_prompt_cache(prompt_cache, current_offset - cache_trim_tokens)
+                sys.stderr.write(
+                    f"KV cache trimmed: {current_offset} → {cache_trim_tokens} tokens\n"
+                )
+                cache_tokens = cache_trim_tokens
+            else:
+                cache_tokens = current_offset
+        else:
+            meta_count = _read_cache_token_count(cache_path) if cache_path else None
+            if meta_count is not None:
+                cache_tokens = meta_count
+            else:
+                # Legacy cache without meta file - skip it for safety
+                sys.stderr.write(
+                    f"WARNING: Cache file exists but no .meta.json found at {cache_path}. "
+                    "Ignoring cache for safety (may be from old implementation).\n"
+                )
+                prompt_cache = None
+                cache_tokens = 0
+        if prompt_cache is not None:
+            sys.stderr.write(
+                f"KV cache loaded: {len(prompt_cache)} layers, {cache_tokens} cached tokens\n"
+            )
+    elif cache_path:
+        sys.stderr.write(f"KV cache load FAILED: {cache_path}\n")
+
+    if not supports_chat_template(tokenizer):
+        prompt = generate_merged_prompt(messages, capabilities)
+        if prompt_cache is not None:
+            sys.stderr.write("KV cache ignored: model does not support chat template\n")
+        _stream_to_stdout(backend, prompt, options, primer=primer)
+        return
+
+    add_generation_prompt = True
+    fmt_messages = list(messages)
+    if primer is not None:
+        fmt_messages.append({"role": "assistant", "content": primer})
+        add_generation_prompt = False
+
+    extra_kwargs = {}
+    if tools is not None:
+        extra_kwargs["tools"] = tools
+    if reasoning_effort is not None:
+        extra_kwargs["reasoning_effort"] = reasoning_effort
+
+    trust_remote_code = options.get("trust_remote_code")
+    if trust_remote_code is not None:
+        extra_kwargs["trust_remote_code"] = trust_remote_code
+
+    try:
+        prompt = tokenizer.apply_chat_template(
+            fmt_messages,
+            add_generation_prompt=add_generation_prompt,
+            tokenize=False,
+            **extra_kwargs,
+        )
+    except TypeError:
+        try:
+            fallback_kwargs = {}
+            if tools is not None:
+                fallback_kwargs["tools"] = tools
+            prompt = tokenizer.apply_chat_template(
+                fmt_messages,
+                add_generation_prompt=add_generation_prompt,
+                tokenize=False,
+                **fallback_kwargs,
+            )
+        except TypeError:
+            prompt = tokenizer.apply_chat_template(
+                fmt_messages,
+                add_generation_prompt=add_generation_prompt,
+                tokenize=False,
+            )
+
+    if primer is not None:
+        prompt = primer.join(prompt.split(primer)[0:-1]) + primer
+
+    if isinstance(prompt, list):
+        sys.stderr.write(f"--- prompt: len={len(prompt)}\n")
+    else:
+        sys.stderr.write(f"--- prompt\n{prompt}\n")
+
+    final_options = dict(options)
+    final_options.pop("trust_remote_code", None)
+
+    effective_prompt = prompt
+    if prompt_cache is not None and cache_tokens > 0 and isinstance(prompt, str):
+        add_special = tokenizer.bos_token is None or not prompt.startswith(
+            tokenizer.bos_token
+        )
+        full_tokens = tokenizer.encode(prompt, add_special_tokens=add_special)
+
+        if cache_tokens < len(full_tokens):
+            effective_prompt = full_tokens[cache_tokens:]
+            sys.stderr.write(
+                f"Prefilled {cache_tokens}/{len(full_tokens)} tokens, "
+                f"generating from {len(effective_prompt)} remaining\n"
+            )
+        else:
+            sys.stderr.write(
+                f"Prefill offset {cache_tokens} >= prompt {len(full_tokens)}, "
+                f"ignoring prefill state\n"
+            )
+            prompt_cache = None
+
+    _stream_to_stdout(backend, effective_prompt, final_options, primer=primer, prompt_cache=prompt_cache)

package/src/mlx-ml/python/handlers/completion.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import os
+import re
+import sys
+
+from backends.base import ModelBackend
+
+
+def handle_completion(
+    backend: ModelBackend,
+    prompt: str | list[int],
+    options: dict | None = None,
+    images: list | None = None,
+    max_image_size: int = 768,
+) -> None:
+    """completion API の処理"""
+    if options is None:
+        options = {}
+
+    final_options = dict(options)
+    if images:
+        final_options["max_image_size"] = max_image_size
+        if os.getenv('MLX_DEBUG'):
+            display_prompt = re.sub(r'(<\|image_pad\|>)+', '<|image_pad|>...', prompt)
+            sys.stderr.write(f"--- vlm completion (images: {len(images)}, max_size: {max_image_size})\n{display_prompt}\n")
+    elif os.getenv('MLX_DEBUG'):
+        if isinstance(prompt, list):
+            sys.stderr.write(f"--- prompt: len={len(prompt)}\n")
+        else:
+            sys.stderr.write(f"--- prompt\n{prompt}\n")
+
+    for response in backend.stream_generate(prompt, final_options, images):
+        print(response.text.replace("\0", ""), end="", flush=True)
+
+    print("\n", end="\0", flush=True)