@modular-prompt/driver 0.14.0 → 0.15.0

package/src/pytorch/python/__main__.py

@@ -0,0 +1,19 @@
+import os
+import sys
+
+from backends import TransformersLmBackend
+from server import Server
+from utils.token_utils import get_capabilities
+
+model_name = sys.argv[1] if len(sys.argv) > 1 else "gpt2"
+device = os.environ.get("PYTORCH_DEVICE", "cpu")
+
+if __name__ == "__main__":
+    backend = TransformersLmBackend(device=device)
+    backend.load(model_name)
+
+    capabilities = get_capabilities(backend.get_tokenizer())
+    capabilities["model_kind"] = "lm"
+
+    server = Server(backend, capabilities)
+    server.run()

package/src/pytorch/python/backends/__init__.py

@@ -0,0 +1,3 @@
+from backends.transformers_lm import TransformersLmBackend
+
+__all__ = ["TransformersLmBackend"]

package/src/pytorch/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

package/src/pytorch/python/backends/transformers_lm.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from threading import Thread
+from typing import Any, Iterator
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+
+from backends.base import ModelBackend
+from utils.token_utils import is_eod_token
+
+
+@dataclass
+class StreamChunk:
+    text: str
+    prompt_tokens: int | None = None
+    generation_tokens: int | None = None
+    finish_reason: str | None = None
+
+
+class TransformersLmBackend(ModelBackend):
+    """Transformers causal LM backend (text-only, CPU-first)."""
+
+    def __init__(self, device: str | None = None) -> None:
+        self.model: Any | None = None
+        self.tokenizer: Any | None = None
+        self._device_name = device or os.environ.get("PYTORCH_DEVICE", "cpu")
+        self._device = torch.device(self._device_name)
+
+    def load(self, model_name: str) -> None:
+        trust_remote_code = os.environ.get("PYTORCH_TRUST_REMOTE_CODE", "").lower() in (
+            "1",
+            "true",
+            "yes",
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            trust_remote_code=trust_remote_code,
+        )
+        if self.tokenizer.pad_token is None and self.tokenizer.eos_token is not None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        dtype = torch.float32 if self._device.type == "cpu" else torch.float16
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            trust_remote_code=trust_remote_code,
+            torch_dtype=dtype,
+        )
+        self.model.to(self._device)
+        self.model.eval()
+
+    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[StreamChunk]:
+        if images:
+            raise ValueError("TransformersLmBackend does not support vision input")
+        if self.model is None or self.tokenizer is None:
+            raise RuntimeError("Model is not loaded")
+
+        final_options = {"max_tokens": 256, **options}
+        max_new_tokens = int(final_options.pop("max_tokens", 256))
+        temperature = float(final_options.pop("temperature", 1.0))
+        top_p = final_options.pop("top_p", None)
+        top_k = final_options.pop("top_k", None)
+
+        if isinstance(prompt, list):
+            input_ids = torch.tensor([prompt], device=self._device)
+            prompt_token_count = len(prompt)
+        else:
+            encoded = self.tokenizer(prompt, return_tensors="pt")
+            input_ids = encoded["input_ids"].to(self._device)
+            prompt_token_count = int(input_ids.shape[-1])
+
+        do_sample = temperature > 0
+        gen_kwargs: dict[str, Any] = {
+            "input_ids": input_ids,
+            "max_new_tokens": max_new_tokens,
+            "do_sample": do_sample,
+        }
+        if do_sample:
+            gen_kwargs["temperature"] = temperature
+        if top_p is not None:
+            gen_kwargs["top_p"] = float(top_p)
+        if top_k is not None:
+            gen_kwargs["top_k"] = int(top_k)
+
+        streamer = TextIteratorStreamer(
+            self.tokenizer,
+            skip_special_tokens=True,
+            skip_prompt=True,
+        )
+        gen_kwargs["streamer"] = streamer
+
+        thread = Thread(target=self.model.generate, kwargs=gen_kwargs)
+        thread.start()
+
+        generation_tokens = 0
+        for text in streamer:
+            generation_tokens += 1
+            chunk = StreamChunk(
+                text=text,
+                prompt_tokens=prompt_token_count if generation_tokens == 1 else None,
+                generation_tokens=generation_tokens if generation_tokens == 1 else None,
+            )
+            if is_eod_token(chunk, self.tokenizer):
+                chunk.finish_reason = "stop"
+                yield chunk
+                break
+            yield chunk
+
+        thread.join()
+
+    def supports_vision(self) -> bool:
+        return False
+
+    @property
+    def model_kind(self) -> str:
+        return "lm"

package/src/pytorch/python/handlers/__init__.py

@@ -0,0 +1,6 @@
+from handlers.capabilities import handle_capabilities
+from handlers.completion import handle_completion
+from handlers.format_test import handle_format_test
+from handlers.generate import handle_generate
+from handlers.render import handle_render
+from handlers.tokenize import handle_tokenize

package/src/pytorch/python/handlers/cancel.py

@@ -0,0 +1,53 @@
+"""Cancel request handling for in-flight streaming generation."""
+
+from __future__ import annotations
+
+import json
+import select
+import sys
+
+_cancel_requested = False
+
+
+def request_cancel() -> None:
+    global _cancel_requested
+    _cancel_requested = True
+
+
+def reset_cancel() -> None:
+    global _cancel_requested
+    _cancel_requested = False
+
+
+def is_cancel_requested() -> bool:
+    return _cancel_requested
+
+
+def poll_cancel() -> bool:
+    """Non-blocking check for a cancel command on stdin during streaming."""
+    global _cancel_requested
+    if _cancel_requested:
+        return True
+
+    try:
+        ready, _, _ = select.select([sys.stdin], [], [], 0)
+    except (ValueError, OSError):
+        return False
+
+    if not ready:
+        return False
+
+    line = sys.stdin.readline()
+    if not line:
+        return False
+
+    try:
+        req = json.loads(line)
+    except json.JSONDecodeError:
+        return False
+
+    if req.get("method") == "cancel":
+        _cancel_requested = True
+        return True
+
+    return False

package/src/pytorch/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/pytorch/python/handlers/completion.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from backends.base import ModelBackend
+from handlers.generate import handle_generate
+
+
+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（後方互換）— generate に委譲"""
+    handle_generate(backend, prompt, options, images, max_image_size)

package/src/pytorch/python/handlers/format_test.py

@@ -0,0 +1,70 @@
+import json
+
+from backends.base import ModelBackend
+from utils.prompt_builder import generate_merged_prompt, supports_chat_template
+
+
+def handle_format_test(
+    backend: ModelBackend,
+    capabilities: dict,
+    messages: list,
+    options: dict | None = None,
+    tools: list | None = None,
+) -> None:
+    """フォーマットテスト API の処理（実際に生成せずフォーマットのみ）"""
+    if options is None:
+        options = {}
+
+    tokenizer = backend.get_tokenizer()
+    result = {
+        "formatted_prompt": None,
+        "template_applied": False,
+        "model_specific_processing": None,
+        "error": None,
+    }
+
+    try:
+        if supports_chat_template(tokenizer):
+            result["model_specific_processing"] = messages
+
+            primer = options.get("primer")
+            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:
+                formatted_prompt = tokenizer.apply_chat_template(
+                    fmt_messages,
+                    tools=tools,
+                    add_generation_prompt=add_generation_prompt,
+                    tokenize=False,
+                )
+            except TypeError:
+                formatted_prompt = tokenizer.apply_chat_template(
+                    fmt_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
+                )
+
+            result["formatted_prompt"] = formatted_prompt
+            result["template_applied"] = True
+        else:
+            formatted_prompt = generate_merged_prompt(messages, capabilities)
+            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)

package/src/pytorch/python/handlers/generate.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import json
+
+from backends.base import ModelBackend
+from handlers.cancel import poll_cancel
+
+
+def _stream_to_stdout(
+    backend: ModelBackend,
+    prompt: str | list[int],
+    options: dict,
+    images: list | None = None,
+    primer: str | None = None,
+) -> None:
+    if images:
+        raise ValueError("PyTorch LIP backend does not support images in Phase 6")
+
+    if primer is not None:
+        print(primer, end="", flush=True)
+
+    last_response = None
+    for response in backend.stream_generate(prompt, options, images):
+        if poll_cancel():
+            break
+        print(response.text.replace("\0", "").replace("\x1e", ""), end="", flush=True)
+        last_response = response
+
+    meta: dict = {}
+    if last_response is not None:
+        if last_response.prompt_tokens is not None:
+            meta["prompt_tokens"] = last_response.prompt_tokens
+        if last_response.generation_tokens is not None:
+            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_generate(
+    backend: ModelBackend,
+    prompt: str | list[int],
+    options: dict | None = None,
+    images: list | None = None,
+    max_image_size: int = 768,
+    primer: str | None = None,
+    cache_path: str | None = None,
+    cache_trim_tokens: int | None = None,
+) -> None:
+    """LIP generate: 整形済み prompt のストリーム推論（KV キャッシュ非対応）"""
+    if cache_path or cache_trim_tokens is not None:
+        raise ValueError("PyTorch LIP backend does not support prompt caching")
+
+    if options is None:
+        options = {}
+
+    final_options = dict(options)
+    final_options.pop("trust_remote_code", None)
+
+    _stream_to_stdout(
+        backend,
+        prompt,
+        final_options,
+        images=images,
+        primer=primer,
+    )

package/src/pytorch/python/handlers/render.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+import json
+
+from backends.base import ModelBackend
+from utils.template_render import apply_chat_template_prompt
+
+
+def handle_render(
+    backend: ModelBackend,
+    messages: list,
+    options: dict | None = None,
+    tools: list | None = None,
+    reasoning_effort: str | None = None,
+) -> None:
+    """LIP render: apply_chat_template のみ（推論しない）"""
+    if options is None:
+        options = {}
+
+    result: dict = {
+        "formatted_prompt": None,
+        "error": None,
+    }
+
+    try:
+        trust_remote_code = options.get("trust_remote_code")
+        primer = options.get("primer")
+        prompt = apply_chat_template_prompt(
+            backend,
+            messages,
+            primer=primer,
+            tools=tools,
+            reasoning_effort=reasoning_effort,
+            trust_remote_code=trust_remote_code,
+        )
+        result["formatted_prompt"] = prompt
+    except Exception as e:
+        result["error"] = str(e)
+
+    print(json.dumps(result), end="\0", flush=True)

package/src/pytorch/python/handlers/tokenize.py

@@ -0,0 +1,63 @@
+import json
+
+from backends.base import ModelBackend
+from utils.prompt_builder import generate_merged_prompt, supports_chat_template
+
+
+def handle_tokenize(
+    backend: ModelBackend,
+    capabilities: dict,
+    messages: list,
+    tools: list | None = None,
+    reasoning_effort: str | None = None,
+) -> None:
+    """メッセージをchat template適用後にトークン化して返す"""
+    tokenizer = backend.get_tokenizer()
+
+    result = {
+        "token_ids": None,
+        "token_count": 0,
+        "error": None,
+    }
+
+    try:
+        # apply_chat_templateのfallbackパターン (chat.py L165-188 と同じ)
+        # 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
+
+        if supports_chat_template(tokenizer):
+            # chat.py と同じfallbackチェーン
+            prompt = None
+            for kwargs in [extra_kwargs, {k: v for k, v in extra_kwargs.items() if k == "tools"}, {}]:
+                try:
+                    prompt = tokenizer.apply_chat_template(
+                        messages,
+                        add_generation_prompt=False,
+                        tokenize=False,
+                        **kwargs,
+                    )
+                    break
+                except TypeError:
+                    continue
+
+            if prompt is None:
+                prompt = str(messages)
+        else:
+            prompt = generate_merged_prompt(messages, capabilities)
+
+        # トークン化
+        add_special = tokenizer.bos_token is None or not prompt.startswith(
+            tokenizer.bos_token or ""
+        )
+        token_ids = tokenizer.encode(prompt, add_special_tokens=add_special)
+
+        result["token_ids"] = token_ids
+        result["token_count"] = len(token_ids)
+    except Exception as e:
+        result["error"] = str(e)
+
+    print(json.dumps(result), end="\0", flush=True)

package/src/pytorch/python/pyproject.toml

@@ -0,0 +1,36 @@
+[project]
+name = "pytorch_driver"
+version = "0.1.0"
+description = "PyTorch (Transformers) driver for modular-prompt — cpu-minimal runtime"
+requires-python = ">=3.10,<3.14"
+dependencies = [
+    "safetensors==0.7.0",
+    "tokenizers==0.22.2",
+    "transformers==4.57.6",
+]
+
+[dependency-groups]
+dev = ["pytest>=9.0"]
+
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[tool.setuptools]
+py-modules = ["__main__", "server"]
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["backends*", "handlers*", "utils*"]
+
+# setup-pytorch は CPU wheel を明示インストールする。手動で CUDA 等に差し替える場合はドキュメント参照。
+[[tool.uv.index]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
+
+[tool.uv.sources]
+torch = { index = "pytorch-cpu" }