diff --git a/src/heretic/model.py b/src/heretic/model.py index d50b245..0d0644c 100644 --- a/src/heretic/model.py +++ b/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( diff --git a/tests/test_chat_template.py b/tests/test_chat_template.py new file mode 100644 index 0000000..0ba0bdb --- /dev/null +++ b/tests/test_chat_template.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: AGPL-3.0-or-later +# Copyright (C) 2025-2026 Philipp Emanuel Weidmann + 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 %}" + "{{ message['role'] }}\n{{ message['content'] }}\n" + "{% endfor %}" + "{% if add_generation_prompt %}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()