From aa2ea4f3a6ff7537d67bb59048ca89683344804c Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Thu, 4 Jun 2026 14:01:05 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20SFT=20=E5=A2=9E=E5=8A=A0=20position=5Fi?= =?UTF-8?q?ds=20=E8=BE=B9=E7=95=8C=E5=A4=84=E7=90=86?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - SFTProcessor 输出 per-sample position_ids（torch.arange），每个样本从 0 开始 - position_ids 与 sequence/loss_mask 一同打包，边界处自然重置 - PT 路径不生成 position_ids - 新增 TestPositionIds 测试及 SFTProcessor 相关测试 --- pipeline/io/export.py | 1 - pipeline/processors/sft.py | 72 ++++++++++++++++++++++++----- tests/test_io.py | 47 ++++++++++++++++++- tests/test_processors.py | 93 +++++++++++++++++++++++++++++++++++--- 4 files changed, 193 insertions(+), 20 deletions(-) diff --git a/pipeline/io/export.py b/pipeline/io/export.py index b45e0ee..2061f42 100644 --- a/pipeline/io/export.py +++ b/pipeline/io/export.py @@ -5,7 +5,6 @@ import logging import os from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Union - from datasets import Dataset from tqdm import tqdm diff --git a/pipeline/processors/sft.py b/pipeline/processors/sft.py index ee21166..c11464c 100644 --- a/pipeline/processors/sft.py +++ b/pipeline/processors/sft.py @@ -15,15 +15,20 @@ from pipeline.processors.factory import ProcessorFactory class SFTProcessor(BaseProcessor): """Supervised fine-tuning data processor. - Processes query-response pairs into tokenized sequences with loss masks. - - Input schema: - - query: str - User query/prompt - - response: str - Assistant response + Supports two 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. + 2. legacy query/response: + ``{"query": "...", "response": "..."}`` + Falls back to the configured PromptStrategy (ChatML by default). Output schema: - - sequence: int32 tensor - Combined token IDs (query + response) + - 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 """ def __init__( @@ -32,28 +37,71 @@ class SFTProcessor(BaseProcessor): strategy: Optional[PromptStrategy] = None, ): self.tokenizer = tokenizer - self.strategy = strategy or ChatMLStrategy(tokenizer) + self.strategy = strategy @property def schema(self) -> ProcessorSchema: return ProcessorSchema( - input_fields={"query": str, "response": str}, + input_fields={ + "messages": list, + "query": str, + "response": str, + }, output_fields={ "sequence": torch.int32, "loss_mask": torch.bool, + "position_ids": torch.int32, }, ) def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]: + 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) + raise KeyError( + "Input must contain 'messages' or 'query'/'response' pair" + ) + + def _process_messages(self, messages: List[Dict[str, str]]) -> Dict[str, Tensor]: + if not messages: + raise ValueError("Messages list is empty") + 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 = self.strategy.assemble_prompt(query_tokens) - response = self.strategy.assemble_response(response_tokens) + prompt = strategy.assemble_prompt(query_tokens) + response = strategy.assemble_response(response_tokens) tokens, loss_mask = encode_with_mask(prompt, response) - return {"sequence": tokens, "loss_mask": loss_mask} + position_ids = torch.arange(len(tokens), dtype=torch.int32) + return {"sequence": tokens, "loss_mask": loss_mask, "position_ids": position_ids} @property def output_keys(self) -> List[str]: - return ["sequence", "loss_mask"] + return ["sequence", "loss_mask", "position_ids"] diff --git a/tests/test_io.py b/tests/test_io.py index a22ac0e..a496722 100644 --- a/tests/test_io.py +++ b/tests/test_io.py @@ -7,7 +7,7 @@ import torch import h5py from pathlib import Path -from pipeline.io import FileScanner, HDF5Handler +from pipeline.io import FileScanner, HDF5Handler, cache_jsonl class TestFileScanner: @@ -137,3 +137,48 @@ class TestHDF5Handler: h5_path = os.path.join(tmpdir, "meta.h5") metadata = HDF5Handler.get_metadata(h5_path) assert metadata["data"] == 5 + + +class DummyTokenizer: + im_end = "<|im_end|>" + + def encode(self, text: str, add_special_tokens: bool = False): + return [ord(c) for c in text] + + 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 TestPositionIds: + def test_example_specific_position_ids(self): + with tempfile.TemporaryDirectory() as tmpdir: + jsonl_path = os.path.join(tmpdir, "data.jsonl") + with open(jsonl_path, "w") as f: + f.write('{"messages": [{"role": "user", "content": "a"}, {"role": "assistant", "content": "bc"}]}\n') + f.write('{"messages": [{"role": "user", "content": "def"}, {"role": "assistant", "content": "g"}]}\n') + + from pipeline.processors import SFTProcessor + + processor = SFTProcessor(DummyTokenizer()) + out_dir = os.path.join(tmpdir, "cached") + cache_jsonl([jsonl_path], out_dir, processor, pack_size=-1) + + h5_path = os.path.join(out_dir, "data.h5") + loaded = HDF5Handler.load(h5_path, share_memory=False) + + assert "position_ids" in loaded + assert len(loaded["position_ids"]) == 2 + assert len(loaded["position_ids"]) == len(loaded["sequence"]) + + for i, pos in enumerate(loaded["position_ids"]): + seq_len = len(loaded["sequence"][i]) + assert len(pos) == seq_len + assert pos[0].item() == 0 + assert (pos == torch.arange(seq_len, dtype=torch.int32)).all() diff --git a/tests/test_processors.py b/tests/test_processors.py index e9d6dab..e10ba3d 100644 --- a/tests/test_processors.py +++ b/tests/test_processors.py @@ -13,9 +13,21 @@ from pipeline.processors import ( class DummyTokenizer: + im_end = "<|im_end|>" + def encode(self, text: str, add_special_tokens: bool = False): return [ord(c) for c in text] + 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 TestBaseProcessor: def test_abstract_class_cannot_be_instantiated(self): @@ -41,16 +53,18 @@ class TestPreTrainProcessor: class TestSFTProcessor: def test_output_keys(self): - assert SFTProcessor(DummyTokenizer()).output_keys == ["sequence", "loss_mask"] + keys = SFTProcessor(DummyTokenizer()).output_keys + assert "sequence" in keys + assert "loss_mask" in keys + assert "position_ids" in keys - def test_process_returns_both_keys(self): + def test_process_returns_all_keys(self): result = SFTProcessor(DummyTokenizer()).process( {"query": "hello", "response": "world"} ) - assert "sequence" in result - assert "loss_mask" in result - assert isinstance(result["sequence"], torch.Tensor) - assert isinstance(result["loss_mask"], torch.Tensor) + for key in ["sequence", "loss_mask", "position_ids"]: + assert key in result + assert isinstance(result[key], torch.Tensor) def test_loss_mask_correct_length(self): result = SFTProcessor(DummyTokenizer()).process( @@ -64,6 +78,73 @@ class TestSFTProcessor: ) assert result["loss_mask"].dtype == torch.bool + def test_position_ids_start_from_zero(self): + result = SFTProcessor(DummyTokenizer()).process( + {"query": "abc", "response": "de"} + ) + seq_len = len(result["sequence"]) + expected = torch.arange(seq_len, dtype=torch.int32) + assert torch.equal(result["position_ids"], expected) + + def test_messages_single_turn(self): + result = SFTProcessor(DummyTokenizer()).process({ + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "bye"}, + ] + }) + for key in ["sequence", "loss_mask", "position_ids"]: + assert key in result + assert len(result["sequence"]) == len(result["loss_mask"]) + + def test_messages_loss_on_last_assistant_only(self): + result = SFTProcessor(DummyTokenizer()).process({ + "messages": [ + {"role": "user", "content": "q1"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "q2"}, + {"role": "assistant", "content": "a2"}, + ] + }) + mask = result["loss_mask"] + first_true = mask.tolist().index(True) + assert not mask[:first_true].any() + assert mask[-1].item() is True + + def test_messages_with_system_prompt(self): + result = SFTProcessor(DummyTokenizer()).process({ + "messages": [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] + }) + assert "sequence" in result + + def test_messages_empty_raises(self): + with pytest.raises(ValueError, match="Messages list is empty"): + SFTProcessor(DummyTokenizer()).process({"messages": []}) + + def test_messages_last_not_assistant_raises(self): + with pytest.raises(ValueError, match="Last message must"): + SFTProcessor(DummyTokenizer()).process({ + "messages": [{"role": "user", "content": "hi"}] + }) + + def test_missing_fields_raises(self): + with pytest.raises(KeyError): + SFTProcessor(DummyTokenizer()).process({"foo": "bar"}) + + def test_position_ids_start_from_zero(self): + result = SFTProcessor(DummyTokenizer()).process( + {"query": "hi", "response": "ok"} + ) + pos_ids = result["position_ids"] + assert pos_ids.dtype == torch.int32 + assert len(pos_ids) == len(result["sequence"]) + assert pos_ids[0].item() == 0 + assert (pos_ids == torch.arange(len(pos_ids))).all() + class TestDPOProcessor: def test_output_keys(self):