111 lines
3.1 KiB
Python
111 lines
3.1 KiB
Python
from abc import ABC, abstractmethod
|
|
from typing import Dict, List, Callable
|
|
import torch
|
|
from torch import Tensor
|
|
|
|
|
|
class BaseProcessor(ABC):
|
|
"""处理器抽象基类 - 策略模式"""
|
|
|
|
@abstractmethod
|
|
def process(self, input_dict: dict) -> dict:
|
|
"""处理单个数据项"""
|
|
pass
|
|
|
|
@property
|
|
@abstractmethod
|
|
def output_keys(self) -> List[str]:
|
|
"""输出字段列表"""
|
|
pass
|
|
|
|
|
|
class PreTrainProcessor(BaseProcessor):
|
|
"""预训练数据处理器"""
|
|
|
|
def __init__(self, tokenizer):
|
|
self.tokenizer = tokenizer
|
|
|
|
def process(self, input_dict: dict) -> dict:
|
|
segment = input_dict["text"]
|
|
tokens = self.tokenizer.encode(f"{segment}<eos>")
|
|
return {'sequence': torch.tensor(tokens, dtype=torch.int32)}
|
|
|
|
@property
|
|
def output_keys(self) -> List[str]:
|
|
return ["sequence"]
|
|
|
|
|
|
class SFTProcessor(BaseProcessor):
|
|
"""监督微调数据处理器"""
|
|
|
|
def __init__(self, tokenizer):
|
|
self.tokenizer = tokenizer
|
|
|
|
def process(self, input_dict: dict) -> dict:
|
|
query, response = input_dict["query"], input_dict["response"]
|
|
q = self.tokenizer.encode(
|
|
f"<|im_start|>user\n{query}<|im_end|>\n<|im_start|>assistant\n"
|
|
)
|
|
a = self.tokenizer.encode(f"{response}<|im_end|>\n<eos>")
|
|
|
|
tokens = torch.tensor(q + a, dtype=torch.int32)
|
|
loss_mask = torch.zeros_like(tokens, dtype=torch.bool)
|
|
loss_mask[len(q):] = True
|
|
|
|
return {"sequence": tokens, "loss_mask": loss_mask}
|
|
|
|
@property
|
|
def output_keys(self) -> List[str]:
|
|
return ["sequence", "loss_mask"]
|
|
|
|
|
|
class DPOProcessor(BaseProcessor):
|
|
"""DPO偏好学习数据处理器"""
|
|
|
|
def __init__(self, tokenizer):
|
|
self.tokenizer = tokenizer
|
|
|
|
def process(self, input_dict: dict) -> dict:
|
|
# TODO: 实现DPO处理逻辑
|
|
return None
|
|
|
|
@property
|
|
def output_keys(self) -> List[str]:
|
|
return ["chosen", "chosen_mask", "rejected", "rejected_mask"]
|
|
|
|
|
|
class ProcessorFactory:
|
|
"""处理器工厂 - 工厂模式"""
|
|
|
|
_processors = {
|
|
"pt": PreTrainProcessor,
|
|
"sft": SFTProcessor,
|
|
"dpo": DPOProcessor,
|
|
}
|
|
|
|
@classmethod
|
|
def create(cls, processor_type: str, tokenizer) -> BaseProcessor:
|
|
"""创建处理器实例"""
|
|
if processor_type not in cls._processors:
|
|
raise ValueError(f"Invalid processor type: {processor_type}")
|
|
return cls._processors[processor_type](tokenizer)
|
|
|
|
@classmethod
|
|
def register(cls, processor_type: str, processor_class: type):
|
|
"""注册新的处理器类型"""
|
|
cls._processors[processor_type] = processor_class
|
|
|
|
|
|
# 向后兼容的函数接口
|
|
def get_pt_processor(tokenizer) -> Callable:
|
|
processor = PreTrainProcessor(tokenizer)
|
|
return processor.process
|
|
|
|
def get_sft_processor(tokenizer) -> Callable:
|
|
processor = SFTProcessor(tokenizer)
|
|
return processor.process
|
|
|
|
def get_dpo_processor(tokenizer) -> Callable:
|
|
processor = DPOProcessor(tokenizer)
|
|
return processor.process
|