perf: 优化 processors、cache、packing 模块性能并简化 README

README.md

@@ -1,6 +1,6 @@
 # DataPipeline
 
-用于训练 KHAOSZ 模型的数据集处理工具。提供文本导出、Tokenize、序列打包、H5 存储等独立工具，支持预训练 / SFT / DPO 三种训练范式。
+数据集处理工具，支持预训练 / SFT / DPO 三种训练范式。
 
 ## 项目结构
 
@@ -8,231 +8,104 @@
 pipeline/
 ├── tokenizer.py      # BPE 分词器
 ├── text.py           # 文本规范化
-├── packing.py        # 序列打包（bin-packing）
+├── packing.py        # 序列打包
 ├── io.py             # 文件/HDF5 读写
 ├── processors.py     # PT / SFT / DPO 处理器
-├── export.py         # Dataset → JSONL 导出
+├── export.py         # Dataset → JSONL
-└── cache.py          # JSONL → Tokenize → H5 缓存
+└── cache.py          # JSONL → Tokenize → H5
-
-pre_train/                          # 预训练数据处理脚本
-supervised_finetuning/              # SFT 数据处理脚本
-reforce_learning/                   # DPO 数据处理脚本
 ```
 
 ## 设计理念
 
-每个模块**独立可用、零互相依赖**，调用者按需组合：
+模块**独立可用**，通过磁盘文件解耦，按需组合：
 
 ```
-HuggingFace Hub
+Dataset → export_dataset() → JSONL → cache_jsonl() → HDF5
-      │
+                              ↑
-      ▼  load_dataset()
+                        processors.py
-  DatasetDict
+                        packing.py
-      │
+                        io.py
-      ▼  export_dataset()          ← pipeline/export.py
-  JSONL 文件
-      │
-      ▼  cache_jsonl()            ← pipeline/cache.py
-      │     ├─ Processor.process()    ← pipeline/processors.py
-      │     ├─ SequencePacker.pack()  ← pipeline/packing.py
-      │     └─ IOHandler.save_h5()    ← pipeline/io.py
-  HDF5 张量文件
 ```
 
-> 各阶段之间通过磁盘文件解耦。你可以只执行阶段 1（导出 JSONL），也可以继续执行阶段 2（tokenize 并缓存为 H5），按需选择。
+## 使用方法
 
-## 快速开始
+### 1. 导出数据集
-
-### 安装依赖
-
-```bash
-pip install -e .
-```
-
-### 阶段 1：导出数据集为 JSONL
 
 ```python
 from datasets import load_dataset
 from pipeline.export import export_dataset
 
 dataset = load_dataset("your-dataset")
-export_dataset(
-    dataset=dataset["train"],        # 直接传 Dataset，不传 DatasetDict
-    output_dir="./dataset",
-    output_prefix="my-data",
-    max_chunks=5,                    # 可选，限制 chunk 数量（调试用）
-)
-```
-
-**自定义转换函数：**
-
-```python
-def process_func(example):
-    # 提取字段、转换格式、展开多轮对话等
-    return {"query": example["instruction"], "response": example["output"]}
-
 export_dataset(
     dataset=dataset["train"],
-    output_dir="./dataset",
+    output_dir="./data",
-    output_prefix="my-sft",
+    output_prefix="train",
-    process_func=process_func,
+    process_func=lambda x: {"text": x["content"]},  # 可选
 )
 ```
 
-> `process_func` 返回单个 `dict` 或 `list[dict]`（一条样本可展开为多条）。
+### 2. Tokenize 并缓存
-
-**使用文本规范化：**
 
 ```python
-from pipeline.text import TextNormalizer
+from pipeline import BpeTokenizer, ProcessorFactory, cache_jsonl
-
-normalizer = TextNormalizer()
-
-def process_func(example):
-    return {"text": normalizer.normalize(example["content"])}
-
-export_dataset(
-    dataset=dataset["train"],
-    output_dir="./dataset",
-    output_prefix="my-pretrain",
-    process_func=process_func,
-)
-```
-
-### 阶段 2：Tokenize 并缓存为 HDF5
-
-```python
-from pipeline.tokenizer import BpeTokenizer
-from pipeline.processors import ProcessorFactory
-from pipeline.cache import cache_jsonl
 
 tokenizer = BpeTokenizer("tokenizer.json")
-processor = ProcessorFactory.create("pt", tokenizer)
+processor = ProcessorFactory.create("pt", tokenizer)  # "pt" | "sft" | "dpo"
 
 cache_jsonl(
-    files=["./dataset/my-pretrain_chunk_0.jsonl"],
+    files=["./data/train.jsonl"],
     output_dir="./cached",
     processor=processor,
-    pack_size=4096,       # 可选，打包长度；<=0 不打包
+    pack_size=4096,  # <=0 不打包
     pad_value=1,
 )
 ```
 
-**处理器类型：**
+### 3. 处理器类型
 
-| 类型 | 工厂 key | 输入格式 | 输出 keys |
+| 类型 | key | 输入 | 输出 |
-|------|----------|---------|-----------|
+|------|-----|------|------|
 | 预训练 | `"pt"` | `{"text": "..."}` | `["sequence"]` |
 | SFT | `"sft"` | `{"query": "...", "response": "..."}` | `["sequence", "loss_mask"]` |
-| DPO | `"dpo"` | 待定 | `["chosen", "chosen_mask", "rejected", "rejected_mask"]` |
+| DPO | `"dpo"` | `{"query": "...", "chosen": "...", "rejected": "..."}` | `["chosen", "chosen_mask", "rejected", "rejected_mask"]` |
 
-## 独立工具参考
+## 参数参考
 
-### BpeTokenizer
+### export_dataset()
 
-```python
+| 参数 | 类型 | 默认 | 说明 |
-from pipeline.tokenizer import BpeTokenizer
+|------|------|------|------|
+| `dataset` | Dataset | 必填 | HuggingFace Dataset |
+| `output_dir` | str | 必填 | 输出目录 |
+| `output_prefix` | str | 必填 | 文件名前缀 |
+| `chunk_size` | int | 1_000_000 | 每个文件的最大样本数 |
+| `max_chunks` | int | None | 最大 chunk 数量 |
+| `process_func` | callable | None | 样本转换函数 |
+| `column` | str | "text" | 默认文本列名 |
 
-tokenizer = BpeTokenizer("tokenizer.json")
+### cache_jsonl()
-ids = tokenizer.encode("hello world")           # → [1234, 5678, 1]
-text = tokenizer.decode(ids)                     # → "hello world"
-len(tokenizer)                                   # → 词表大小
-```
 
-### TextNormalizer
+| 参数 | 类型 | 默认 | 说明 |
-
+|------|------|------|------|
-```python
+| `files` | List[str] | 必填 | JSONL 文件列表 |
-from pipeline.text import TextNormalizer
+| `output_dir` | str | 必填 | 输出目录 |
-
+| `processor` | BaseProcessor | 必填 | 处理器实例 |
-normalizer = TextNormalizer()
+| `pack_size` | int | -1 | 打包长度，<=0 不打包 |
-text = normalizer.normalize(raw_text)
+| `pad_value` | int | 1 | 填充值 |
-```
-
-替换规则包括：全角引号 → 半角、各种短横线统一、不间断空格 → 普通空格等。支持自定义规则：
-
-```python
-normalizer = TextNormalizer(custom_rules={"旧词": "新词"})
-```
 
 ### SequencePacker
 
 ```python
-from pipeline.packing import SequencePacker
-
 packer = SequencePacker(pack_size=4096, pad_value=0)
-packed = packer.pack(list_of_tensors)    # → List[Tensor]，每个长度为 pack_size
+packed = packer.pack([tensor1, tensor2, ...])  # → List[Tensor]
 ```
 
 ### IOHandler
 
 ```python
-from pipeline.io import IOHandler
-
 # 保存
-IOHandler.save_h5("./output", "my_data", {"sequence": [tensor1, tensor2]})
+IOHandler.save_h5("./out", "name", {"key": [tensor1, tensor2]})
 
 # 加载
-data = IOHandler.load_h5("./output")     # → {"sequence": [tensor1, tensor2, ...]}
+data = IOHandler.load_h5("./out")  # → {"key": [tensor1, ...]}
-
-# 遍历文件
-files = IOHandler.fetch_files("./dataset")
-folders = IOHandler.fetch_folders("./dataset")
-```
-
-### 自定义 Processor
-
-```python
-from pipeline.processors import BaseProcessor, ProcessorFactory
-import torch
-
-class MyProcessor(BaseProcessor):
-    def __init__(self, tokenizer):
-        self.tokenizer = tokenizer
-
-    def process(self, input_dict):
-        tokens = self.tokenizer.encode(input_dict["text"])
-        return {"sequence": torch.tensor(tokens, dtype=torch.int32)}
-
-    @property
-    def output_keys(self):
-        return ["sequence"]
-
-ProcessorFactory.register("my_type", MyProcessor)
-```
-
-## 运行脚本
-
-```bash
-# 预训练
-python pre_train/chinese-c4.py
-python pre_train/english-wiki.py
-
-# SFT
-python supervised_finetuning/sft_belle.py
-python supervised_finetuning/sft_coder.py
-
-# DPO
-python reforce_learning/dpp_chinese_dpo_pairs.py
-```
-
-## 输出格式
-
-**JSONL**（阶段 1 输出）：
-
-```jsonl
-{"text": "训练文本内容..."}
-{"query": "问题", "response": "答案"}
-```
-
-**HDF5**（阶段 2 输出）：
-
-```
-my_data.h5
-├── sequence/
-│   ├── data_0    # Tensor (4096,) int32
-│   ├── data_1    # Tensor (4096,) int32
-│   └── ...
-└── loss_mask/     # 仅 SFT
-    ├── data_0    # Tensor (4096,) bool
-    └── ...
 ```

pipeline/cache.py

@@ -2,7 +2,7 @@
 import json
 import os
 import logging
-from typing import List
+from typing import List, Dict
 from pathlib import Path
 
 from tqdm import tqdm
@@ -39,33 +39,40 @@ def cache_jsonl(
     """
     os.makedirs(output_dir, exist_ok=True)
     output_files: List[str] = []
+    # Cache output_keys to avoid repeated attribute access
+    output_keys = processor.output_keys
 
     for file_path in files:
         file_name = Path(file_path).stem
 
-        arrows = []
+        # Pre-allocate lists for each output key
+        arrows: Dict[str, List] = {key: [] for key in output_keys}
+        
+        # Read and process all lines
         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:
-                    arrow = processor.process(json.loads(line))
+                    result = processor.process(json.loads(line))
+                    if result is not None:
+                        # Batch append: add each key's tensor to corresponding list
+                        for key in output_keys:
+                            arrows[key].append(result[key])
                 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 arrow is not None:
-                    arrows.append(arrow)
 
-        package = {key: [a[key] for a in arrows] for key in processor.output_keys}
+        # Convert lists to tensors once per key
-
-        output = {}
-        for key in processor.output_keys:
         if pack_size > 0:
-                packer = SequencePacker(pack_size, pad_value)  # independent instance per key
+            output = {}
-                output[key] = packer.pack(package[key])
+            for key in output_keys:
+                packer = SequencePacker(pack_size, pad_value)
+                output[key] = packer.pack(arrows[key])
         else:
-                output[key] = package[key]
+            # No packing: directly use the arrow tensors
+            output = arrows
 
         IOHandler.save_h5(output_dir, file_name, output)
         h5_path = os.path.join(output_dir, f"{file_name}.h5")

pipeline/io.py

@@ -34,6 +34,7 @@ class IOHandler:
     def save_h5(output_dir: str, file_name: str, tensor_group: Dict[str, List[Tensor]]) -> None:
         os.makedirs(output_dir, exist_ok=True)
         full_path = os.path.join(output_dir, f"{file_name}.h5")
+        
         with h5py.File(full_path, 'w') as f:
             for key, tensors in tensor_group.items():
                 grp = f.create_group(key)

pipeline/packing.py

@@ -1,5 +1,5 @@
 import logging
-from typing import List
+from typing import List, Optional
 import torch
 from torch import Tensor
 
@@ -14,14 +14,23 @@ class SequencePacker:
         self.pack_size = pack_size
         self.pad_value = pad_value
         self.dtype = dtype
+        # Pre-allocate buffer for better performance
+        self._buffer: Optional[Tensor] = None
         self._reset()
 
     def _reset(self) -> None:
         """Reset internal state for instance reuse."""
-        self._current_pack = torch.full(
+        # Reuse buffer instead of creating new tensors
+        if self._buffer is None or self._buffer.shape[0] != self.pack_size:
+            self._buffer = torch.full(
                 (self.pack_size,), self.pad_value, dtype=self.dtype
             )
+        else:
+            self._buffer.fill_(self.pad_value)
         self._current_pos = 0
+        self._packages: List[Tensor] = []
+        # Backward compatibility: maintain _current_pack reference
+        self._current_pack = self._buffer
 
     @error_handler()
     def pack(self, sequences: List[Tensor]) -> List[Tensor]:
@@ -37,53 +46,63 @@ class SequencePacker:
         # Input validation
         if not sequences:
             return []
+        
+        # Validate and cache tensor sizes in one pass
+        tensor_sizes = []
         for i, seq in enumerate(sequences):
             if seq.dim() != 1:
                 raise ValueError(
                     f"Expected 1D tensor at index {i}, got {seq.dim()}D tensor with shape {seq.shape}"
                 )
-            # Check dtype compatibility and warn if mismatched
+            tensor_sizes.append(seq.numel())
             if seq.dtype != self.dtype:
                 logger.warning(
                     f"Input tensor dtype {seq.dtype} does not match packer dtype {self.dtype}, "
                     f"will be converted. This may affect packing efficiency."
                 )
 
-        packages = []
+        # Reset state for new packing
-        # Sort by length in descending order to improve packing efficiency
+        self._packages = []
-        # Use sorted() to avoid modifying the input list
+        self._reset()
-        sorted_sequences = sorted(sequences, key=lambda x: x.numel(), reverse=True)
         
-        for tensor in sorted_sequences:
+        # Combine sequences with their sizes for sorting
+        indexed_seqs = list(zip(sequences, tensor_sizes))
+        # Sort by size descending (First-Fit Decreasing algorithm)
+        indexed_seqs.sort(key=lambda x: x[1], reverse=True)
+
+        for tensor, tensor_size in indexed_seqs:
             # Truncate sequences that exceed pack_size
-            if tensor.numel() > self.pack_size:
+            if tensor_size > self.pack_size:
                 logger.warning(
-                    f"Sequence length {tensor.numel()} exceeds pack_size {self.pack_size}, truncating"
+                    f"Sequence length {tensor_size} exceeds pack_size {self.pack_size}, truncating"
                 )
+                tensor_size = self.pack_size
                 tensor = tensor[: self.pack_size]
-            tensor_size = tensor.numel()
 
             # Current package is full, create a new one
             if self._current_pos + tensor_size > self.pack_size:
-                packages.append(self._current_pack)
+                # Finish current package (pad to pack_size)
-                self._current_pack = torch.full(
+                package = self._buffer.clone()
-                    (self.pack_size,), self.pad_value, dtype=self.dtype
+                self._packages.append(package)
-                )
+                # Reset buffer for reuse
+                self._buffer.fill_(self.pad_value)
                 self._current_pos = 0
 
-            # Place tensor in current package (remaining positions stay as pad_value)
+            # Place tensor in current package
-            self._current_pack[self._current_pos : self._current_pos + tensor_size] = (
+            self._buffer[self._current_pos : self._current_pos + tensor_size] = tensor
-                tensor
-            )
             self._current_pos += tensor_size
 
-        # Handle the last package
+        # Handle the last package (pad to pack_size)
         if self._current_pos > 0:
-            packages.append(self._current_pack)
+            package = self._buffer.clone()
+            self._packages.append(package)
+
+        # Clear buffer and reset state for backward compatibility
+        self._buffer = None
         self._current_pack = None
         self._current_pos = 0
         
-        return packages
+        return self._packages
 
     def reset(self) -> None:
         """Reset packer state for reuse. More efficient than creating a new instance."""

pipeline/processors.py

@@ -41,14 +41,18 @@ class SFTProcessor(BaseProcessor):
         self.tokenizer = tokenizer
 
     def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
-        query, response = input_dict["query"], input_dict["response"]
+        query = input_dict["query"]
+        response = input_dict["response"]
+        
         q = self.tokenizer.encode(
             f"<|im_start|>user\n{query}<|im_end|>\n<|im_start|>assistant\n"
         )
         a = self.tokenizer.encode(f"{response}<|im_end|>\n<eos>")
+        
+        q_len = len(q)
         tokens = torch.tensor(q + a, dtype=torch.int32)
-        loss_mask = torch.zeros_like(tokens, dtype=torch.bool)
+        loss_mask = torch.zeros(q_len + len(a), dtype=torch.bool)
-        loss_mask[len(q):] = True
+        loss_mask[q_len:] = True
         return {"sequence": tokens, "loss_mask": loss_mask}
 
     @property
@@ -72,14 +76,19 @@ class DPOProcessor(BaseProcessor):
         )
 
         chosen = self.tokenizer.encode(f"{chosen_response}<|im_end|>\n<eos>")
+        q_len = len(q)
+        chosen_len = len(chosen)
+        
         chosen_tokens = torch.tensor(q + chosen, dtype=torch.int32)
-        chosen_mask = torch.zeros_like(chosen_tokens, dtype=torch.bool)
+        chosen_mask = torch.zeros(q_len + chosen_len, dtype=torch.bool)
-        chosen_mask[len(q):] = True
+        chosen_mask[q_len:] = True
 
         rejected = self.tokenizer.encode(f"{rejected_response}<|im_end|>\n<eos>")
+        rejected_len = len(rejected)
+        
         rejected_tokens = torch.tensor(q + rejected, dtype=torch.int32)
-        rejected_mask = torch.zeros_like(rejected_tokens, dtype=torch.bool)
+        rejected_mask = torch.zeros(q_len + rejected_len, dtype=torch.bool)
-        rejected_mask[len(q):] = True
+        rejected_mask[q_len:] = True
 
         return {
             "chosen": chosen_tokens,