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

scripts/filter_short.py

@@ -0,0 +1,92 @@
+import argparse
+import json
+import os
+import tempfile
+import shutil
+
+MIN_LEN = 15
+
+
+def filter_sft(input_path: str) -> tuple[int, int]:
+    """Filter SFT JSONL (messages format), remove if any msg content < MIN_LEN chars."""
+    kept, total = 0, 0
+    tmp_fd, tmp_path = tempfile.mkstemp(dir=os.path.dirname(input_path))
+    try:
+        with open(input_path, encoding="utf-8") as fin, open(tmp_fd, "w", encoding="utf-8") as fout:
+            for line in fin:
+                line = line.strip()
+                if not line:
+                    continue
+                total += 1
+                try:
+                    obj = json.loads(line)
+                except json.JSONDecodeError:
+                    continue
+                messages = obj.get("messages", [])
+                short = any(len(m.get("content", "")) < MIN_LEN for m in messages)
+                if not short:
+                    fout.write(line + "\n")
+                    kept += 1
+        shutil.move(tmp_path, input_path)
+    except Exception:
+        if os.path.exists(tmp_path):
+            os.unlink(tmp_path)
+        raise
+    return kept, total
+
+
+def filter_pretrain(input_path: str) -> tuple[int, int]:
+    """Filter pretrain JSONL (text format), remove if text < MIN_LEN chars."""
+    kept, total = 0, 0
+    tmp_fd, tmp_path = tempfile.mkstemp(dir=os.path.dirname(input_path))
+    try:
+        with open(input_path, encoding="utf-8") as fin, open(tmp_fd, "w", encoding="utf-8") as fout:
+            for line in fin:
+                line = line.strip()
+                if not line:
+                    continue
+                total += 1
+                try:
+                    obj = json.loads(line)
+                except json.JSONDecodeError:
+                    continue
+                text = obj.get("text", "")
+                if len(text) >= MIN_LEN:
+                    fout.write(line + "\n")
+                    kept += 1
+        shutil.move(tmp_path, input_path)
+    except Exception:
+        if os.path.exists(tmp_path):
+            os.unlink(tmp_path)
+        raise
+    return kept, total
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Filter short samples from JSONL datasets")
+    parser.add_argument("input_dir", help="Directory containing JSONL files")
+    parser.add_argument("--type", choices=["sft", "pt"], required=True, help="Dataset type")
+    args = parser.parse_args()
+
+    from pipeline import FileScanner
+
+    jsonl_files = FileScanner.scan(args.input_dir, suffix=".jsonl")
+    if not jsonl_files:
+        print(f"No JSONL files found in {args.input_dir}")
+        return
+
+    filter_fn = filter_sft if args.type == "sft" else filter_pretrain
+
+    total_kept, total_lines = 0, 0
+    for fpath in jsonl_files:
+        kept, lines = filter_fn(fpath)
+        total_kept += kept
+        total_lines += lines
+        removed = lines - kept
+        print(f"  {os.path.basename(fpath)}: {lines} -> {kept} (removed {removed})")
+
+    print(f"\nTotal: {total_lines} -> {total_kept} (removed {total_lines - total_kept})")
+
+
+if __name__ == "__main__":
+    main()