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