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))
-                    if result is not None:
-                        for key in output_keys:
-                            arrows[key].append(result[key])
+                    batch.append((line_num, json.loads(line)))
+                    if len(batch) >= batch_size:
+                        append_batch(batch)
+                        batch = []
                 except json.JSONDecodeError as e:
                     logger.warning(
                         f"JSON decode error in {file_path} line {line_num}: {e}. Skipping line."
                     )
-                    continue
-                except Exception as e:
-                    logger.warning(
-                        f"Unexpected error processing line {line_num} in {file_path}: {e}. Skipping line."
-                    )
-                    continue
+            if batch:
+                append_batch(batch)
 
         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"]
-        im_end = getattr(self.tokenizer, "im_end", "<|im_end|>")
-        resp_tokens = self.tokenizer.encode(
-            f"{resp_content}{im_end}\n", add_special_tokens=False
-        )
+        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]
 
-        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 jinja_template.render(**variables)
+        return self._compiled.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 jinja_template.render(**variables)
+        return self._compiled.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: