| """Regression tests for FrozenModelBackend.format_prompt's chat-template handling.""" |
|
|
| import unittest |
|
|
| from owmi.backends import FrozenModelBackend |
| from owmi.types import ExperimentConfig |
|
|
|
|
| def _make_config(system_prompt="You are a careful introspection assistant."): |
| return ExperimentConfig(model_name="fake/model", system_prompt=system_prompt) |
|
|
|
|
| class _AcceptsSystemRoleTokenizer: |
| """Mimics a normal chat template that accepts a dedicated system turn.""" |
|
|
| def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=True): |
| assert not tokenize |
| assert add_generation_prompt |
| rendered = [] |
| for m in messages: |
| rendered.append(f"<{m['role']}>{m['content']}") |
| return "\n".join(rendered) + "\n<assistant>" |
|
|
|
|
| class _RejectsSystemRoleTokenizer: |
| """Mimics Gemma-2's chat template, which raises when given a system turn.""" |
|
|
| def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=True): |
| roles = [m["role"] for m in messages] |
| if "system" in roles: |
| raise Exception("System role not supported") |
| rendered = [] |
| for m in messages: |
| rendered.append(f"<{m['role']}>{m['content']}") |
| return "\n".join(rendered) + "\n<assistant>" |
|
|
|
|
| class _RaisesUnrelatedErrorTokenizer: |
| """A tokenizer whose apply_chat_template fails for an unrelated reason.""" |
|
|
| def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=True): |
| raise Exception("tokenizer is not fully initialized") |
|
|
|
|
| def _backend_with_tokenizer(tokenizer, system_prompt="You are a careful introspection assistant."): |
| backend = FrozenModelBackend.__new__(FrozenModelBackend) |
| backend.tokenizer = tokenizer |
| backend.config = _make_config(system_prompt) |
| return backend |
|
|
|
|
| class FormatPromptSystemRoleTests(unittest.TestCase): |
| def test_uses_system_role_when_supported(self): |
| backend = _backend_with_tokenizer(_AcceptsSystemRoleTokenizer()) |
| out = backend.format_prompt("What happened?") |
| self.assertIn("<system>You are a careful introspection assistant.", out) |
| self.assertIn("<user>What happened?", out) |
|
|
| def test_falls_back_to_merged_user_turn_when_system_role_rejected(self): |
| backend = _backend_with_tokenizer(_RejectsSystemRoleTokenizer()) |
| out = backend.format_prompt("What happened?") |
| |
| self.assertNotIn("<system>", out) |
| self.assertIn("You are a careful introspection assistant.", out) |
| self.assertIn("What happened?", out) |
| self.assertIn("<user>", out) |
|
|
| def test_unrelated_tokenizer_errors_still_raise(self): |
| backend = _backend_with_tokenizer(_RaisesUnrelatedErrorTokenizer()) |
| with self.assertRaises(Exception) as ctx: |
| backend.format_prompt("What happened?") |
| self.assertIn("not fully initialized", str(ctx.exception)) |
|
|
| def test_no_tokenizer_falls_back_to_plain_format(self): |
| backend = _backend_with_tokenizer(None) |
| out = backend.format_prompt("What happened?") |
| self.assertEqual( |
| out, |
| "System: You are a careful introspection assistant.\n" |
| "User: What happened?\nAssistant:", |
| ) |
|
|
