diff --git a/pipeline/processors/sft.py b/pipeline/processors/sft.py index acad467..3f70002 100644 --- a/pipeline/processors/sft.py +++ b/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"] diff --git a/scripts/filter_short.py b/scripts/filter_short.py new file mode 100644 index 0000000..88e70fc --- /dev/null +++ b/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()