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

pipeline/processors/sft.py

@@ -7,7 +7,7 @@ from torch import Tensor
 
 from pipeline.tokenize import AutoTokenizer
 from pipeline.strategies import PromptStrategy, ChatMLStrategy
-from pipeline.processors.base import BaseProcessor, ProcessorSchema, encode_with_mask
+from pipeline.processors.base import BaseProcessor, ProcessorSchema
 from pipeline.processors.factory import ProcessorFactory
 
 
@@ -15,20 +15,20 @@ from pipeline.processors.factory import ProcessorFactory
 class SFTProcessor(BaseProcessor):
     """Supervised fine-tuning data processor.
 
-    Supports two input formats:
+    Input formats:
       1. messages (recommended):
          ``{"messages": [{"role": "user", "content": "..."},
                          {"role": "assistant", "content": "..."}]}``
-         Multi-turn and system prompts are supported.
+         Multi-turn and system prompts are supported.  Each assistant
-         The tokenizer's ``apply_chat_template`` is used for rendering.
+         turn gets ``loss_mask = 1``; all other roles get 0.
       2. legacy query/response:
          ``{"query": "...", "response": "..."}``
-         Falls back to the configured PromptStrategy (ChatML by default).
+         Internally converted to messages.
 
     Output schema:
-        - sequence: int32 tensor - Combined token IDs (prompt + response)
+        - sequence: int32 tensor - Combined token IDs
-        - loss_mask: bool tensor - True for response tokens (compute loss)
+        - loss_mask: bool tensor - True for assistant response tokens
-        - position_ids: int32 tensor - Per-sample position IDs starting from 0
+        - position_ids: int32 tensor - Per-sample position IDs, start from 0
     """
 
     def __init__(
@@ -58,7 +58,10 @@ class SFTProcessor(BaseProcessor):
         if "messages" in input_dict:
             return self._process_messages(input_dict["messages"])
         if "query" in input_dict and "response" in input_dict:
-            return self._process_legacy(input_dict)
+            return self._process_messages([
+                {"role": "user", "content": input_dict["query"]},
+                {"role": "assistant", "content": input_dict["response"]},
+            ])
         raise KeyError(
             "Input must contain 'messages' or 'query'/'response' pair"
         )
@@ -69,38 +72,19 @@ class SFTProcessor(BaseProcessor):
         if messages[-1]["role"] != "assistant":
             raise ValueError("Last message must have role 'assistant'")
 
-        last_asst_idx = max(
-            i for i, m in enumerate(messages) if m["role"] == "assistant"
-        )
-
-        prompt_tokens = self.tokenizer.apply_chat_template(
-            messages[:last_asst_idx],
-            add_generation_prompt=True,
-            tokenize=True,
-        )
-
-        resp_content = messages[last_asst_idx]["content"]
-        im_end = getattr(self.tokenizer, "im_end", "<|im_end|>")
-        resp_tokens = self.tokenizer.encode(
-            f"{resp_content}{im_end}\n", add_special_tokens=False
-        )
-
-        tokens, loss_mask = encode_with_mask(prompt_tokens, resp_tokens)
-        position_ids = torch.arange(len(tokens), dtype=torch.int32)
-        return {"sequence": tokens, "loss_mask": loss_mask, "position_ids": position_ids}
-
-    def _process_legacy(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
         strategy = self.strategy or ChatMLStrategy(self.tokenizer)
 
-        query_tokens = self.tokenizer.encode(input_dict["query"])
+        prompt, resp = strategy.format_messages(messages)
-        response_tokens = self.tokenizer.encode(input_dict["response"])
 
-        prompt = strategy.assemble_prompt(query_tokens)
+        sequence = torch.tensor(prompt + resp, dtype=torch.int32)
-        response = strategy.assemble_response(response_tokens)
+        loss_mask = torch.zeros(len(sequence), dtype=torch.bool)
-
+        loss_mask[len(prompt) :] = True
-        tokens, loss_mask = encode_with_mask(prompt, response)
+        position_ids = torch.arange(len(sequence), dtype=torch.int32)
-        position_ids = torch.arange(len(tokens), dtype=torch.int32)
+        return {
-        return {"sequence": tokens, "loss_mask": loss_mask, "position_ids": position_ids}
+            "sequence": sequence,
+            "loss_mask": loss_mask,
+            "position_ids": position_ids,
+        }
 
     @property
     def output_keys(self) -> List[str]:

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__(
+    Renders messages using the tokenizer's jinja chat_template from
-        self,
+    ``tokenizer_config.json``.  Falls back to DEFAULT_CHATML_TEMPLATE
-        tokenizer: AutoTokenizer,
+    when no template is configured.
-        user_start: str = "<｜im▁start｜>user",
+
-        user_end: str = "<｜im▁end｜>",
+    The strategy does **not** hard-code any special tokens – all
-        assistant_start: str = "<｜im▁start｜>assistant",
+    formatting is driven by the jinja template.
-        assistant_end: str = "<｜im▁end｜>",
+    """
-    ):
+
+    def __init__(self, tokenizer: AutoTokenizer):
         super().__init__(tokenizer)
-        nl_id = tokenizer.encode("a\nb", add_special_tokens=False)[1]
+        if tokenizer._chat_template is None:
-
+            tokenizer.set_chat_template(DEFAULT_CHATML_TEMPLATE)
-        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]
 
     @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 (
+        text = self.tokenizer.decode(query_tokens)
-            self._user_start_ids
+        return self.tokenizer.apply_chat_template(
-            + query_tokens
+            [{"role": "user", "content": text}],
-            + self._user_end_ids
+            add_generation_prompt=True,
-            + self._assistant_start_ids
+            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) :]

scripts/cache_h5.py

@@ -34,8 +34,8 @@ def main():
     parser.add_argument(
         "-t",
         "--tokenizer",
-        default="./tokenizer.json",
+        default="./tokenizer",
-        help="Tokenizer path (default: ./tokenizer.json)",
+        help="Tokenizer dir (default: ./tokenizer)",
     )
     parser.add_argument(
         "-s",

scripts/supervised_finetuning/sft_alpaca_gpt4.py

@@ -6,10 +6,13 @@ def process_func(input_dict: dict):
     instruction = input_dict["instruction"]
     inp = input_dict.get("input", "")
     if inp:
-        query = instruction + "\n" + inp
+        content = instruction + "\n" + inp
     else:
-        query = instruction
+        content = instruction
-    return {"query": query, "response": input_dict["output"]}
+    return {"messages": [
+        {"role": "user", "content": content},
+        {"role": "assistant", "content": input_dict["output"]},
+    ]}
 
 
 if __name__ == "__main__":

scripts/supervised_finetuning/sft_firefly-1.1m-rephrased.py

@@ -3,7 +3,10 @@ from pipeline import export_dataset
 
 
 def process_func(input_dict: dict):
-    return {"query": input_dict["instruction"], "response": input_dict["output"]}
+    return {"messages": [
+        {"role": "user", "content": input_dict["instruction"]},
+        {"role": "assistant", "content": input_dict["output"]},
+    ]}
 
 
 if __name__ == "__main__":

scripts/supervised_finetuning/sft_magicoder.py

@@ -3,7 +3,10 @@ from pipeline import export_dataset
 
 
 def process_func(input_dict: dict):
-    return {"query": input_dict["instruction"], "response": input_dict["response"]}
+    return {"messages": [
+        {"role": "user", "content": input_dict["instruction"]},
+        {"role": "assistant", "content": input_dict["response"]},
+    ]}
 
 
 if __name__ == "__main__":

scripts/supervised_finetuning/sft_metamathqa.py

@@ -3,7 +3,10 @@ from pipeline import export_dataset
 
 
 def process_func(sample: dict) -> dict:
-    return {"query": sample["query"], "response": sample["response"]}
+    return {"messages": [
+        {"role": "user", "content": sample["query"]},
+        {"role": "assistant", "content": sample["response"]},
+    ]}
 
 
 if __name__ == "__main__":

scripts/supervised_finetuning/sft_openhermes.py

@@ -2,13 +2,30 @@ from datasets import load_dataset
 from pipeline import export_dataset
 
 
+ROLE_MAP = {"system": "system", "human": "user", "gpt": "assistant"}
+
+
 def process_func(input_dict: dict):
     conversations = input_dict["conversations"]
+
+    system_msgs = []
+    idx = 0
+    if conversations and conversations[0]["from"] == "system":
+        system_msgs.append({
+            "role": "system",
+            "content": conversations[0]["value"],
+        })
+        idx = 1
+
     examples = []
-    for i in range(0, len(conversations) - 1, 2):
+    for i in range(idx, len(conversations) - 1, 2):
-        user_msg = conversations[i]["value"]
+        user_msg = conversations[i]
-        assistant_msg = conversations[i + 1]["value"]
+        assistant_msg = conversations[i + 1]
-        examples.append({"query": user_msg, "response": assistant_msg})
+        messages = system_msgs + [
+            {"role": ROLE_MAP[user_msg["from"]], "content": user_msg["value"]},
+            {"role": ROLE_MAP[assistant_msg["from"]], "content": assistant_msg["value"]},
+        ]
+        examples.append({"messages": messages})
     return examples
 
 

tests/test_io.py

@@ -140,22 +140,28 @@ class TestHDF5Handler:
 
 
 class DummyTokenizer:
-    im_end = "<|im_end|>"
+    def __init__(self):
+        self._special_token_map = {}
+        self._chat_template = None
 
     def encode(self, text: str, add_special_tokens: bool = False):
         return [ord(c) for c in text]
 
+    def decode(self, tokens, skip_special_tokens=True):
+        return "".join(chr(t) for t in tokens)
+
     def token_to_id(self, token: str):
         return ord(token)
 
-    def apply_chat_template(
+    def set_chat_template(self, template):
-        self, messages, add_generation_prompt=True, tokenize=True
+        self._chat_template = template
-    ):
+
+    def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=True):
         text = ""
         for m in messages:
-            text += f"<|im_start|>{m['role']}\n{m['content']}<|im_end|>\n"
+            text += f"<｜im▁start｜>{m['role']}\n{m['content']}<｜im▁end｜>\n"
         if add_generation_prompt:
-            text += "<|im_start|>assistant\n"
+            text += "<｜im▁start｜>assistant\n"
         return self.encode(text) if tokenize else text
 
 

tests/test_processors.py

@@ -13,22 +13,28 @@ from pipeline.processors import (
 
 
 class DummyTokenizer:
-    im_end = "<|im_end|>"
+    def __init__(self):
+        self._special_token_map = {}
+        self._chat_template = None
 
     def encode(self, text: str, add_special_tokens: bool = False):
         return [ord(c) for c in text]
 
+    def decode(self, tokens, skip_special_tokens=True):
+        return "".join(chr(t) for t in tokens)
+
     def token_to_id(self, token: str):
         return ord(token)
 
-    def apply_chat_template(
+    def set_chat_template(self, template):
-        self, messages, add_generation_prompt=True, tokenize=True
+        self._chat_template = template
-    ):
+
+    def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=True):
         text = ""
         for m in messages:
-            text += f"<|im_start|>{m['role']}\n{m['content']}<|im_end|>\n"
+            text += f"<｜im▁start｜>{m['role']}\n{m['content']}<｜im▁end｜>\n"
         if add_generation_prompt:
-            text += "<|im_start|>assistant\n"
+            text += "<｜im▁start｜>assistant\n"
         return self.encode(text) if tokenize else text
 
 

tests/test_strategies.py

@@ -10,12 +10,30 @@ from pipeline.strategies import (
 
 
 class DummyTokenizer:
+    def __init__(self):
+        self._special_token_map = {}
+        self._chat_template = None
+
     def encode(self, text: str, add_special_tokens: bool = False):
         return [ord(c) for c in text]
 
+    def decode(self, tokens, skip_special_tokens=True):
+        return "".join(chr(t) for t in tokens)
+
     def token_to_id(self, token: str):
         return ord(token)
 
+    def set_chat_template(self, template):
+        self._chat_template = template
+
+    def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=True):
+        text = ""
+        for m in messages:
+            text += f"<｜im▁start｜>{m['role']}\n{m['content']}<｜im▁end｜>\n"
+        if add_generation_prompt:
+            text += "<｜im▁start｜>assistant\n"
+        return self.encode(text) if tokenize else text
+
 
 class DummyStrategy(PromptStrategy):
     def __init__(self, tokenizer):
@@ -65,11 +83,8 @@ class TestChatMLStrategy:
         tk = DummyTokenizer()
         strategy = ChatMLStrategy(tk)
         prompt = strategy.assemble_prompt(tk.encode("hi"))
-        # prompt 末尾应该是 assistant_start 的 token ids
+        assistant_start = tk.encode("<｜im▁start｜>assistant\n")
-        assert (
+        assert prompt[-len(assistant_start):] == assistant_start
-            prompt[-len(strategy._assistant_start_ids) :]
-            == strategy._assistant_start_ids
-        )
 
 
 class TestAlpacaStrategy: