diff --git a/pipeline/processors/sft.py b/pipeline/processors/sft.py index c11464c..acad467 100644 --- a/pipeline/processors/sft.py +++ b/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. - 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]: diff --git a/pipeline/strategies/chatml.py b/pipeline/strategies/chatml.py index b4f1493..947e1f6 100644 --- a/pipeline/strategies/chatml.py +++ b/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) :] diff --git a/scripts/cache_h5.py b/scripts/cache_h5.py index 47b7ef9..7b32bb1 100644 --- a/scripts/cache_h5.py +++ b/scripts/cache_h5.py @@ -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", diff --git a/scripts/supervised_finetuning/sft_alpaca_gpt4.py b/scripts/supervised_finetuning/sft_alpaca_gpt4.py index a4ad21c..0ae2def 100644 --- a/scripts/supervised_finetuning/sft_alpaca_gpt4.py +++ b/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 - return {"query": query, "response": input_dict["output"]} + content = instruction + return {"messages": [ + {"role": "user", "content": content}, + {"role": "assistant", "content": input_dict["output"]}, + ]} if __name__ == "__main__": diff --git a/scripts/supervised_finetuning/sft_firefly-1.1m-rephrased.py b/scripts/supervised_finetuning/sft_firefly-1.1m-rephrased.py index 69b7c20..76d6e90 100644 --- a/scripts/supervised_finetuning/sft_firefly-1.1m-rephrased.py +++ b/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__": diff --git a/scripts/supervised_finetuning/sft_magicoder.py b/scripts/supervised_finetuning/sft_magicoder.py index be029e5..f6c59fe 100644 --- a/scripts/supervised_finetuning/sft_magicoder.py +++ b/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__": diff --git a/scripts/supervised_finetuning/sft_metamathqa.py b/scripts/supervised_finetuning/sft_metamathqa.py index a26ff59..0e8aec1 100644 --- a/scripts/supervised_finetuning/sft_metamathqa.py +++ b/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__": diff --git a/scripts/supervised_finetuning/sft_openhermes.py b/scripts/supervised_finetuning/sft_openhermes.py index b84d79c..c1f40ae 100644 --- a/scripts/supervised_finetuning/sft_openhermes.py +++ b/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): - 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 diff --git a/tests/test_io.py b/tests/test_io.py index b07099d..9368d23 100644 --- a/tests/test_io.py +++ b/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( - 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 diff --git a/tests/test_processors.py b/tests/test_processors.py index bec2400..41c1ff3 100644 --- a/tests/test_processors.py +++ b/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( - 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 diff --git a/tests/test_strategies.py b/tests/test_strategies.py index 2447024..809a7ee 100644 --- a/tests/test_strategies.py +++ b/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 - 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: