ViperEkura/DataPipeline
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
ViperEkura
2026-07-04 14:32:35 +08:00
parent 06735b9cb3
commit 816c02dab0
11 changed files with 177 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pipeline/processors/sft.py
+22 -38
View File
@@ -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.
The tokenizer's ``apply_chat_template`` is used for rendering.
Multi-turn and system prompts are supported. Each assistant
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)
- loss_mask: bool tensor - True for response tokens (compute loss)
- position_ids: int32 tensor - Per-sample position IDs starting from 0
- sequence: int32 tensor - Combined token IDs
- loss_mask: bool tensor - True for assistant response tokens
- 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"])
response_tokens = self.tokenizer.encode(input_dict["response"])
prompt, resp = strategy.format_messages(messages)
prompt = strategy.assemble_prompt(query_tokens)
response = strategy.assemble_response(response_tokens)
tokens, loss_mask = encode_with_mask(prompt, response)
position_ids = torch.arange(len(tokens), dtype=torch.int32)
return {"sequence": tokens, "loss_mask": loss_mask, "position_ids": position_ids}
sequence = torch.tensor(prompt + resp, dtype=torch.int32)
loss_mask = torch.zeros(len(sequence), dtype=torch.bool)
loss_mask[len(prompt) :] = True
position_ids = torch.arange(len(sequence), dtype=torch.int32)
return {
"sequence": sequence,
"loss_mask": loss_mask,
"position_ids": position_ids,
}
@property
def output_keys(self) -> List[str]:
pipeline/strategies/chatml.py
+70 -22
View File
@@ -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) :]
scripts/cache_h5.py
+2 -2
View File
@@ -34,8 +34,8 @@ def main():
parser.add_argument(
"-t",
"--tokenizer",
default="./tokenizer.json",
help="Tokenizer path (default: ./tokenizer.json)",
default="./tokenizer",
help="Tokenizer dir (default: ./tokenizer)",
)
parser.add_argument(
"-s",
scripts/supervised_finetuning/sft_alpaca_gpt4.py
+6 -3
View File
@@ -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
return {"query": query, "response": input_dict["output"]}
content = instruction
return {"messages": [
{"role": "user", "content": content},
{"role": "assistant", "content": input_dict["output"]},
]}
if __name__ == "__main__":
scripts/supervised_finetuning/sft_firefly-1.1m-rephrased.py
+4 -1
View File
@@ -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
+4 -1
View File
@@ -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
+4 -1
View File
@@ -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
+21 -4
View File
@@ -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):
user_msg = conversations[i]["value"]
assistant_msg = conversations[i + 1]["value"]
examples.append({"query": user_msg, "response": assistant_msg})
for i in range(idx, len(conversations) - 1, 2):
user_msg = conversations[i]
assistant_msg = conversations[i + 1]
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
+12 -6
View File
@@ -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(
self, messages, add_generation_prompt=True, tokenize=True
):
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"
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
+12 -6
View File
@@ -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(
self, messages, add_generation_prompt=True, tokenize=True
):
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"
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
+20 -5
View File
@@ -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
assert (
prompt[-len(strategy._assistant_start_ids) :]
== strategy._assistant_start_ids
)
assistant_start = tk.encode("<｜im▁start｜>assistant\n")
assert prompt[-len(assistant_start):] == assistant_start
class TestAlpacaStrategy: