From 68e9f6e931a2292dcd65f26627c203ee5a3150f0 Mon Sep 17 00:00:00 2001 From: Moustafa Saleh Date: Sat, 23 May 2026 17:09:39 -0700 Subject: [PATCH] fix: enter thinking block for Kimi generation prompts --- tests/test_kimi_k25_parser.py | 8 ++++++++ torchspec/data/parse.py | 2 +- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/test_kimi_k25_parser.py b/tests/test_kimi_k25_parser.py index 93ebedb6..842cb33d 100644 --- a/tests/test_kimi_k25_parser.py +++ b/tests/test_kimi_k25_parser.py @@ -138,6 +138,14 @@ def test_single_turn_conversation(self, mock_tokenizer, kimi_template): assert len(input_ids) == len(loss_mask) assert loss_mask.sum() > 0 + def test_generation_prompt_enters_thinking_block(self, mock_tokenizer, kimi_template): + parser = KimiK25Parser(mock_tokenizer, kimi_template) + conversation = [{"role": "user", "content": "Hello!"}] + + formatted = parser.format(conversation, add_generation_prompt=True) + + assert formatted.endswith("<|im_assistant|>assistant<|im_middle|>") + def test_multi_turn_conversation(self, mock_tokenizer, kimi_template): parser = KimiK25Parser(mock_tokenizer, kimi_template) conversation = [ diff --git a/torchspec/data/parse.py b/torchspec/data/parse.py index c56cc2cc..8f2a9868 100644 --- a/torchspec/data/parse.py +++ b/torchspec/data/parse.py @@ -478,7 +478,7 @@ def format(self, conversation: "Conversation", **kwargs) -> str: parts.append(f"{self.TOOL_HEADER}{content}{self.END_TOKEN}") if add_generation_prompt: - parts.append(self.ASSISTANT_HEADER) + parts.append(f"{self.ASSISTANT_HEADER}") return "".join(parts)