reafactor: 重构项目

pipeline/processors/pretrain.py

@@ -0,0 +1,26 @@
+"""Pre-training data processor."""
+from typing import Dict, List, Any
+
+import torch
+from torch import Tensor
+
+from pipeline.tokenizer import BpeTokenizer
+from pipeline.processors.base import BaseProcessor
+from pipeline.processors.factory import ProcessorFactory
+
+
+@ProcessorFactory.register("pt")
+class PreTrainProcessor(BaseProcessor):
+    """Pre-training data processor."""
+
+    def __init__(self, tokenizer: BpeTokenizer):
+        self.tokenizer = tokenizer
+
+    def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
+        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"]