feat: processors 支持批量 tokenize，优化性能并缓存 chat template

2026-08-05 12:21:25 +08:00
parent aa2ea4f3a6
commit 65dadac10f
10 changed files with 290 additions and 25 deletions
@@ -43,6 +43,14 @@ class PreTrainProcessor(BaseProcessor):
        tokens = self.tokenizer.encode(f"{segment}{self._eos_token}")
        return {"sequence": torch.tensor(tokens, dtype=torch.int32)}

+    def process_batch(self, input_dicts: List[Dict[str, Any]]) -> List[Dict[str, Tensor]]:
+        texts = [f"{item['text']}{self._eos_token}" for item in input_dicts]
+        encoded = self.tokenizer.encode(texts)
+        return [
+            {"sequence": torch.tensor(tokens, dtype=torch.int32)}
+            for tokens in encoded
+        ]
+
    @property
    def output_keys(self) -> List[str]:
        return ["sequence"]