refactor: 重构流水线架构,添加Pipeline抽象并拆分IOHandler
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user