feat: add SFT process_batch for parallel tokenization + short QA filter script

pipeline/processors/sft.py

@@ -66,6 +66,16 @@ class SFTProcessor(BaseProcessor):
             "Input must contain 'messages' or 'query'/'response' pair"
         )
 
+    def _extract_messages(self, input_dict: Dict[str, Any]) -> Optional[List[Dict[str, str]]]:
+        if "messages" in input_dict:
+            return input_dict["messages"]
+        if "query" in input_dict and "response" in input_dict:
+            return [
+                {"role": "user", "content": input_dict["query"]},
+                {"role": "assistant", "content": input_dict["response"]},
+            ]
+        return None
+
     def _process_messages(self, messages: List[Dict[str, str]]) -> Dict[str, Tensor]:
         if not messages:
             raise ValueError("Messages list is empty")
@@ -86,6 +96,54 @@ class SFTProcessor(BaseProcessor):
             "position_ids": position_ids,
         }
 
+    def process_batch(self, input_dicts: List[Dict[str, Any]]) -> List[Optional[Dict[str, Tensor]]]:
+        strategy = self.strategy or ChatMLStrategy(self.tokenizer)
+
+        prompts_text: List[str] = []
+        fulls_text: List[str] = []
+        indices: List[int] = []
+        results: List[Optional[Dict[str, Tensor]]] = [None] * len(input_dicts)
+
+        for i, d in enumerate(input_dicts):
+            try:
+                messages = self._extract_messages(d)
+                if not messages or messages[-1]["role"] != "assistant":
+                    continue
+                last_asst = max(j for j, m in enumerate(messages) if m["role"] == "assistant")
+                prompt_text = self.tokenizer.apply_chat_template(
+                    messages[:last_asst], add_generation_prompt=True, tokenize=False
+                )
+                full_text = self.tokenizer.apply_chat_template(
+                    messages[: last_asst + 1], add_generation_prompt=False, tokenize=False
+                )
+                prompts_text.append(prompt_text)
+                fulls_text.append(full_text)
+                indices.append(i)
+            except Exception:
+                continue
+
+        if not prompts_text:
+            return results
+
+        prompt_tokens_list = self.tokenizer.encode(prompts_text)
+        full_tokens_list = self.tokenizer.encode(fulls_text)
+
+        for j, idx in enumerate(indices):
+            prompt_tokens = prompt_tokens_list[j]
+            full_tokens = full_tokens_list[j]
+            resp_tokens = full_tokens[len(prompt_tokens):]
+            sequence = torch.tensor(prompt_tokens + resp_tokens, dtype=torch.int32)
+            loss_mask = torch.zeros(len(sequence), dtype=torch.bool)
+            loss_mask[len(prompt_tokens):] = True
+            position_ids = torch.arange(len(sequence), dtype=torch.int32)
+            results[idx] = {
+                "sequence": sequence,
+                "loss_mask": loss_mask,
+                "position_ids": position_ids,
+            }
+
+        return results
+
     @property
     def output_keys(self) -> List[str]:
         return ["sequence", "loss_mask", "position_ids"]