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

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