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

pipeline/io/export.py

@@ -83,6 +83,7 @@ def cache_jsonl(
     *,
     pack_size: int = -1,
     pad_value: int = 0,
+    batch_size: int = 256,
 ) -> List[str]:
     """Tokenize JSONL files and pack them into HDF5 storage.
 
@@ -92,6 +93,7 @@ def cache_jsonl(
         processor: Initialized Processor instance.
         pack_size: Packing length, <=0 means no packing.
         pad_value: Padding value.
+        batch_size: Number of records passed to the processor at once.
 
     Returns:
         List of generated H5 file paths.
@@ -105,25 +107,48 @@ def cache_jsonl(
 
         arrows: Dict[str, List] = {key: [] for key in output_keys}
 
+        def append_batch(batch):
+            items = [item for _, item in batch]
+            try:
+                results = processor.process_batch(items)
+                if len(results) != len(items):
+                    raise RuntimeError(
+                        "Batch processor returned a different number of results"
+                    )
+            except Exception:
+                results = []
+                for line_num, item in batch:
+                    try:
+                        results.append(processor.process(item))
+                    except Exception as e:
+                        logger.warning(
+                            f"Unexpected error processing line {line_num} "
+                            f"in {file_path}: {e}. Skipping line."
+                        )
+                        results.append(None)
+
+            for result in results:
+                if result is not None:
+                    for key in output_keys:
+                        arrows[key].append(result[key])
+
+        batch = []
+        batch_size = max(1, batch_size)
         with open(file_path, "r", encoding="utf-8") as f:
             for line_num, line in enumerate(
                 tqdm(f, desc=f"Processing {file_name}", leave=False), start=1
             ):
                 try:
-                    result = processor.process(json.loads(line))
+                    batch.append((line_num, json.loads(line)))
-                    if result is not None:
+                    if len(batch) >= batch_size:
-                        for key in output_keys:
+                        append_batch(batch)
-                            arrows[key].append(result[key])
+                        batch = []
                 except json.JSONDecodeError as e:
                     logger.warning(
                         f"JSON decode error in {file_path} line {line_num}: {e}. Skipping line."
                     )
-                    continue
+            if batch:
-                except Exception as e:
+                append_batch(batch)
-                    logger.warning(
-                        f"Unexpected error processing line {line_num} in {file_path}: {e}. Skipping line."
-                    )
-                    continue
 
         if pack_size > 0:
             dtypes = (

pipeline/processors/base.py

@@ -76,6 +76,12 @@ class BaseProcessor(ABC):
         """
         pass
 
+    def process_batch(
+        self, input_dicts: List[Dict[str, Any]]
+    ) -> List[Dict[str, Tensor]]:
+        """Process a batch, falling back to the single-record implementation."""
+        return [self.process(input_dict) for input_dict in input_dicts]
+
     @property
     @abstractmethod
     def output_keys(self) -> List[str]:

pipeline/processors/dpo.py

@@ -74,6 +74,33 @@ class DPOProcessor(BaseProcessor):
             "rejected_mask": rejected_m,
         }
 
+    def process_batch(self, input_dicts: List[Dict[str, Any]]) -> List[Dict[str, Tensor]]:
+        query_batch = self.tokenizer.encode([item["query"] for item in input_dicts])
+        chosen_batch = self.tokenizer.encode([item["chosen"] for item in input_dicts])
+        rejected_batch = self.tokenizer.encode(
+            [item["rejected"] for item in input_dicts]
+        )
+        results = []
+        for query_tokens, chosen_tokens, rejected_tokens in zip(
+            query_batch, chosen_batch, rejected_batch
+        ):
+            prompt = self.strategy.assemble_prompt(query_tokens)
+            chosen_t, chosen_m = encode_with_mask(
+                prompt, self.strategy.assemble_response(chosen_tokens)
+            )
+            rejected_t, rejected_m = encode_with_mask(
+                prompt, self.strategy.assemble_response(rejected_tokens)
+            )
+            results.append(
+                {
+                    "chosen": chosen_t,
+                    "chosen_mask": chosen_m,
+                    "rejected": rejected_t,
+                    "rejected_mask": rejected_m,
+                }
+            )
+        return results
+
     @property
     def output_keys(self) -> List[str]:
         return ["chosen", "chosen_mask", "rejected", "rejected_mask"]

pipeline/processors/pretrain.py

@@ -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"]

pipeline/processors/sft.py

@@ -29,15 +29,20 @@ class SFTProcessor(BaseProcessor):
         - sequence: int32 tensor - Combined token IDs (prompt + response)
         - loss_mask: bool tensor - True for response tokens (compute loss)
         - position_ids: int32 tensor - Per-sample position IDs starting from 0
+
+    Only the final assistant message is trained (mask_history behavior).
+    All earlier turns are context/prompt and masked from loss.
     """
 
     def __init__(
         self,
         tokenizer: AutoTokenizer,
         strategy: Optional[PromptStrategy] = None,
+        max_seq_len: Optional[int] = None,
     ):
         self.tokenizer = tokenizer
         self.strategy = strategy
+        self.max_seq_len = max_seq_len
 
     @property
     def schema(self) -> ProcessorSchema:
@@ -63,6 +68,35 @@ class SFTProcessor(BaseProcessor):
             "Input must contain 'messages' or 'query'/'response' pair"
         )
 
+    def process_batch(self, input_dicts: List[Dict[str, Any]]) -> List[Dict[str, Tensor]]:
+        results: List[Optional[Dict[str, Tensor]]] = [None] * len(input_dicts)
+        message_indices = [i for i, item in enumerate(input_dicts) if "messages" in item]
+        legacy_indices = [
+            i
+            for i, item in enumerate(input_dicts)
+            if "messages" not in item and "query" in item and "response" in item
+        ]
+        if len(message_indices) + len(legacy_indices) != len(input_dicts):
+            raise KeyError("Input must contain 'messages' or 'query'/'response' pair")
+
+        if message_indices:
+            items = [input_dicts[i] for i in message_indices]
+            batch_results = self._process_messages_batch(
+                [item["messages"] for item in items]
+            )
+            for index, result in zip(message_indices, batch_results):
+                results[index] = result
+
+        if legacy_indices:
+            items = [input_dicts[i] for i in legacy_indices]
+            batch_results = self._process_legacy_batch(items)
+            for index, result in zip(legacy_indices, batch_results):
+                results[index] = result
+
+        if any(result is None for result in results):
+            raise RuntimeError("Batch processing did not produce all results")
+        return results
+
     def _process_messages(self, messages: List[Dict[str, str]]) -> Dict[str, Tensor]:
         if not messages:
             raise ValueError("Messages list is empty")
@@ -73,21 +107,80 @@ class SFTProcessor(BaseProcessor):
             i for i, m in enumerate(messages) if m["role"] == "assistant"
         )
 
-        prompt_tokens = self.tokenizer.apply_chat_template(
+        full_text = self.tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=False
+        )
+        full_ids = self.tokenizer.encode(full_text, add_special_tokens=False)
+
+        prompt_text = self.tokenizer.apply_chat_template(
             messages[:last_asst_idx],
+            tokenize=False,
             add_generation_prompt=True,
-            tokenize=True,
         )
+        prompt_ids = self.tokenizer.encode(prompt_text, add_special_tokens=False)
 
-        resp_content = messages[last_asst_idx]["content"]
+        resp_ids = full_ids[len(prompt_ids) :]
-        im_end = getattr(self.tokenizer, "im_end", "<|im_end|>")
+        if not resp_ids:
-        resp_tokens = self.tokenizer.encode(
+            raise ValueError("Empty assistant response")
-            f"{resp_content}{im_end}\n", add_special_tokens=False
+
-        )
+        tokens, loss_mask = encode_with_mask(prompt_ids, list(resp_ids))
+
+        if self.max_seq_len and len(tokens) > self.max_seq_len:
+            tokens = tokens[: self.max_seq_len]
+            loss_mask = loss_mask[: self.max_seq_len]
 
-        tokens, loss_mask = encode_with_mask(prompt_tokens, resp_tokens)
         position_ids = torch.arange(len(tokens), dtype=torch.int32)
-        return {"sequence": tokens, "loss_mask": loss_mask, "position_ids": position_ids}
+        return {
+            "sequence": tokens,
+            "loss_mask": loss_mask,
+            "position_ids": position_ids,
+        }
+
+    def _process_messages_batch(
+        self, conversations: List[List[Dict[str, str]]]
+    ) -> List[Dict[str, Tensor]]:
+        for messages in conversations:
+            if not messages:
+                raise ValueError("Messages list is empty")
+            if messages[-1]["role"] != "assistant":
+                raise ValueError("Last message must have role 'assistant'")
+
+        assistant_indices = [
+            max(i for i, message in enumerate(messages) if message["role"] == "assistant")
+            for messages in conversations
+        ]
+        full_texts = [
+            self.tokenizer.apply_chat_template(
+                messages, tokenize=False, add_generation_prompt=False
+            )
+            for messages in conversations
+        ]
+        prompt_texts = [
+            self.tokenizer.apply_chat_template(
+                messages[:assistant_idx], tokenize=False, add_generation_prompt=True
+            )
+            for messages, assistant_idx in zip(conversations, assistant_indices)
+        ]
+        full_ids_batch = self.tokenizer.encode(full_texts, add_special_tokens=False)
+        prompt_ids_batch = self.tokenizer.encode(prompt_texts, add_special_tokens=False)
+
+        results = []
+        for full_ids, prompt_ids in zip(full_ids_batch, prompt_ids_batch):
+            resp_ids = full_ids[len(prompt_ids) :]
+            if not resp_ids:
+                raise ValueError("Empty assistant response")
+            tokens, loss_mask = encode_with_mask(prompt_ids, list(resp_ids))
+            if self.max_seq_len and len(tokens) > self.max_seq_len:
+                tokens = tokens[: self.max_seq_len]
+                loss_mask = loss_mask[: self.max_seq_len]
+            results.append(
+                {
+                    "sequence": tokens,
+                    "loss_mask": loss_mask,
+                    "position_ids": torch.arange(len(tokens), dtype=torch.int32),
+                }
+            )
+        return results
 
     def _process_legacy(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
         strategy = self.strategy or ChatMLStrategy(self.tokenizer)
@@ -102,6 +195,28 @@ class SFTProcessor(BaseProcessor):
         position_ids = torch.arange(len(tokens), dtype=torch.int32)
         return {"sequence": tokens, "loss_mask": loss_mask, "position_ids": position_ids}
 
+    def _process_legacy_batch(
+        self, input_dicts: List[Dict[str, Any]]
+    ) -> List[Dict[str, Tensor]]:
+        strategy = self.strategy or ChatMLStrategy(self.tokenizer)
+        query_batch = self.tokenizer.encode([item["query"] for item in input_dicts])
+        response_batch = self.tokenizer.encode(
+            [item["response"] for item in input_dicts]
+        )
+        results = []
+        for query_tokens, response_tokens in zip(query_batch, response_batch):
+            prompt = strategy.assemble_prompt(query_tokens)
+            response = strategy.assemble_response(response_tokens)
+            tokens, loss_mask = encode_with_mask(prompt, response)
+            results.append(
+                {
+                    "sequence": tokens,
+                    "loss_mask": loss_mask,
+                    "position_ids": torch.arange(len(tokens), dtype=torch.int32),
+                }
+            )
+        return results
+
     @property
     def output_keys(self) -> List[str]:
         return ["sequence", "loss_mask", "position_ids"]

pipeline/tokenize/chat_template.py

@@ -3,6 +3,7 @@ Chat template module with Jinja2 rendering support.
 """
 
 from dataclasses import dataclass, field
+from functools import cached_property
 from typing import Any, Dict, List, Optional
 
 from jinja2 import Template
@@ -32,6 +33,10 @@ class ChatTemplate:
     default_variables: Dict[str, Any] = field(default_factory=dict)
     special_tokens: Dict[str, str] = field(default_factory=dict)
 
+    @cached_property
+    def _compiled(self) -> Template:
+        return Template(self.template_str)
+
     @classmethod
     def from_string(
         cls,
@@ -79,8 +84,7 @@ class ChatTemplate:
         if system_prompt is not None:
             variables["system_prompt"] = system_prompt
 
-        jinja_template = Template(self.template_str)
+        return self._compiled.render(**variables)
-        return jinja_template.render(**variables)
 
 
 # Default ChatML template

pipeline/tokenize/tokenizer.py

@@ -3,6 +3,7 @@ Tokenizer module with BPE implementation and auto-loading support.
 """
 
 from dataclasses import dataclass
+from functools import cached_property
 import json
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Union
@@ -102,6 +103,10 @@ class ChatTemplate:
         if self.special_tokens is None:
             self.special_tokens = {}
 
+    @cached_property
+    def _compiled(self) -> Template:
+        return Template(self.template_str)
+
     @classmethod
     def from_string(
         cls,
@@ -142,8 +147,7 @@ class ChatTemplate:
         if system_prompt is not None:
             variables["system_prompt"] = system_prompt
 
-        jinja_template = Template(self.template_str)
+        return self._compiled.render(**variables)
-        return jinja_template.render(**variables)
 
 
 
@@ -326,7 +330,9 @@ class AutoTokenizer:
             KeyError: If template name is not registered.
         """
         if isinstance(template, str):
-            self._chat_template = ChatTemplate.from_string(template)
+            self._chat_template = ChatTemplate.from_string(
+                template, special_tokens=self._special_token_map
+            )
         elif isinstance(template, ChatTemplate):
             self._chat_template = template
         else:

scripts/cache_h5.py

@@ -53,6 +53,12 @@ def main():
     parser.add_argument(
         "--pad-value", type=int, default=0, help="Padding value (default: 0)"
     )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=256,
+        help="Records tokenized per batch (default: 256)",
+    )
     parser.add_argument(
         "--log-level",
         default="INFO",
@@ -105,6 +111,7 @@ def main():
         processor=processor,
         pack_size=args.pack_size,
         pad_value=args.pad_value,
+        batch_size=args.batch_size,
     )
     print(f"\nDone! Output saved to {output_dir}")
 

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):