@modular-prompt/driver 0.16.0 → 0.17.0

package/src/pytorch/templates/cuda/tests/test_server.py

@@ -0,0 +1,145 @@
+import pytest
+
+pytest.importorskip("torch")
+
+import json
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+
+class _Tokenizer:
+    bos_token = "<bos>"
+    chat_template = "template"
+
+    def apply_chat_template(self, messages, **kwargs):
+        del messages, kwargs
+        return "prefix"
+
+
+class _Backend:
+    model_kind = "lm"
+
+    def __init__(self):
+        self.tokenizer = _Tokenizer()
+        self.cache = None
+        self.generate_prompt = None
+        self.pending_cache_write_tokens = 0
+
+    def get_tokenizer(self):
+        return self.tokenizer
+
+    def cache_prefill(self, cache_path, prompt, **kwargs):
+        del prompt, kwargs
+        self.cache = object()
+        self.pending_cache_write_tokens = 2
+        return {
+            "cache_path": cache_path,
+            "token_count": 2,
+            "cache_write_tokens": 2,
+        }
+
+    def load_cache_from_file(self, cache_path, **kwargs):
+        del kwargs
+        return self.cache if cache_path == "memory://prefix" else None
+
+    def get_cache_offset(self, prompt_cache):
+        return 2 if prompt_cache is self.cache else 0
+
+    def consume_cache_write_tokens(self, cache_path):
+        del cache_path
+        result = self.pending_cache_write_tokens
+        self.pending_cache_write_tokens = 0
+        return result
+
+    def tokenize_prompt(self, prompt):
+        assert prompt == "prefix suffix"
+        return [1, 2, 3]
+
+    def stream_generate(self, prompt, options, images=None, prompt_cache=None):
+        del options, images
+        self.generate_prompt = prompt
+        assert prompt_cache is self.cache
+        yield SimpleNamespace(
+            text="ok",
+            prompt_tokens=3,
+            generation_tokens=1,
+            cache_read_tokens=2,
+        )
+
+
+def _make_server():
+    from server import Server
+
+    backend = MagicMock()
+    return Server(backend, {"methods": ["capabilities"], "model_kind": "lm"})
+
+
+def test_cache_prefill_dispatches_to_handler(capsys):
+    server = _make_server()
+
+    with patch("server.handle_cache_prefill") as mock_prefill:
+        server._dispatch(
+            {
+                "method": "cache_prefill",
+                "cache_path": "memory://prefix",
+                "messages": [{"role": "user", "content": "hello"}],
+            }
+        )
+
+    mock_prefill.assert_called_once()
+    assert mock_prefill.call_args.args[2:] == (
+        "memory://prefix",
+        [{"role": "user", "content": "hello"}],
+    )
+    assert capsys.readouterr().out == ""
+
+
+def test_cache_prefill_requires_path_and_messages(capsys):
+    server = _make_server()
+
+    server._dispatch({"method": "cache_prefill"})
+
+    assert capsys.readouterr().out.endswith("\0")
+
+
+def test_cache_prefill_then_generate_with_cache_via_server(capsys):
+    from server import Server
+
+    backend = _Backend()
+    server = Server(backend, {"methods": ["cache_prefill"], "model_kind": "lm"})
+
+    server._dispatch(
+        {
+            "method": "cache_prefill",
+            "cache_path": "memory://prefix",
+            "messages": [{"role": "user", "content": "hello"}],
+        }
+    )
+    prefill_output = capsys.readouterr().out
+    assert json.loads(prefill_output.split("\0", 1)[0]) == {
+        "cache_path": "memory://prefix",
+        "token_count": 2,
+        "cache_write_tokens": 2,
+    }
+
+    server._dispatch(
+        {
+            "method": "generate",
+            "prompt": "prefix suffix",
+            "options": {"max_tokens": 1},
+            "cache_path": "memory://prefix",
+        }
+    )
+    generate_output = capsys.readouterr().out
+    assert generate_output.startswith("ok")
+    meta = json.loads(
+        generate_output.split("\x1e__META__:", 1)[1].split("\0", 1)[0]
+    )
+    assert meta == {
+        "prompt_tokens": 3,
+        "generation_tokens": 1,
+        "cache_read_tokens": 2,
+        "cache_write_tokens": 2,
+        "cache_loaded": True,
+    }
+    assert backend.generate_prompt == [3]

package/src/pytorch/templates/cuda/tests/test_transformers_errors.py

@@ -0,0 +1,89 @@
+from types import SimpleNamespace
+
+import pytest
+
+from utils.transformers_errors import (
+    MIN_TRANSFORMERS_VERSION,
+    extract_unsupported_model_type,
+    unsupported_model_type_error,
+)
+
+
+def test_extracts_model_type_from_transformers_value_error():
+    error = ValueError(
+        "The checkpoint has model type `qwen3_5` but Transformers does not "
+        "recognize this architecture."
+    )
+
+    assert extract_unsupported_model_type(error) == "qwen3_5"
+
+
+def test_extracts_model_type_from_older_registry_key_error():
+    assert extract_unsupported_model_type(KeyError("qwen3_5")) == "qwen3_5"
+
+
+def test_leaves_unrelated_errors_unchanged():
+    assert extract_unsupported_model_type(ValueError("invalid weights")) is None
+    assert extract_unsupported_model_type(KeyError("model_type")) is None
+
+
+def test_error_includes_runtime_requirement_and_setup_guidance():
+    error = unsupported_model_type_error("Qwen/Qwen3.5-0.8B", "qwen3_5", "4.57.6")
+    message = str(error)
+
+    assert "qwen3_5" in message
+    assert "transformers 4.57.6" in message
+    assert f"transformers>={MIN_TRANSFORMERS_VERSION}" in message
+    assert "setup-pytorch" in message
+    assert "modular-prompt-runtime sync pytorch" in message
+
+
+def test_backend_wraps_unknown_architecture_error(monkeypatch):
+    pytest.importorskip("torch")
+    import backends.transformers_lm as backend_module
+
+    tokenizer = SimpleNamespace(pad_token=None, eos_token="<eos>")
+    backend = backend_module.TransformersLmBackend(device="cpu")
+
+    def load_tokenizer(*args, **kwargs):
+        return tokenizer
+
+    def load_model(*args, **kwargs):
+        raise ValueError(
+            "The checkpoint has model type `qwen3_5` but Transformers does not "
+            "recognize this architecture."
+        )
+
+    monkeypatch.setattr(backend_module.AutoTokenizer, "from_pretrained", load_tokenizer)
+    monkeypatch.setattr(
+        backend_module.AutoModelForCausalLM,
+        "from_pretrained",
+        load_model,
+    )
+
+    with pytest.raises(RuntimeError, match="qwen3_5") as raised:
+        backend.load("Qwen/Qwen3.5-0.8B")
+
+    assert f"transformers {backend_module.transformers.__version__}" in str(raised.value)
+    assert "transformers>=5.14.0" in str(raised.value)
+
+
+def test_backend_wraps_unknown_architecture_error_from_tokenizer(monkeypatch):
+    pytest.importorskip("torch")
+    import backends.transformers_lm as backend_module
+
+    backend = backend_module.TransformersLmBackend(device="cpu")
+
+    def load_tokenizer(*args, **kwargs):
+        raise ValueError(
+            "The checkpoint has model type `qwen3_5` but Transformers does not "
+            "recognize this architecture."
+        )
+
+    monkeypatch.setattr(backend_module.AutoTokenizer, "from_pretrained", load_tokenizer)
+
+    with pytest.raises(RuntimeError, match="qwen3_5") as raised:
+        backend.load("Qwen/Qwen3.5-0.8B")
+
+    assert f"transformers {backend_module.transformers.__version__}" in str(raised.value)
+    assert "transformers>=5.14.0" in str(raised.value)

package/src/pytorch/templates/cuda/tests/test_transformers_lm_cache.py

@@ -0,0 +1,288 @@
+import json
+from types import SimpleNamespace
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from backends.transformers_lm import TransformersLmBackend
+from handlers.generate import handle_generate
+from tokenizers import Tokenizer
+from tokenizers.models import WordLevel
+from tokenizers.pre_tokenizers import Whitespace
+from transformers import GPT2Config, GPT2LMHeadModel, PreTrainedTokenizerFast
+
+
+class _Tokenizer:
+    bos_token = "<bos>"
+    pad_token = "<pad>"
+    eos_token = "<eos>"
+
+    def __init__(self):
+        self.prompts = {
+            "prefix": [1, 2],
+            "prefix suffix": [1, 2, 3],
+            "different suffix": [9, 8, 7],
+        }
+        self.token_text = {
+            10: "a",
+            11: "b",
+            12: "c",
+            13: "d",
+        }
+
+    def encode(self, prompt, add_special_tokens):
+        del add_special_tokens
+        return self.prompts[prompt]
+
+    def decode(self, token_ids, **kwargs):
+        del kwargs
+        if isinstance(token_ids, torch.Tensor):
+            token_ids = token_ids.flatten().tolist()
+        return "".join(self.token_text.get(int(token_id), "x") for token_id in token_ids)
+
+
+class _Model:
+    def __init__(self, past_key_values, generated_texts=None, generated_token_batches=None):
+        self.past_key_values = past_key_values
+        self.generated_texts = generated_texts or ["generated"]
+        self.generated_token_batches = generated_token_batches
+        self.forward_calls = []
+        self.generate_calls = []
+
+    def __call__(self, **kwargs):
+        self.forward_calls.append(kwargs)
+        return SimpleNamespace(past_key_values=self.past_key_values)
+
+    def _emit_stream(self, streamer, input_ids, token_batches):
+        streamer.put(input_ids.cpu())
+        for token_batch in token_batches:
+            streamer.put(torch.tensor(token_batch, dtype=torch.long))
+        streamer.end()
+
+    def generate(self, **kwargs):
+        self.generate_calls.append(kwargs)
+        streamer = kwargs["streamer"]
+        if self.generated_token_batches is not None:
+            self._emit_stream(streamer, kwargs["input_ids"], self.generated_token_batches)
+            return
+        token_batches = [[10 + index] for index in range(len(self.generated_texts))]
+        self._emit_stream(streamer, kwargs["input_ids"], token_batches)
+
+
+def _backend(generated_texts=None, generated_token_batches=None):
+    # Legacy tuple-shaped caches remain supported by Transformers and make the
+    # test independent of a specific Cache class implementation.
+    key = torch.zeros((1, 1, 2, 4))
+    value = torch.zeros((1, 1, 2, 4))
+    past_key_values = ((key, value),)
+
+    backend = TransformersLmBackend(device="cpu")
+    backend.tokenizer = _Tokenizer()
+    backend.model = _Model(past_key_values, generated_texts, generated_token_batches)
+    return backend, past_key_values
+
+
+def _tiny_gpt2_backend():
+    vocab = {
+        "<pad>": 0,
+        "<unk>": 1,
+        "hello": 2,
+        "alpha": 3,
+        "beta": 4,
+        "gamma": 5,
+    }
+    tokenizer = Tokenizer(WordLevel(vocab=vocab, unk_token="<unk>"))
+    tokenizer.pre_tokenizer = Whitespace()
+    fast_tokenizer = PreTrainedTokenizerFast(
+        tokenizer_object=tokenizer,
+        unk_token="<unk>",
+        pad_token="<pad>",
+    )
+    config = GPT2Config(
+        vocab_size=len(vocab),
+        n_positions=32,
+        n_ctx=32,
+        n_embd=16,
+        n_layer=1,
+        n_head=1,
+        pad_token_id=vocab["<pad>"],
+        eos_token_id=None,
+    )
+    model = GPT2LMHeadModel(config)
+    # Make greedy output deterministic and non-special so the streamer emits
+    # several chunks without downloading a model in CI.
+    model.lm_head = torch.nn.Linear(config.n_embd, config.vocab_size, bias=True)
+    with torch.no_grad():
+        model.lm_head.weight.zero_()
+        model.lm_head.bias.fill_(-1)
+        model.lm_head.bias[vocab["alpha"]] = 1
+    model.eval()
+
+    backend = TransformersLmBackend()
+    backend.tokenizer = fast_tokenizer
+    backend.model = model
+    return backend
+
+
+def test_cache_prefill_keeps_past_key_values_in_process():
+    backend, past_key_values = _backend()
+
+    result = backend.cache_prefill("memory://prefix", "prefix")
+
+    assert result == {
+        "cache_path": "memory://prefix",
+        "token_count": 2,
+        "cache_write_tokens": 2,
+    }
+    assert backend.load_cache_from_file("memory://prefix") is past_key_values
+    assert backend.consume_cache_write_tokens("memory://prefix") == 2
+    assert backend.consume_cache_write_tokens("memory://prefix") == 0
+    call = backend.model.forward_calls[0]
+    assert call["input_ids"].tolist() == [[1, 2]]
+    assert call["input_ids"].device == backend._device
+    assert call["use_cache"] is True
+
+
+def test_load_cache_validates_prompt_prefix_without_reading_a_file():
+    backend, past_key_values = _backend()
+    backend.cache_prefill("memory://prefix", "prefix")
+
+    assert backend.load_cache_from_file(
+        "memory://prefix", prompt="prefix suffix"
+    ) is past_key_values
+    assert backend.load_cache_from_file(
+        "memory://prefix", prompt="different suffix"
+    ) is None
+    assert backend.load_cache_from_file("/tmp/not-created.cache") is None
+
+
+def test_stream_generate_passes_cached_prefix_and_reports_usage():
+    backend, past_key_values = _backend()
+    backend.cache_prefill("memory://prefix", "prefix")
+
+    chunks = list(
+        backend.stream_generate(
+            [3],
+            {"max_tokens": 1, "temperature": 0},
+            prompt_cache=past_key_values,
+        )
+    )
+
+    assert chunks[-1].text == "a"
+    assert chunks[-1].generation_tokens == 1
+    call = backend.model.generate_calls[0]
+    assert call["input_ids"].tolist() == [[3]]
+    assert call["past_key_values"] is not past_key_values
+    assert torch.equal(call["past_key_values"][0][0], past_key_values[0][0])
+    assert call["cache_position"].tolist() == [2]
+    assert call["attention_mask"].tolist() == [[1, 1, 1]]
+    first_meta_chunk = next(chunk for chunk in chunks if chunk.prompt_tokens is not None)
+    assert first_meta_chunk.prompt_tokens == 3
+    assert first_meta_chunk.cache_read_tokens == 2
+
+
+def test_stream_generate_reports_cumulative_generation_tokens_for_multiple_chunks():
+    backend, past_key_values = _backend(["a", "b", "c"])
+    backend.cache_prefill("memory://prefix", "prefix")
+
+    chunks = list(
+        backend.stream_generate(
+            [3],
+            {"max_tokens": 3, "temperature": 0},
+            prompt_cache=past_key_values,
+        )
+    )
+
+    assert chunks[-1].text == "abc"
+    assert chunks[-1].generation_tokens == 3
+    first_meta_chunk = next(chunk for chunk in chunks if chunk.prompt_tokens is not None)
+    assert first_meta_chunk.prompt_tokens == 3
+    assert first_meta_chunk.cache_read_tokens == 2
+
+
+def test_stream_generate_counts_token_ids_not_empty_text_chunks(capsys):
+    backend, _ = _backend(generated_token_batches=[[10], [11], [12], [13]])
+
+    chunks = list(
+        backend.stream_generate(
+            "prefix suffix",
+            {"max_tokens": 4, "temperature": 0},
+        )
+    )
+
+    # TextIteratorStreamer emits an empty chunk while buffering each token and
+    # one final text chunk, so chunk count is five for four generated tokens.
+    assert len(chunks) == 5
+    assert chunks[0].text == ""
+    assert chunks[-1].text == "abcd"
+    assert chunks[-1].generation_tokens == 4
+
+    handle_generate(
+        backend,
+        "prefix suffix",
+        options={"max_tokens": 4, "temperature": 0},
+    )
+    output = capsys.readouterr().out
+    meta = json.loads(output.split("\x1e__META__:", 1)[1].split("\0", 1)[0])
+    assert meta["prompt_tokens"] == 3
+    assert meta["generation_tokens"] == 4
+
+
+def test_generate_handler_reports_backend_prefill_write_usage(capsys):
+    backend, _ = _backend()
+    backend.cache_prefill("memory://prefix", "prefix")
+
+    handle_generate(
+        backend,
+        "prefix suffix",
+        options={"max_tokens": 1, "temperature": 0},
+        cache_path="memory://prefix",
+    )
+
+    output = capsys.readouterr().out
+    assert output.startswith("a")
+    meta = output.split("\x1e__META__:", 1)[1].split("\0", 1)[0]
+    assert json.loads(meta) == {
+        "prompt_tokens": 3,
+        "generation_tokens": 1,
+        "cache_read_tokens": 2,
+        "cache_write_tokens": 2,
+        "cache_loaded": True,
+    }
+
+
+def test_generate_handler_preserves_usage_for_tiny_gpt2_multiple_chunks(capsys):
+    backend = _tiny_gpt2_backend()
+    options = {"max_tokens": 4, "temperature": 0}
+
+    expected_chunks = list(backend.stream_generate("hello", options))
+    # The real TextIteratorStreamer buffers the generated word, so the first
+    # text chunk is empty and the final chunk flushes all four generated
+    # tokens.  Token usage must not be derived from this five-chunk sequence.
+    assert len(expected_chunks) == 5
+    assert expected_chunks[0].text == ""
+    expected_generation_tokens = expected_chunks[-1].generation_tokens
+    assert expected_generation_tokens == 4
+
+    handle_generate(backend, "hello", options=options)
+
+    output = capsys.readouterr().out
+    meta = json.loads(output.split("\x1e__META__:", 1)[1].split("\0", 1)[0])
+    assert meta["prompt_tokens"] == len(backend.tokenize_prompt("hello"))
+    assert meta["generation_tokens"] == expected_generation_tokens
+
+
+def test_cache_prefill_rejects_phase_two_arguments():
+    backend, _ = _backend()
+
+    try:
+        backend.cache_prefill(
+            "memory://prefix",
+            "prefix",
+            base_cache_path="memory://base",
+        )
+    except ValueError as error:
+        assert "incremental prefill" in str(error)
+    else:
+        raise AssertionError("Phase 2 incremental prefill must be rejected")

package/src/pytorch/templates/cuda/utils/chat_template_constraints.py

@@ -0,0 +1,164 @@
+"""
+チャットテンプレートの制約検出
+
+tokenizerのapply_chat_templateを使用して、
+モデルがサポートするメッセージパターンの制約を検出する。
+"""
+
+
+def detect_chat_restrictions(tokenizer) -> dict:
+    """
+    チャットテンプレートの制約を検出
+
+    Args:
+        tokenizer: HuggingFace tokenizer (apply_chat_template対応)
+
+    Returns:
+        dict: chat_restrictions情報
+            {
+                "single_system_at_start": bool,
+                "max_system_messages": int,
+                "alternating_turns": bool,
+                "requires_user_last": bool,
+                "allow_empty_messages": bool
+            }
+    """
+    if not hasattr(tokenizer, 'apply_chat_template'):
+        return None
+
+    # テストパターンを実行
+    test_results = {}
+    for pattern in _get_test_patterns():
+        try:
+            tokenizer.apply_chat_template(
+                pattern['messages'],
+                tokenize=False,
+                add_generation_prompt=False
+            )
+            test_results[pattern['name']] = {'success': True}
+        except Exception as e:
+            test_results[pattern['name']] = {'error': str(e)}
+
+    # テスト結果から制約を推論
+    return _infer_restrictions_from_results(test_results)
+
+
+def _get_test_patterns():
+    """テストパターンの定義"""
+    return [
+        # 基本パターン
+        {
+            'name': 'basic',
+            'messages': [
+                {'role': 'user', 'content': 'Hello'}
+            ]
+        },
+
+        # システムメッセージ付き
+        {
+            'name': 'with-system',
+            'messages': [
+                {'role': 'system', 'content': 'You are a helpful assistant.'},
+                {'role': 'user', 'content': 'Hello'}
+            ]
+        },
+
+        # 複数システムメッセージ
+        {
+            'name': 'multi-system',
+            'messages': [
+                {'role': 'system', 'content': 'First system message.'},
+                {'role': 'system', 'content': 'Second system message.'},
+                {'role': 'user', 'content': 'Hello'}
+            ]
+        },
+
+        # 連続ユーザーメッセージ
+        {
+            'name': 'consecutive-user',
+            'messages': [
+                {'role': 'user', 'content': 'First question'},
+                {'role': 'user', 'content': 'Second question'}
+            ]
+        },
+
+        # アシスタントで終わる
+        {
+            'name': 'assistant-last',
+            'messages': [
+                {'role': 'user', 'content': 'Hello'},
+                {'role': 'assistant', 'content': 'Hi there!'}
+            ]
+        },
+
+        # 交互の会話
+        {
+            'name': 'alternating',
+            'messages': [
+                {'role': 'user', 'content': 'Question 1'},
+                {'role': 'assistant', 'content': 'Answer 1'},
+                {'role': 'user', 'content': 'Question 2'}
+            ]
+        },
+
+        # 空メッセージ
+        {
+            'name': 'empty-message',
+            'messages': [
+                {'role': 'user', 'content': ''}
+            ]
+        },
+
+        # システムメッセージが途中にある
+        {
+            'name': 'system-middle',
+            'messages': [
+                {'role': 'user', 'content': 'First'},
+                {'role': 'system', 'content': 'System in middle'},
+                {'role': 'user', 'content': 'Second'}
+            ]
+        }
+    ]
+
+
+def _infer_restrictions_from_results(test_results: dict) -> dict:
+    """
+    テスト結果から制約を推論
+
+    Args:
+        test_results: テストパターン名をキーとした結果の辞書
+
+    Returns:
+        dict: 検出された制約
+    """
+    restrictions = {}
+
+    # システムメッセージの制約を検出
+    with_system = test_results.get('with-system')
+    multi_system = test_results.get('multi-system')
+
+    if with_system and 'error' in with_system:
+        # 単独のsystemメッセージもエラー → systemロール自体がサポートされていない
+        restrictions['max_system_messages'] = 0
+    elif multi_system and 'error' in multi_system:
+        # 複数はエラーだが単独は成功 → 最大1つまで
+        restrictions['single_system_at_start'] = True
+        restrictions['max_system_messages'] = 1
+    # それ以外（両方成功）→ max_system_messagesキーを設定しない（無制限）
+
+    # 連続ユーザーメッセージのテスト
+    consecutive_user = test_results.get('consecutive-user')
+    if consecutive_user and 'error' in consecutive_user:
+        restrictions['alternating_turns'] = True
+
+    # アシスタントで終わるテスト
+    assistant_last = test_results.get('assistant-last')
+    if assistant_last and 'error' in assistant_last:
+        restrictions['requires_user_last'] = True
+
+    # 空メッセージのテスト
+    empty_message = test_results.get('empty-message')
+    if empty_message and 'error' in empty_message:
+        restrictions['allow_empty_messages'] = False
+
+    return restrictions if restrictions else None