refactor: 重构流水线架构,添加Pipeline抽象并拆分IOHandler

This commit is contained in:
2026-04-23 19:45:57 +08:00
parent a38334f4ce
commit cb6bfcb976
14 changed files with 886 additions and 184 deletions
+22 -4
View File
@@ -1,13 +1,13 @@
"""Supervised fine-tuning data processor."""
from typing import Dict, List, Any, Optional
from typing import Any, Dict, List, Optional
import torch
from torch import Tensor
from pipeline.tokenize import AutoTokenizer
from pipeline.strategies import PromptStrategy, ChatMLStrategy
from pipeline.processors.base import BaseProcessor, _encode_with_mask
from pipeline.processors.base import BaseProcessor, ProcessorSchema, encode_with_mask
from pipeline.processors.factory import ProcessorFactory
@@ -15,7 +15,15 @@ from pipeline.processors.factory import ProcessorFactory
class SFTProcessor(BaseProcessor):
"""Supervised fine-tuning data processor.
Supports custom prompt strategy via constructor parameter.
Processes query-response pairs into tokenized sequences with loss masks.
Input schema:
- query: str - User query/prompt
- response: str - Assistant response
Output schema:
- sequence: int32 tensor - Combined token IDs (query + response)
- loss_mask: bool tensor - True for response tokens (compute loss)
"""
def __init__(
@@ -26,6 +34,16 @@ class SFTProcessor(BaseProcessor):
self.tokenizer = tokenizer
self.strategy = strategy or ChatMLStrategy(tokenizer)
@property
def schema(self) -> ProcessorSchema:
return ProcessorSchema(
input_fields={"query": str, "response": str},
output_fields={
"sequence": torch.int32,
"loss_mask": torch.bool,
},
)
def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
query_tokens = self.tokenizer.encode(input_dict["query"])
response_tokens = self.tokenizer.encode(input_dict["response"])
@@ -33,7 +51,7 @@ class SFTProcessor(BaseProcessor):
prompt = self.strategy.assemble_prompt(query_tokens)
response = self.strategy.assemble_response(response_tokens)
tokens, loss_mask = _encode_with_mask(prompt, response)
tokens, loss_mask = encode_with_mask(prompt, response)
return {"sequence": tokens, "loss_mask": loss_mask}
@property