@modular-prompt/driver 0.11.15 → 0.13.1

package/src/mlx-ml/python/__main__.py

@@ -1,442 +1,58 @@
 import sys
-import json
-from vlm_utils import detect_model_kind, load_and_resize_images
-from token_utils import get_capabilities, is_eod_token
+
+from backends import MlxLmBackend, MlxVlmBackend
+from utils.token_utils import get_capabilities
+from utils.vlm_utils import detect_model_kind
+from server import Server
 
 model_name = sys.argv[1] if len(sys.argv) > 1 else "mlx-community/gemma-3-270m-it-qat-4bit"
 text_only = "--text-only" in sys.argv
 
-# モデル種別の判定とロード
-model_kind = "lm" if text_only else detect_model_kind(model_name)
-
-if model_kind == "vlm":
-    from mlx_vlm import load as vlm_load, stream_generate as vlm_stream_generate
-    try:
-        model, processor = vlm_load(model_name)
-        tokenizer = processor  # capabilities取得用（VLMのprocessorもtokenizer互換）
-    except (ValueError, Exception) as e:
-        # mlx_vlm.models にモジュールが存在しても、実際のモデルに vision コンポーネントが
-        # ない場合（例: Qwen3.5 テキストモデルが qwen2_vl として認識される）にフォールバック
-        sys.stderr.write(f"VLM load failed, falling back to LM: {e}\n")
-        model_kind = "lm"
-        from mlx_lm import load, stream_generate
-        from mlx_lm.sample_utils import make_sampler
-        model, tokenizer = load(model_name)
-else:
-    from mlx_lm import load, stream_generate
-    from mlx_lm.sample_utils import make_sampler
-    model, tokenizer = load(model_name)
-
-# Capabilities情報の取得
-capabilities = get_capabilities(tokenizer)
-capabilities["model_kind"] = model_kind
+drafter_model = None
+if "--drafter" in sys.argv:
+    idx = sys.argv.index("--drafter")
+    if idx + 1 < len(sys.argv):
+        drafter_model = sys.argv[idx + 1]
 
-def read():
-    lines = []
-    data = None
-    eof = False
-    while not eof:
-        line = sys.stdin.readline()
-        # sys.stderr.write('line:' + line + '\n')
-        if not line:
-            eof = True
-        else:
-            lines.append(line)
+draft_block_size = None
+if "--draft-block-size" in sys.argv:
+    idx = sys.argv.index("--draft-block-size")
+    if idx + 1 < len(sys.argv):
         try:
-            data = json.loads(''.join(lines))
-        except json.JSONDecodeError as e:
-            data = None
-            continue
-        break
-    return data
-
-
-def supports_chat_template():
-    """
-    チャットテンプレートがサポートされているかを判定
-
-    apply_chat_templateメソッドの存在と、tokenizer.chat_templateの両方を確認する。
-    tokenizer.chat_templateが設定されていない場合、apply_chat_templateを呼んでも
-    エラーになるため、両方の条件をチェックする必要がある。
-
-    Returns:
-        bool: チャットテンプレートがサポートされている場合True
-    """
-    return (hasattr(tokenizer, 'apply_chat_template') and
-            hasattr(tokenizer, 'chat_template') and
-            tokenizer.chat_template is not None)
-
-
-def handle_capabilities():
-    """capabilities API の処理"""
-    print(json.dumps(capabilities), end='\0', flush=True)
-
-
-def handle_format_test(messages, options=None, tools=None):
-    """フォーマットテスト API の処理（実際に生成せずフォーマットのみ）"""
-    if options is None:
-        options = {}
-
-    result = {
-        "formatted_prompt": None,
-        "template_applied": False,
-        "model_specific_processing": None,
-        "error": None
-    }
-
-    try:
-        # チャットテンプレートが利用可能かチェック
-        if supports_chat_template():
-            # messagesはTypeScript側で既にモデル固有処理済み
-            result["model_specific_processing"] = messages
-
-            # プロンプト生成（フォーマットのみ）
-            primer = options.get('primer')
-            add_generation_prompt = True
-            tokenize = False  # 常にテキストで返す
-
-            if primer is not None:
-                messages.append({'role': 'assistant', 'content': primer})
-                add_generation_prompt = False
-
-            # tools対応を試みる（テンプレートが対応していなければtools無しで実行）
-            try:
-                formatted_prompt = tokenizer.apply_chat_template(
-                    messages,
-                    tools=tools,
-                    add_generation_prompt=add_generation_prompt,
-                    tokenize=tokenize,
-                )
-            except TypeError:
-                formatted_prompt = tokenizer.apply_chat_template(
-                    messages,
-                    add_generation_prompt=add_generation_prompt,
-                    tokenize=tokenize,
-                )
-
-            if primer is not None:
-                formatted_prompt = primer.join(formatted_prompt.split(primer)[0:-1]) + primer
-
-            result["formatted_prompt"] = formatted_prompt
-            result["template_applied"] = True
-        else:
-            # チャットテンプレートがない場合はcompletionフォーマット
-            formatted_prompt = generate_merged_prompt(messages)
-            primer = options.get('primer')
-            if primer is not None:
-                formatted_prompt += primer
-
-            result["formatted_prompt"] = formatted_prompt
-            result["template_applied"] = False
-
-    except Exception as e:
-        result["error"] = str(e)
-
-    print(json.dumps(result), end='\0', flush=True)
-
-def handle_chat(messages, primer=None, options=None, tools=None):
-    """chat API の処理"""
-    if options is None:
-        options = {}
-
-    # チャットテンプレートが利用可能かチェック
-    if not supports_chat_template():
-        # チャットテンプレートがない場合はcompletionフォーマットに変換
-        prompt = generate_merged_prompt(messages)
-        if primer is not None:
-            print(primer, end='', flush=True)
-        generate_text(prompt, options)
-        return
-
-    # プロンプト生成
-    add_generation_prompt = True
-    tokenize = False
-
-    if primer is not None:
-        messages.append({'role': 'assistant', 'content': primer})
-        add_generation_prompt = False
-        tokenize = False
-
-    # tools対応を試みる（テンプレートが対応していなければtools無しで実行）
-    try:
-        prompt = tokenizer.apply_chat_template(
-            messages,
-            tools=tools,
-            add_generation_prompt=add_generation_prompt,
-            tokenize=tokenize,
-        )
-    except TypeError:
-        prompt = tokenizer.apply_chat_template(
-            messages,
-            add_generation_prompt=add_generation_prompt,
-            tokenize=tokenize,
-        )
-
-    if primer is not None:
-        prompt = primer.join(prompt.split(primer)[0:-1]) + primer
-        print(primer, end='', flush=True)
-
-    generate_text(prompt, options)
-
-
-def generate_merged_prompt(messages):
-    """apply_chat_templateがない場合のプロンプト生成"""
-    # messagesはTypeScript側で既にmergeSystemMessages処理済み
-    # TypeScript側のformatterと同じフォーマットを維持
-
-    prompt_parts = []
-    special_tokens = capabilities.get('special_tokens', {})
-
-    for msg in messages:
-        role = msg['role']  # 小文字のまま
-        role_upper = role.upper()
-
-        # 1. 専用のspecial_tokenを探す
-        role_token = special_tokens.get(role)
+            draft_block_size = int(sys.argv[idx + 1])
+        except ValueError:
+            sys.stderr.write(f"Invalid --draft-block-size value: {sys.argv[idx + 1]}\n")
+            sys.exit(1)
 
-        if role_token and isinstance(role_token, dict) and 'start' in role_token:
-            # 専用トークンがある場合
-            start_token = role_token['start']['text']
-            end_token = role_token['end']['text']
-            prompt_parts.extend([
-                start_token,
-                msg['content'].strip(),
-                end_token,
-                ''  # 空行で区切る
-            ])
-        else:
-            # 2. 専用トークンがない場合、汎用blockトークンを探す
-            # blockやcontextなどの汎用的なペアトークンを探す
-            block_token = None
-            for candidate in ['block', 'context', 'quote', 'section']:
-                token = special_tokens.get(candidate)
-                if token and isinstance(token, dict) and 'start' in token:
-                    block_token = token
-                    break
 
-            if block_token:
-                # 汎用blockトークンがある場合: {block_begin}{role}:\n...{block_end}
-                start_token = block_token['start']['text']
-                end_token = block_token['end']['text']
-                prompt_parts.extend([
-                    f'{start_token}{role_upper}:\n{msg["content"].strip()}',
-                    end_token,
-                    ''  # 空行で区切る
-                ])
-            else:
-                # 3. どちらもない場合は、HTMLコメント形式（フォールバック）
-                prompt_parts.extend([
-                    f'<!-- begin of {role_upper} -->',
-                    msg['content'].strip(),
-                    f'<!-- end of {role_upper} -->',
-                    ''  # 空行で区切る
-                ])
+def create_backend(model_name: str, text_only: bool = False):
+    model_kind = "lm" if text_only else detect_model_kind(model_name)
 
-    # 最後の空行を削除して、ダブル改行で結合
-    return '\n'.join(prompt_parts[:-1])
-
-
-def handle_completion(prompt, options=None, images=None, max_image_size=768):
-    """completion API の処理
-
-    VLMモデルの場合、TypeScript側でプロンプトにimageトークンが挿入済み。
-    images が渡された場合は VLM 生成を使用する。
-    """
-    if options is None:
-        options = {}
-
-    # promptはTypeScript側で既にモデル固有処理済み
-
-    if images:
-        pil_images = load_and_resize_images(images, max_image_size)
-
-        import re
-        display_prompt = re.sub(r'(<\|image_pad\|>)+', '<|image_pad|>...', prompt)
-        sys.stderr.write(f"--- vlm completion (images: {len(pil_images)}, max_size: {max_image_size})\n{display_prompt}\n")
-
-        generate_text_vlm(prompt, pil_images, options)
-    else:
-        generate_text(prompt, options)
-
-
-def handle_chat_vlm(messages, images, options=None, max_image_size=768, tools=None, primer=None):
-    """VLMモデル用のチャット処理
-
-    messages: TypeScript側で画像プレースホルダー({type: "image"})が挿入済み
-    images: 画像ファイルパスの配列（プレースホルダーと位置が対応）
-    tools: ツール定義（テンプレートが対応している場合のみ使用）
-    primer: アシスタント応答のプリフィックス
-    """
-    if options is None:
-        options = {}
-
-    # primer処理
-    add_generation_prompt = True
-    if primer is not None:
-        messages.append({'role': 'assistant', 'content': primer})
-        add_generation_prompt = False
-
-    # processorのapply_chat_templateを直接使用
-    # systemメッセージのマージはTypeScript側でchat_restrictionsに基づき処理済み
-    # tools対応を試みる（テンプレートが対応していなければtools無しで実行）
-    try:
-        formatted_prompt = processor.apply_chat_template(
-            messages,
-            tools=tools,
-            add_generation_prompt=add_generation_prompt,
-            tokenize=False,
-        )
-    except TypeError:
-        formatted_prompt = processor.apply_chat_template(
-            messages,
-            add_generation_prompt=add_generation_prompt,
-            tokenize=False,
-        )
-
-    if primer is not None:
-        formatted_prompt = primer.join(formatted_prompt.split(primer)[0:-1]) + primer
-        print(primer, end='', flush=True)
-
-    # 画像ファイルを読み込み・リサイズ
-    pil_images = load_and_resize_images(images, max_image_size)
-
-    # image_padトークンを省略して表示（大量のパディングで読みづらいため）
-    import re
-    display_prompt = re.sub(r'(<\|image_pad\|>)+', '<|image_pad|>...', formatted_prompt)
-    sys.stderr.write(f"--- vlm prompt (images: {len(pil_images)}, max_size: {max_image_size})\n{display_prompt}\n")
-
-    generate_text_vlm(formatted_prompt, pil_images, options)
-
-
-def generate_text_vlm(prompt, images, options, stop_token_ids=None):
-    """VLMストリーミング生成"""
-    temperature = options.pop('temperature', 1.0) if 'temperature' in options else 1.0
-    max_tokens = options.pop('max_tokens', 1000) if 'max_tokens' in options else 1000
-    top_p = options.pop('top_p', 0.0) if 'top_p' in options else 0.0
-    top_k = options.pop('top_k', 0) if 'top_k' in options else 0
-
-    for response in vlm_stream_generate(
-        model, processor, prompt,
-        image=images if images else None,
-        max_tokens=max_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        top_k=top_k,
-    ):
-        # 追加 stop token チェック（tool call end 等）
-        if stop_token_ids and hasattr(response, 'token') and int(response.token) in stop_token_ids:
-            sys.stderr.write(f"--- stop token detected (vlm): {int(response.token)}\n")
-            print('\n', end='\0', flush=True)
-            return
-        print(response.text.replace('\0', ''), end='', flush=True)
-
-    print('\n', end='\0', flush=True)
-
-
-def generate_text(prompt, options):
-    """テキスト生成の共通処理
-
-    注意: optionsはTypeScript側で事前にバリデーション済み
-    - temperatureパラメータはsamplerオブジェクトに変換
-    - サポートされていないパラメータはTS側でフィルタリング
-    """
-    # デフォルトオプションの設定
-    default_options = {'max_tokens': 1000}
-
-    # temperatureパラメータを抽出してsamplerを作成
-    temperature = options.pop('temperature', 1.0) if 'temperature' in options else 1.0
-    top_p = options.pop('top_p', 0.0) if 'top_p' in options else 0.0
-    top_k = options.pop('top_k', 0) if 'top_k' in options else 0
-
-    # samplerオブジェクトを作成
-    sampler = make_sampler(temp=temperature, top_p=top_p, top_k=top_k)
-
-    # 残りのオプションとマージ
-    final_options = {**default_options, **options, 'sampler': sampler}
-
-    if isinstance(prompt, list):  # tokenized
-        sys.stderr.write(f"--- prompt: len={len(prompt)}\n")
-    else:
-        sys.stderr.write(f"--- prompt\n{prompt}\n")
-
-    eos_detected = False
-    for response in stream_generate(model, tokenizer, prompt, **final_options):
-        # トークンIDによるEOS判定（より確実）
-        if is_eod_token(response, tokenizer):
-            eos_detected = True
-            print('\n', end='\0', flush=True)
-            break
-        if not eos_detected:
-            print(response.text.replace('\0', ''), end='', flush=True)
-
-    if not eos_detected:
-        print('\n', end='\0', flush=True)
-
-def main():
-    while True:
-        req = read()
-        if req is None:
-            break
-        
-        method = req.get('method')
-        if not method:
-            sys.stderr.write("Error: 'method' field is required\n")
-            print('\n', end='\0', flush=True)
-            continue
-        
+    if model_kind == "vlm":
+        backend = MlxVlmBackend()
         try:
-            if method == 'capabilities':
-                handle_capabilities()
-            
-            elif method == 'format_test':
-                messages = req.get('messages')
-                if not messages:
-                    sys.stderr.write("Error: 'messages' field is required for format_test method\n")
-                    print('\n', end='\0', flush=True)
-                    continue
+            backend.load(model_name)
+            return backend, "vlm"
+        except (ValueError, Exception) as e:
+            sys.stderr.write(f"VLM load failed, falling back to LM: {e}\n")
 
-                options = req.get('options', {})
-                tools = req.get('tools')
-                handle_format_test(messages, options, tools)
+    backend = MlxLmBackend()
+    backend.load(model_name)
+    return backend, "lm"
 
-            elif method == 'chat':
-                messages = req.get('messages')
-                if not messages:
-                    sys.stderr.write("Error: 'messages' field is required for chat method\n")
-                    print('\n', end='\0', flush=True)
-                    continue
 
-                primer = req.get('primer')
-                options = req.get('options', {})
-                tools = req.get('tools')
-                images = req.get('images', [])
+if __name__ == "__main__":
+    backend, model_kind = create_backend(model_name, text_only)
 
-                if model_kind == "vlm":
-                    max_image_size = req.get('maxImageSize', 768)
-                    handle_chat_vlm(messages, images, options, max_image_size, tools, primer)
-                else:
-                    handle_chat(messages, primer, options, tools)
-            
-            elif method == 'completion':
-                prompt = req.get('prompt')
-                if not prompt:
-                    sys.stderr.write("Error: 'prompt' field is required for completion method\n")
-                    print('\n', end='\0', flush=True)
-                    continue
-                
-                options = req.get('options', {})
-                images = req.get('images', [])
-                max_image_size = req.get('maxImageSize', 768)
-                handle_completion(prompt, options, images if images else None, max_image_size)
-            
-            else:
-                sys.stderr.write(f"Error: Unknown method '{method}'\n")
-                print('\n', end='\0', flush=True)
-        
-        except Exception as e:
-            sys.stderr.write(f"Error processing request: {e}\n")
-            print('\n', end='\0', flush=True)
+    if drafter_model:
+        backend.load_drafter(drafter_model)
+        if draft_block_size is not None and hasattr(backend, 'draft_block_size'):
+            backend.draft_block_size = draft_block_size
 
+    capabilities = get_capabilities(backend.get_tokenizer())
+    capabilities["model_kind"] = model_kind
+    if model_kind == "lm":
+        capabilities["methods"].append("cache_prefill")
 
-if __name__ == "__main__":
-    main()
+    server = Server(backend, capabilities)
+    server.run()

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

@@ -0,0 +1,3 @@
+from backends.base import ModelBackend
+from backends.mlx_lm import MlxLmBackend
+from backends.mlx_vlm import MlxVlmBackend

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

@@ -0,0 +1,84 @@
+from abc import ABC, abstractmethod
+from typing import Any, Iterator
+
+
+class ModelBackend(ABC):
+    """Abstract base class for model backends."""
+
+    @abstractmethod
+    def load(self, model_name: str) -> None:
+        """Load the target model."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_tokenizer(self) -> Any:
+        """Return the tokenizer or processor."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def stream_generate(
+        self, prompt: str | list[int], options: dict, images: list | None = None,
+        prompt_cache: list | None = None,
+    ) -> Iterator[Any]:
+        """Stream generation results."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def supports_vision(self) -> bool:
+        """Return whether image input is supported."""
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def model_kind(self) -> str:
+        """Return "lm" or "vlm"."""
+        raise NotImplementedError
+
+    def load_drafter(self, drafter_model: str) -> None:
+        """Load a drafter model for speculative decoding."""
+        raise NotImplementedError(
+            f"{type(self).__name__} does not support drafter models"
+        )
+
+    def has_drafter(self) -> bool:
+        """Return whether a drafter model is loaded."""
+        return False
+
+    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:
+        """Build a KV cache from a prompt prefix."""
+        raise NotImplementedError(
+            f"{type(self).__name__} does not support prompt caching"
+        )
+
+    def load_cache_from_file(self, cache_path: str) -> list | None:
+        """Load a prompt cache from file, or None."""
+        return None
+
+    def get_cache_offset(self, prompt_cache: list) -> int:
+        """Get the number of tokens stored in a loaded prompt cache."""
+        if not prompt_cache:
+            return 0
+        layer0 = prompt_cache[0]
+        if hasattr(layer0, 'offset'):
+            off = layer0.offset
+            return int(off.item() if hasattr(off, 'item') else off)
+        if hasattr(layer0, 'caches'):
+            for c in layer0.caches:
+                if hasattr(c, 'offset'):
+                    off = c.offset
+                    return int(off.item() if hasattr(off, 'item') else off)
+        try:
+            return int(layer0[0].shape[2])
+        except Exception:
+            pass
+        if hasattr(layer0, 'keys') and layer0.keys is not None:
+            return int(layer0.keys.shape[2])
+        return 0