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
+26 -5
View File
@@ -1,25 +1,46 @@
"""Pre-training data processor."""
from typing import Dict, List, Any
from typing import Any, Dict, List
import torch
from torch import Tensor
from pipeline.tokenize import AutoTokenizer
from pipeline.processors.base import BaseProcessor
from pipeline.processors.base import BaseProcessor, ProcessorSchema
from pipeline.processors.factory import ProcessorFactory
@ProcessorFactory.register("pt")
class PreTrainProcessor(BaseProcessor):
"""Pre-training data processor."""
"""Pre-training data processor.
def __init__(self, tokenizer: AutoTokenizer):
Processes raw text into tokenized sequences with EOS tokens.
Input schema:
- text: str - Raw text string to tokenize
Output schema:
- sequence: int32 tensor - Token IDs with EOS appended
"""
def __init__(
self,
tokenizer: AutoTokenizer,
eos_token: str = "<end▁of▁sentence>",
):
self.tokenizer = tokenizer
self._eos_token = eos_token
@property
def schema(self) -> ProcessorSchema:
return ProcessorSchema(
input_fields={"text": str},
output_fields={"sequence": torch.int32},
)
def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
segment = input_dict["text"]
tokens = self.tokenizer.encode(f"{segment}<end▁of▁sentence>")
tokens = self.tokenizer.encode(f"{segment}{self._eos_token}")
return {"sequence": torch.tensor(tokens, dtype=torch.int32)}
@property