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

tests/test_cache.py

@@ -29,6 +29,16 @@ class DummyProcessor(BaseProcessor):
         }
 
 
+class BatchTrackingProcessor(DummyProcessor):
+    def __init__(self):
+        super().__init__()
+        self.batch_sizes = []
+
+    def process_batch(self, items):
+        self.batch_sizes.append(len(items))
+        return super().process_batch(items)
+
+
 class TestCacheJsonl:
     def test_basic_cache_functionality(self):
         with tempfile.TemporaryDirectory() as tmpdir:
@@ -121,3 +131,19 @@ class TestCacheJsonl:
                 pad_value=0,
             )
             assert len(output_files) == 1
+
+    def test_uses_configured_batch_size(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            jsonl_path = os.path.join(tmpdir, "test.jsonl")
+            with open(jsonl_path, "w", encoding="utf-8") as f:
+                for i in range(5):
+                    f.write(json.dumps({"text": str(i)}) + "\n")
+
+            processor = BatchTrackingProcessor()
+            cache_jsonl(
+                files=[jsonl_path],
+                output_dir=tmpdir,
+                processor=processor,
+                batch_size=2,
+            )
+            assert processor.batch_sizes == [2, 2, 1]

tests/test_processors.py

@@ -15,7 +15,9 @@ from pipeline.processors import (
 class DummyTokenizer:
     im_end = "<|im_end|>"
 
-    def encode(self, text: str, add_special_tokens: bool = False):
+    def encode(self, text, add_special_tokens: bool = False):
+        if isinstance(text, list):
+            return [[ord(c) for c in item] for item in text]
         return [ord(c) for c in text]
 
     def apply_chat_template(
@@ -50,6 +52,16 @@ class TestPreTrainProcessor:
         result = PreTrainProcessor(DummyTokenizer()).process({"text": "a"})
         assert len(result["sequence"]) > 0
 
+    def test_process_batch_matches_single(self):
+        processor = PreTrainProcessor(DummyTokenizer())
+        items = [{"text": "hello"}, {"text": "world"}]
+        batch = processor.process_batch(items)
+        single = [processor.process(item) for item in items]
+        assert all(
+            torch.equal(batch_item["sequence"], single_item["sequence"])
+            for batch_item, single_item in zip(batch, single)
+        )
+
 
 class TestSFTProcessor:
     def test_output_keys(self):
@@ -121,6 +133,23 @@ class TestSFTProcessor:
         })
         assert "sequence" in result
 
+    def test_process_batch_matches_single(self):
+        processor = SFTProcessor(DummyTokenizer())
+        items = [
+            {
+                "messages": [
+                    {"role": "user", "content": "q1"},
+                    {"role": "assistant", "content": "a1"},
+                ]
+            },
+            {"query": "q2", "response": "a2"},
+        ]
+        batch = processor.process_batch(items)
+        single = [processor.process(item) for item in items]
+        for batch_item, single_item in zip(batch, single):
+            for key in processor.output_keys:
+                assert torch.equal(batch_item[key], single_item[key])
+
     def test_messages_empty_raises(self):
         with pytest.raises(ValueError, match="Messages list is empty"):
             SFTProcessor(DummyTokenizer()).process({"messages": []})
@@ -173,6 +202,18 @@ class TestDPOProcessor:
         assert result["chosen_mask"].dtype == torch.bool
         assert result["rejected_mask"].dtype == torch.bool
 
+    def test_process_batch_matches_single(self):
+        processor = DPOProcessor(DummyTokenizer())
+        items = [
+            {"query": "q1", "chosen": "yes", "rejected": "no"},
+            {"query": "q2", "chosen": "good", "rejected": "bad"},
+        ]
+        batch = processor.process_batch(items)
+        single = [processor.process(item) for item in items]
+        for batch_item, single_item in zip(batch, single):
+            for key in processor.output_keys:
+                assert torch.equal(batch_item[key], single_item[key])
+
 
 class TestProcessorFactory:
     def test_create_pre_train_processor(self):