fix: support models whose chat template rejects the system role (#498)

Models like Gemma 2 and its fine-tunes reject the system role in their
chat template. generate() always prepends a system message, so the
template raises 'System role not supported' and the dtype loop gives up,
making these models impossible to load.

Detect once (right after the tokenizer loads) whether the template
accepts a system role, by rendering a probe chat and checking a sentinel
system message survives in the output (this catches both templates that
raise and templates that silently drop the role). When the role isn't
supported, fold every system message into the following user message
(dropping empty system prompts), so a system role anywhere in the chat
is removed before rendering, not just a leading one. Templates that
support the system role are unaffected, so their rendered chats and
resulting weights are unchanged.

Supersedes #483.

Co-authored-by: jayhemnani9910 <jayhemnani9910@users.noreply.github.com>

src/heretic/model.py

@@ -51,6 +51,30 @@ def strftime_fixed_date(format: str) -> str:
     return datetime(2026, 1, 1, tzinfo=timezone.utc).strftime(format)
 
 
+def chat_template_supports_system_role(tokenizer: PreTrainedTokenizerBase) -> bool:
+    # Some chat templates (e.g. Gemma 2 and its fine-tunes) reject the system role
+    # outright, raising when one is present; others silently drop it. The only
+    # reliable way to tell is to render a probe chat and check whether a sentinel
+    # system message survives into the output. This is also how llama.cpp's template
+    # engine (minja) detects system-role support.
+    sentinel = "HERETIC_SYSTEM_ROLE_PROBE"
+    try:
+        rendered = cast(
+            str,
+            tokenizer.apply_chat_template(
+                [
+                    {"role": "system", "content": sentinel},
+                    {"role": "user", "content": "Hi"},
+                ],
+                add_generation_prompt=True,
+                tokenize=False,
+            ),
+        )
+    except Exception:
+        return False
+    return sentinel in rendered
+
+
 # The list contains one element per layer.
 # Each element maps from the component name to a (possibly sparse) mapping
 # from the module index to an (input, output) tuple containing the I/O
@@ -65,6 +89,8 @@ class Model:
     processor: ProcessorMixin | None
     peft_config: LoraConfig
     dtype: torch.dtype
+    # Whether the model's chat template accepts a system role (see __init__).
+    supports_system_role: bool
 
     def __init__(self, settings: Settings):
         self.settings = settings
@@ -104,6 +130,11 @@ class Model:
         #           after the prompt and thinks the sequence is complete.
         self.tokenizer.padding_side = "left"
 
+        # Detect once whether this tokenizer's chat template accepts a system
+        # role. Templates that don't (e.g. Gemma 2) need the system prompt folded
+        # into the first user message instead (see _apply_system_role).
+        self.supports_system_role = chat_template_supports_system_role(self.tokenizer)
+
         self.model = None
         self.max_memory = (
             {int(k) if k.isdigit() else k: v for k, v in settings.max_memory.items()}
@@ -460,16 +491,46 @@ class Model:
 
         return sorted(components)
 
+    def _apply_system_role(
+        self, messages: list[dict[str, str]]
+    ) -> list[dict[str, str]]:
+        # Templates that support the system role get the messages unchanged, so the
+        # rendered chat (and therefore the resulting weights) is byte-for-byte
+        # identical to before. For templates that reject it (e.g. Gemma 2), fold each
+        # system message into the next user message (what Gemma's own later templates
+        # do), so a system role anywhere in the chat is removed before rendering, not
+        # just a leading one. Empty system prompts are dropped; a system message with
+        # no following user message becomes a user turn.
+        if self.supports_system_role:
+            return messages
+        result: list[dict[str, str]] = []
+        pending: list[str] = []
+        for message in messages:
+            if message["role"] == "system":
+                if message["content"]:
+                    pending.append(message["content"])
+                continue
+            if pending and message["role"] == "user":
+                message = dict(message)
+                message["content"] = "\n\n".join([*pending, message["content"]])
+                pending = []
+            result.append(message)
+        if pending:
+            result.append({"role": "user", "content": "\n\n".join(pending)})
+        return result
+
     def generate(
         self,
         prompts: list[Prompt],
         **kwargs: Any,
     ) -> tuple[BatchEncoding, GenerateDecoderOnlyOutput | LongTensor]:
         chats = [
-            [
-                {"role": "system", "content": prompt.system},
-                {"role": "user", "content": prompt.user},
-            ]
+            self._apply_system_role(
+                [
+                    {"role": "system", "content": prompt.system},
+                    {"role": "user", "content": prompt.user},
+                ]
+            )
             for prompt in prompts
         ]
 
@@ -814,6 +875,7 @@ class Model:
         return torch.cat(logits, dim=0)
 
     def stream_chat_response(self, chat: list[dict[str, str]]) -> str:
+        chat = self._apply_system_role(chat)
         # This cast is valid because str is the return type
         # for single-chat operation with tokenize=False.
         chat_prompt = cast(

tests/test_chat_template.py

@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: AGPL-3.0-or-later
+# Copyright (C) 2025-2026  Philipp Emanuel Weidmann <pew@worldwidemann.com> + contributors
+
+import unittest
+
+from tokenizers import Tokenizer
+from tokenizers.models import WordLevel
+from tokenizers.pre_tokenizers import Whitespace
+from transformers import PreTrainedTokenizerFast
+
+from heretic.model import Model, chat_template_supports_system_role
+
+# A template that rejects the system role the way Gemma 2's does.
+GEMMA_STYLE_TEMPLATE = (
+    "{% for message in messages %}"
+    "{% if message['role'] == 'system' %}"
+    "{{ raise_exception('System role not supported') }}"
+    "{% endif %}"
+    "<start_of_turn>{{ message['role'] }}\n{{ message['content'] }}<end_of_turn>\n"
+    "{% endfor %}"
+    "{% if add_generation_prompt %}<start_of_turn>model\n{% endif %}"
+)
+
+# A template that renders the system role normally.
+SYSTEM_TEMPLATE = (
+    "{% for message in messages %}"
+    "<|{{ message['role'] }}|>\n{{ message['content'] }}\n"
+    "{% endfor %}"
+    "{% if add_generation_prompt %}<|assistant|>\n{% endif %}"
+)
+
+
+def _tokenizer(template: str) -> PreTrainedTokenizerFast:
+    backend = Tokenizer(WordLevel(vocab={"[UNK]": 0}, unk_token="[UNK]"))
+    backend.pre_tokenizer = Whitespace()
+    return PreTrainedTokenizerFast(
+        tokenizer_object=backend,
+        unk_token="[UNK]",
+        chat_template=template,
+    )
+
+
+class ChatTemplateSupportTests(unittest.TestCase):
+    def test_detects_unsupported_system_role(self) -> None:
+        self.assertFalse(
+            chat_template_supports_system_role(_tokenizer(GEMMA_STYLE_TEMPLATE))
+        )
+
+    def test_detects_supported_system_role(self) -> None:
+        self.assertTrue(chat_template_supports_system_role(_tokenizer(SYSTEM_TEMPLATE)))
+
+
+class ApplySystemRoleTests(unittest.TestCase):
+    def _model(self, supports: bool) -> Model:
+        model = Model.__new__(Model)
+        model.supports_system_role = supports
+        return model
+
+    def test_passthrough_when_supported(self) -> None:
+        messages = [
+            {"role": "system", "content": "You are helpful."},
+            {"role": "user", "content": "hi"},
+        ]
+        # Same object contents, so the rendered chat is byte-identical to before.
+        self.assertEqual(self._model(True)._apply_system_role(messages), messages)
+
+    def test_folds_into_first_user_message(self) -> None:
+        result = self._model(False)._apply_system_role(
+            [
+                {"role": "system", "content": "You are helpful."},
+                {"role": "user", "content": "hi"},
+            ]
+        )
+        self.assertEqual(
+            result, [{"role": "user", "content": "You are helpful.\n\nhi"}]
+        )
+
+    def test_drops_empty_system_prompt(self) -> None:
+        result = self._model(False)._apply_system_role(
+            [
+                {"role": "system", "content": ""},
+                {"role": "user", "content": "hi"},
+            ]
+        )
+        self.assertEqual(result, [{"role": "user", "content": "hi"}])
+
+    def test_no_system_message_is_unchanged(self) -> None:
+        messages = [{"role": "user", "content": "hi"}]
+        self.assertEqual(self._model(False)._apply_system_role(messages), messages)
+
+    def test_folds_system_role_appearing_later(self) -> None:
+        # A system role anywhere (not just leading) must be removed for templates
+        # that reject it.
+        result = self._model(False)._apply_system_role(
+            [
+                {"role": "user", "content": "hi"},
+                {"role": "assistant", "content": "hello"},
+                {"role": "system", "content": "Be terse."},
+                {"role": "user", "content": "and now?"},
+            ]
+        )
+        self.assertFalse(any(m["role"] == "system" for m in result))
+        self.assertEqual(
+            result,
+            [
+                {"role": "user", "content": "hi"},
+                {"role": "assistant", "content": "hello"},
+                {"role": "user", "content": "Be terse.\n\nand now?"},
+            ],
+        )
+
+    def test_folded_later_system_role_renders_without_raising(self) -> None:
+        # End to end: a later system role would make the Gemma-style template raise;
+        # after folding it must render fine.
+        model = self._model(False)
+        tok = _tokenizer(GEMMA_STYLE_TEMPLATE)
+        chat = model._apply_system_role(
+            [
+                {"role": "user", "content": "hi"},
+                {"role": "system", "content": "Be terse."},
+                {"role": "user", "content": "and now?"},
+            ]
+        )
+        rendered = tok.apply_chat_template(
+            chat, add_generation_prompt=True, tokenize=False
+        )
+        self.assertIn("Be terse.", rendered)
+
+
+if __name__ == "__main__":
+    unittest.main()