refactor: SFT 统一 messages 格式 + ChatML 纯 jinja 渲染

pipeline/strategies/chatml.py

@@ -1,43 +1,91 @@
 """ChatML format strategy."""
 
-from typing import List
+from typing import Dict, List, Tuple
 
 from pipeline.tokenize import AutoTokenizer
 from pipeline.strategies.base import PromptStrategy
 from pipeline.strategies.factory import StrategyFactory
 
+DEFAULT_CHATML_TEMPLATE = (
+    "{% for message in messages %}"
+    "{% if message['role'] == 'system' %}"
+    "{{ '<|im_start|>system\n' + message['content'] + '<|im_end|>\n' }}"
+    "{% elif message['role'] == 'user' %}"
+    "{{ '<|im_start|>user\n' + message['content'] + '<|im_end|>\n' }}"
+    "{% elif message['role'] == 'assistant' %}"
+    "{{ '<|im_start|>assistant\n' + message['content'] + '<|im_end|>\n' }}"
+    "{% endif %}"
+    "{% endfor %}"
+    "{% if add_generation_prompt %}"
+    "{{ '<|im_start|>assistant\n' }}"
+    "{% endif %}"
+)
+
 
 @StrategyFactory.register("chatml")
 class ChatMLStrategy(PromptStrategy):
-    """ChatML format strategy."""
+    """ChatML format strategy.
 
-    def __init__(
-        self,
-        tokenizer: AutoTokenizer,
-        user_start: str = "<｜im▁start｜>user",
-        user_end: str = "<｜im▁end｜>",
-        assistant_start: str = "<｜im▁start｜>assistant",
-        assistant_end: str = "<｜im▁end｜>",
-    ):
+    Renders messages using the tokenizer's jinja chat_template from
+    ``tokenizer_config.json``.  Falls back to DEFAULT_CHATML_TEMPLATE
+    when no template is configured.
+
+    The strategy does **not** hard-code any special tokens – all
+    formatting is driven by the jinja template.
+    """
+
+    def __init__(self, tokenizer: AutoTokenizer):
         super().__init__(tokenizer)
-        nl_id = tokenizer.encode("a\nb", add_special_tokens=False)[1]
-
-        self._user_start_ids = self._encode_format(user_start) + [nl_id]
-        self._user_end_ids = self._encode_format(user_end) + [nl_id]
-        self._assistant_start_ids = self._encode_format(assistant_start) + [nl_id]
-        self._assistant_end_ids = self._encode_format(assistant_end) + [nl_id]
+        if tokenizer._chat_template is None:
+            tokenizer.set_chat_template(DEFAULT_CHATML_TEMPLATE)
 
     @property
     def name(self) -> str:
         return "chatml"
 
+    def format_messages(
+        self,
+        messages: List[Dict[str, str]],
+    ) -> Tuple[List[int], List[int]]:
+        """Render a single-turn messages conversation.
+
+        Returns ``(prompt_tokens, response_tokens)`` where
+        *prompt_tokens* contains everything up to (and including) the
+        last assistant start marker, and *response_tokens* is the
+        assistant content plus the closing markers.
+        """
+        last_asst = max(
+            i for i, m in enumerate(messages) if m["role"] == "assistant"
+        )
+
+        prompt = self.tokenizer.apply_chat_template(
+            messages[:last_asst],
+            add_generation_prompt=True,
+            tokenize=True,
+        )
+        full = self.tokenizer.apply_chat_template(
+            messages[: last_asst + 1],
+            add_generation_prompt=False,
+            tokenize=True,
+        )
+        return prompt, full[len(prompt) :]
+
     def assemble_prompt(self, query_tokens: List[int]) -> List[int]:
-        return (
-            self._user_start_ids
-            + query_tokens
-            + self._user_end_ids
-            + self._assistant_start_ids
+        text = self.tokenizer.decode(query_tokens)
+        return self.tokenizer.apply_chat_template(
+            [{"role": "user", "content": text}],
+            add_generation_prompt=True,
+            tokenize=True,
         )
 
     def assemble_response(self, response_tokens: List[int]) -> List[int]:
-        return response_tokens + self._assistant_end_ids
+        text = self.tokenizer.decode(response_tokens)
+        full = self.tokenizer.apply_chat_template(
+            [{"role": "assistant", "content": text}],
+            add_generation_prompt=False,
+            tokenize=True,
+        )
+        opening = self.tokenizer.apply_chat_template(
+            [], add_generation_prompt=True, tokenize=True
+        )
+        return full[len(opening) :]