refactor: 修改项目结构

.gitignore

@@ -7,6 +7,4 @@
 # Allow specific file types and root files
 !*.py
 !*.md
-!*.png
-!LICENSE
 !pyproject.toml

README.md

@@ -1,174 +1,238 @@
 # DataPipeline
 
-用于训练KHAOSZ模型的数据集处理工具
+用于训练 KHAOSZ 模型的数据集处理工具。提供文本导出、Tokenize、序列打包、H5 存储等独立工具，支持预训练 / SFT / DPO 三种训练范式。
 
 ## 项目结构
 
 ```
-khaosz_dataset/
+pipeline/
-├── modules/                    # 核心模块
+├── tokenizer.py      # BPE 分词器
-│   ├── tokenizer.py           # BPE Tokenizer
+├── text.py           # 文本规范化
-│   └── datapipeline/          # 数据管道模块
+├── packing.py        # 序列打包（bin-packing）
-│       ├── pipeline.py        # 主数据管道
+├── io.py             # 文件/HDF5 读写
-│       ├── processors.py      # 数据处理器（策略模式）
+├── processors.py     # PT / SFT / DPO 处理器
-│       ├── io.py             # 文件IO操作
+├── export.py         # Dataset → JSONL 导出
-│       ├── packing.py        # 序列打包
+└── cache.py          # JSONL → Tokenize → H5 缓存
-│       └── text.py           # 文本规范化
+
-├── tokenizer.json             # Tokenizer配置
+pre_train/                          # 预训练数据处理脚本
-└── pyproject.toml            # 项目依赖
+supervised_finetuning/              # SFT 数据处理脚本
+reforce_learning/                   # DPO 数据处理脚本
 ```
 
-## 架构设计
+## 设计理念
 
-本项目采用模块化设计，应用了多种设计模式：
+每个模块**独立可用、零互相依赖**，调用者按需组合：
 
-### 设计模式
+```
-1. **策略模式** - 不同类型的数据处理器（PT/SFT/DPO）
+HuggingFace Hub
-2. **工厂模式** - 处理器工厂统一创建实例
+      │
-3. **模板方法模式** - 数据管道流程标准化
+      ▼  load_dataset()
-
+  DatasetDict
-### 核心组件
+      │
-
+      ▼  export_dataset()          ← pipeline/export.py
-#### DataPipeline（数据管道）
+  JSONL 文件
-主数据管道，负责数据集的分块处理、格式转换和存储。
+      │
-
+      ▼  cache_jsonl()            ← pipeline/cache.py
-```python
+      │     ├─ Processor.process()    ← pipeline/processors.py
-from modules.datapipeline import DataPipeline
+      │     ├─ SequencePacker.pack()  ← pipeline/packing.py
-
+      │     └─ IOHandler.save_h5()    ← pipeline/io.py
-pipeline = DataPipeline(output_dir="./dataset")
+  HDF5 张量文件
-pipeline.process_dataset(
-    dataset_dict=dataset,
-    output_subdir="my-data"
-)
 ```
 
-#### ProcessorFactory（处理器工厂）
+> 各阶段之间通过磁盘文件解耦。你可以只执行阶段 1（导出 JSONL），也可以继续执行阶段 2（tokenize 并缓存为 H5），按需选择。
-创建不同类型的数据处理器。
 
-```python
+## 快速开始
-from modules.datapipeline import ProcessorFactory
-from modules.tokenizer import BpeTokenizer
-
-tokenizer = BpeTokenizer("tokenizer.json")
-
-# 创建预训练处理器
-processor = ProcessorFactory.create("pt", tokenizer)
-
-# 创建SFT处理器
-processor = ProcessorFactory.create("sft", tokenizer)
-```
-
-#### TextNormalizer（文本规范化）
-文本预处理和规范化。
-
-```python
-from modules.datapipeline import TextNormalizer
-
-normalizer = TextNormalizer()
-normalized_text = normalizer.normalize(text)
-```
-
-## 使用说明
 
 ### 安装依赖
 
 ```bash
-pip install datasets tokenizers tqdm torch h5py
+pip install -e .
 ```
 
-### 运行数据处理
+### 阶段 1：导出数据集为 JSONL
-
-运行所有数据处理脚本：
-
-```bash
-python run.py
-```
-
-运行特定脚本：
-
-```bash
-# 预训练数据处理
-python pre_train/english-wiki.py
-
-# SFT数据处理
-python supervised_finetuning/sft_belle.py
-
-# DPO数据处理
-python reforce_learning/dpp_chinese_dpo_pairs.py
-```
-
-### 自定义数据处理
-
-#### 1. 基础数据处理
 
 ```python
 from datasets import load_dataset
-from modules.datapipeline import DataPipeline
+from pipeline.export import export_dataset
 
-# 加载数据集
 dataset = load_dataset("your-dataset")
-
+export_dataset(
-# 创建管道
+    dataset=dataset["train"],        # 直接传 Dataset，不传 DatasetDict
-pipeline = DataPipeline()
+    output_dir="./dataset",
-
+    output_prefix="my-data",
-# 处理数据
+    max_chunks=5,                    # 可选，限制 chunk 数量（调试用）
-pipeline.process_dataset(
-    dataset_dict=dataset,
-    output_subdir="output-dir",
-    process_func=lambda x: {"text": x["content"]}
 )
 ```
 
-#### 2. 自定义处理器
+**自定义转换函数：**
 
 ```python
-from modules.datapipeline.processors import BaseProcessor
+def process_func(example):
-from modules.datapipeline import ProcessorFactory
+    # 提取字段、转换格式、展开多轮对话等
+    return {"query": example["instruction"], "response": example["output"]}
+
+export_dataset(
+    dataset=dataset["train"],
+    output_dir="./dataset",
+    output_prefix="my-sft",
+    process_func=process_func,
+)
+```
+
+> `process_func` 返回单个 `dict` 或 `list[dict]`（一条样本可展开为多条）。
+
+**使用文本规范化：**
+
+```python
+from pipeline.text import TextNormalizer
+
+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)
+
+cache_jsonl(
+    files=["./dataset/my-pretrain_chunk_0.jsonl"],
+    output_dir="./cached",
+    processor=processor,
+    pack_size=4096,       # 可选，打包长度；<=0 不打包
+    pad_value=1,
+)
+```
+
+**处理器类型：**
+
+| 类型 | 工厂 key | 输入格式 | 输出 keys |
+|------|----------|---------|-----------|
+| 预训练 | `"pt"` | `{"text": "..."}` | `["sequence"]` |
+| SFT | `"sft"` | `{"query": "...", "response": "..."}` | `["sequence", "loss_mask"]` |
+| DPO | `"dpo"` | 待定 | `["chosen", "chosen_mask", "rejected", "rejected_mask"]` |
+
+## 独立工具参考
+
+### BpeTokenizer
+
+```python
+from pipeline.tokenizer import BpeTokenizer
+
+tokenizer = BpeTokenizer("tokenizer.json")
+ids = tokenizer.encode("hello world")           # → [1234, 5678, 1]
+text = tokenizer.decode(ids)                     # → "hello world"
+len(tokenizer)                                   # → 词表大小
+```
+
+### TextNormalizer
+
+```python
+from pipeline.text import TextNormalizer
+
+normalizer = TextNormalizer()
+text = normalizer.normalize(raw_text)
+```
+
+替换规则包括：全角引号 → 半角、各种短横线统一、不间断空格 → 普通空格等。支持自定义规则：
+
+```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
+```
+
+### IOHandler
+
+```python
+from pipeline.io import IOHandler
+
+# 保存
+IOHandler.save_h5("./output", "my_data", {"sequence": [tensor1, tensor2]})
+
+# 加载
+data = IOHandler.load_h5("./output")     # → {"sequence": [tensor1, tensor2, ...]}
+
+# 遍历文件
+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: dict) -> dict:
+    def process(self, input_dict):
-        # 自定义处理逻辑
+        tokens = self.tokenizer.encode(input_dict["text"])
-        return {"processed": data}
+        return {"sequence": torch.tensor(tokens, dtype=torch.int32)}
 
     @property
-    def output_keys(self) -> list:
+    def output_keys(self):
-        return ["processed"]
+        return ["sequence"]
 
-# 注册处理器
 ProcessorFactory.register("my_type", MyProcessor)
 ```
 
-#### 3. 文本规范化
+## 运行脚本
 
-```python
+```bash
-from modules.datapipeline import TextNormalizer
+# 预训练
+python pre_train/chinese-c4.py
+python pre_train/english-wiki.py
 
-# 使用默认规则
+# SFT
-normalizer = TextNormalizer()
+python supervised_finetuning/sft_belle.py
-text = normalizer.normalize(text)
+python supervised_finetuning/sft_coder.py
 
-# 自定义规则
+# DPO
-custom_rules = {"旧词": "新词"}
+python reforce_learning/dpp_chinese_dpo_pairs.py
-normalizer = TextNormalizer(custom_rules)
 ```
 
-## 数据输出格式
+## 输出格式
 
-### JSONL格式
+**JSONL**（阶段 1 输出）：
-每个数据块保存为JSONL文件：
+
-```
+```jsonl
 {"text": "训练文本内容..."}
 {"query": "问题", "response": "答案"}
 ```
 
-### H5格式
+**HDF5**（阶段 2 输出）：
-打包后的张量数据保存为HDF5格式，支持高效加载：
-```python
-from modules.datapipeline import IOHandler
 
-# 加载H5数据
+```
-data = IOHandler.load_h5("./cached_data")
+my_data.h5
+├── sequence/
+│   ├── data_0    # Tensor (4096,) int32
+│   ├── data_1    # Tensor (4096,) int32
+│   └── ...
+└── loss_mask/     # 仅 SFT
+    ├── data_0    # Tensor (4096,) bool
+    └── ...
 ```

modules/__init__.py

@@ -1,17 +0,0 @@
-from .tokenizer import BpeTokenizer
-from .datapipeline import (
-    DataPipeline,
-    ProcessorFactory,
-    IOHandler,
-    TextNormalizer,
-    SequencePacker
-)
-
-__all__ = [
-    'BpeTokenizer',
-    'DataPipeline',
-    'ProcessorFactory',
-    'IOHandler',
-    'TextNormalizer',
-    'SequencePacker'
-]

modules/datapipeline/__init__.py

@@ -1,13 +0,0 @@
-from .pipeline import DataPipeline
-from .processors import ProcessorFactory
-from .io import IOHandler
-from .text import TextNormalizer
-from .packing import SequencePacker
-
-__all__ = [
-    'DataPipeline',
-    'ProcessorFactory', 
-    'IOHandler',
-    'TextNormalizer',
-    'SequencePacker'
-]

modules/datapipeline/pipeline.py

@@ -1,252 +0,0 @@
-from typing import Dict, List, Callable, Union
-from datasets import DatasetDict
-from tqdm import tqdm
-import torch
-import json
-import os
-
-from .processors import ProcessorFactory
-from .packing import SequencePacker
-from .io import IOHandler
-from .text import TextNormalizer
-
-
-class DataPipeline:
-    """数据处理管道 - 模板方法模式"""
-    
-    def __init__(self, output_dir: str = None):
-        self.output_dir = output_dir or os.path.join(os.getcwd(), "dataset")
-    
-    def process_dataset(
-        self,
-        dataset_dict: DatasetDict,
-        output_subdir: str,
-        max_chunk_num: int = None,
-        chunk_size: int = 1000000,
-        split_name: str = "train",
-        column_name: str = "text",
-        process_func: Union[Callable[[dict], dict], Callable[[List[dict]], List[dict]]] = None,
-        normalization_func: Callable[[str], str] = None,
-        output_dir: str = None,
-    ) -> None:
-        """处理数据集的主流程"""
-        dataset = dataset_dict[split_name]
-        total_samples = len(dataset)
-        num_chunks = (total_samples // chunk_size) + 1
-        lim_chunks = min(max_chunk_num, num_chunks) if max_chunk_num else num_chunks
-        
-        output_dir = output_dir or os.path.join(self.output_dir, output_subdir)
-        os.makedirs(output_dir, exist_ok=True)
-        
-        # 处理每个数据块
-        for i in range(lim_chunks):
-            self._process_chunk(
-                dataset=dataset,
-                chunk_idx=i,
-                chunk_size=chunk_size,
-                total_samples=total_samples,
-                output_dir=output_dir,
-                output_subdir=output_subdir,
-                column_name=column_name,
-                process_func=process_func,
-                normalization_func=normalization_func
-            )
-    
-    def _process_chunk(
-        self,
-        dataset,
-        chunk_idx: int,
-        chunk_size: int,
-        total_samples: int,
-        output_dir: str,
-        output_subdir: str,
-        column_name: str,
-        process_func,
-        normalization_func
-    ) -> None:
-        """处理单个数据块"""
-        start_idx = chunk_idx * chunk_size
-        end_idx = min((chunk_idx + 1) * chunk_size, total_samples)
-        chunk = dataset.select(range(start_idx, end_idx))
-        
-        output_path = os.path.join(output_dir, f"{output_subdir}_text_chunk_{chunk_idx}.jsonl")
-        
-        with open(output_path, "w", encoding="utf-8") as f:
-            for example in chunk:
-                processed = self._process_example(
-                    example, column_name, process_func, normalization_func
-                )
-                self._write_processed(f, processed)
-        
-        print(f"Saved text chunk {chunk_idx} to {output_path}")
-    
-    def _process_example(
-        self,
-        example: dict,
-        column_name: str,
-        process_func,
-        normalization_func
-    ) -> Union[dict, List[dict]]:
-        """处理单个样本"""
-        if process_func:
-            return process_func(example)
-        
-        text = example[column_name]
-        if normalization_func:
-            text = normalization_func(text)
-        return {column_name: text}
-    
-    def _write_processed(self, file, processed) -> None:
-        """写入处理后的数据"""
-        if isinstance(processed, dict):
-            file.write(json.dumps(processed, ensure_ascii=False) + "\n")
-        elif isinstance(processed, list):
-            for item in processed:
-                file.write(json.dumps(item, ensure_ascii=False) + "\n")
-    
-    def cache_files(
-        self,
-        tokenizer,
-        files: List[str],
-        base_out_dir: str,
-        cache_type: str,
-        packing_size: int = -1,
-        pad_value: int = 1
-    ) -> None:
-        """缓存文件到H5格式"""
-        processor = ProcessorFactory.create(cache_type, tokenizer)
-        self.dump_files(
-            files=files,
-            base_out_dir=base_out_dir,
-            process_func=processor.process,
-            output_keys=processor.output_keys,
-            packing_size=packing_size,
-            pad_value=pad_value
-        )
-    
-    def dump_files(
-        self,
-        files: List[str],
-        base_out_dir: str,
-        process_func: Callable[[dict], dict],
-        output_keys: List[str],
-        packing_size: int = -1,
-        pad_value: int = 0
-    ) -> None:
-        """转储文件到H5格式"""
-        for file_path in files:
-            self._dump_single_file(
-                file_path=file_path,
-                base_out_dir=base_out_dir,
-                process_func=process_func,
-                output_keys=output_keys,
-                packing_size=packing_size,
-                pad_value=pad_value
-            )
-    
-    def _dump_single_file(
-        self,
-        file_path: str,
-        base_out_dir: str,
-        process_func: Callable[[dict], dict],
-        output_keys: List[str],
-        packing_size: int,
-        pad_value: int
-    ) -> None:
-        """转储单个文件"""
-        os.makedirs(base_out_dir, exist_ok=True)
-        file_name = os.path.basename(file_path)
-        out_file_name = file_name.split(".")[0]
-        
-        # 读取和处理数据
-        with open(file_path, "r") as f:
-            lines = f.readlines()
-        
-        arrows: List[Dict[str, torch.Tensor]] = []
-        for line in tqdm(lines, desc=f"Processing {file_name}", leave=False):
-            line_dict = json.loads(line)
-            arrow = process_func(line_dict)
-            if arrow is not None:
-                arrows.append(arrow)
-        
-        # 组织输出数据
-        package: Dict[str, List[torch.Tensor]] = {
-            key: [arrow[key] for arrow in arrows]
-            for key in output_keys
-        }
-        
-        # 打包序列（如果需要）
-        output_package = {}
-        for key in output_keys:
-            if packing_size > 0:
-                print(f"Packaging key: '{key}'")
-                packer = SequencePacker(packing_size, pad_value)
-                output_package[key] = packer.pack(package[key])
-            else:
-                output_package[key] = package[key]
-        
-        # 保存到H5
-        IOHandler.save_h5(base_out_dir, out_file_name, output_package)
-
-
-# 向后兼容的函数接口
-def process_dataset(
-    dataset_dict: DatasetDict,
-    output_subdir: str,
-    max_chunk_num: int = None,
-    chunk_size: int = 1000000,
-    split_name: str = "train",
-    column_name: str = "text",
-    process_func: Union[Callable[[dict], dict], Callable[[List[dict]], List[dict]]] = None,
-    normalization_func: Callable[[str], str] = None,
-    output_dir: str = None,
-) -> None:
-    """向后兼容的函数接口"""
-    pipeline = DataPipeline(output_dir)
-    normalizer = TextNormalizer() if normalization_func is None else None
-    norm_func = normalization_func or (normalizer.normalize if normalizer else None)
-    
-    return pipeline.process_dataset(
-        dataset_dict=dataset_dict,
-        output_subdir=output_subdir,
-        max_chunk_num=max_chunk_num,
-        chunk_size=chunk_size,
-        split_name=split_name,
-        column_name=column_name,
-        process_func=process_func,
-        normalization_func=norm_func,
-        output_dir=output_dir
-    )
-
-
-def cache_files(tokenizer, files, base_out_dir, cache_type, packing_size: int = -1, pad_value: int = 1):
-    """向后兼容的函数接口"""
-    pipeline = DataPipeline()
-    return pipeline.cache_files(
-        tokenizer=tokenizer,
-        files=files,
-        base_out_dir=base_out_dir,
-        cache_type=cache_type,
-        packing_size=packing_size,
-        pad_value=pad_value
-    )
-
-
-def dump_files(
-    files: List[str],
-    base_out_dir: str,
-    process_func: Callable[[dict], dict],
-    output_keys: List[str],
-    packing_size: int = -1,
-    pad_value: int = 0
-):
-    """向后兼容的函数接口"""
-    pipeline = DataPipeline()
-    return pipeline.dump_files(
-        files=files,
-        base_out_dir=base_out_dir,
-        process_func=process_func,
-        output_keys=output_keys,
-        packing_size=packing_size,
-        pad_value=pad_value
-    )

pipeline/__init__.py

@@ -0,0 +1,18 @@
+from .tokenizer import BpeTokenizer
+from .text import TextNormalizer
+from .packing import SequencePacker
+from .io import IOHandler
+from .processors import ProcessorFactory, BaseProcessor
+from .export import export_dataset
+from .cache import cache_jsonl
+
+__all__ = [
+    'BpeTokenizer',
+    'TextNormalizer',
+    'SequencePacker',
+    'IOHandler',
+    'ProcessorFactory',
+    'BaseProcessor',
+    'export_dataset',
+    'cache_jsonl',
+]

pipeline/cache.py

@@ -0,0 +1,64 @@
+"""将 JSONL 文件 tokenize 后打包存储为 HDF5"""
+import json
+import os
+from typing import List
+from pathlib import Path
+
+from tqdm import tqdm
+
+from .processors import BaseProcessor
+from .packing import SequencePacker
+from .io import IOHandler
+
+
+def cache_jsonl(
+    files: List[str],
+    output_dir: str,
+    processor: BaseProcessor,
+    *,
+    pack_size: int = -1,
+    pad_value: int = 1,
+) -> List[str]:
+    """
+    将 JSONL 文件 tokenize 后打包存储为 HDF5。
+
+    Args:
+        files: JSONL 文件路径列表
+        output_dir: H5 输出目录
+        processor: 已初始化的 Processor 实例
+        pack_size: 打包长度，<=0 表示不打包
+        pad_value: 填充值
+
+    Returns:
+        生成的 H5 文件路径列表
+    """
+    os.makedirs(output_dir, exist_ok=True)
+    output_files: List[str] = []
+
+    for file_path in files:
+        file_name = Path(file_path).stem
+
+        with open(file_path, "r", encoding="utf-8") as f:
+            lines = f.readlines()
+
+        arrows = []
+        for line in tqdm(lines, desc=f"Processing {file_name}", leave=False):
+            arrow = processor.process(json.loads(line))
+            if arrow is not None:
+                arrows.append(arrow)
+
+        package = {key: [a[key] for a in arrows] for key in processor.output_keys}
+
+        output = {}
+        for key in processor.output_keys:
+            if pack_size > 0:
+                output[key] = SequencePacker(pack_size, pad_value).pack(package[key])
+            else:
+                output[key] = package[key]
+
+        IOHandler.save_h5(output_dir, file_name, output)
+        h5_path = os.path.join(output_dir, f"{file_name}.h5")
+        output_files.append(h5_path)
+        print(f"Saved {h5_path}")
+
+    return output_files

pipeline/export.py

@@ -0,0 +1,54 @@
+"""将 HuggingFace Dataset 分块导出为 JSONL 文件"""
+import json
+import os
+from typing import Callable, Optional, List, Union
+
+
+def export_dataset(
+    dataset,
+    output_dir: str,
+    output_prefix: str,
+    *,
+    chunk_size: int = 1_000_000,
+    max_chunks: Optional[int] = None,
+    process_func: Optional[Callable] = None,
+    column: str = "text",
+) -> List[str]:
+    """
+    将 HuggingFace Dataset 分块导出为 JSONL 文件。
+
+    Args:
+        dataset: HuggingFace Dataset 对象
+        output_dir: 输出目录
+        output_prefix: 输出文件名前缀，如 "chinese-c4-pretrain"
+        chunk_size: 每个文件的最大样本数
+        max_chunks: 最多处理几个 chunk（用于调试）
+        process_func: 单条样本的转换函数 (dict) -> dict | list[dict]
+        column: 默认提取的文本列名（仅在 process_func 为 None 时使用）
+
+    Returns:
+        生成的文件路径列表
+    """
+    os.makedirs(output_dir, exist_ok=True)
+    total = len(dataset)
+    num_chunks = (total + chunk_size - 1) // chunk_size
+    lim = min(max_chunks, num_chunks) if max_chunks else num_chunks
+
+    output_files: List[str] = []
+    for i in range(lim):
+        start = i * chunk_size
+        end = min(start + chunk_size, total)
+        chunk = dataset.select(range(start, end))
+
+        path = os.path.join(output_dir, f"{output_prefix}_chunk_{i}.jsonl")
+        with open(path, "w", encoding="utf-8") as f:
+            for example in chunk:
+                processed = process_func(example) if process_func else {column: example[column]}
+                items = processed if isinstance(processed, list) else [processed]
+                for item in items:
+                    f.write(json.dumps(item, ensure_ascii=False) + "\n")
+
+        output_files.append(path)
+        print(f"[{i + 1}/{lim}] Saved {path}")
+
+    return output_files

pipeline/io.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Dict, List, Union
+from typing import Dict, List
 import os
 import h5py
 import torch
@@ -7,11 +7,10 @@ from torch import Tensor
 
 
 class IOHandler:
-    """文件和H5数据存储处理器"""
+    """文件和 HDF5 读写"""
 
     @staticmethod
     def fetch_files(directory: str) -> List[str]:
-        """获取目录下所有文件"""
         return [
             os.path.join(root, f)
             for root, _, files in os.walk(directory)
@@ -20,7 +19,6 @@ class IOHandler:
 
     @staticmethod
     def fetch_folders(root_dir: str, filter_func=None) -> List[str]:
-        """获取目录下所有文件夹"""
         folders = []
         for root, dirs, _ in os.walk(root_dir):
             for dir_name in dirs:
@@ -31,10 +29,8 @@ class IOHandler:
 
     @staticmethod
     def save_h5(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]) -> None:
-        """保存张量组到H5文件"""
         os.makedirs(file_path, exist_ok=True)
         full_path = os.path.join(file_path, f"{file_name}.h5")
-        
         with h5py.File(full_path, 'w') as f:
             for key, tensors in tensor_group.items():
                 grp = f.create_group(key)
@@ -43,11 +39,9 @@ class IOHandler:
 
     @staticmethod
     def load_h5(file_path: str, share_memory: bool = True) -> Dict[str, List[Tensor]]:
-        """从H5文件加载张量组"""
         tensor_group: Dict[str, List[Tensor]] = {}
         root_path = Path(file_path)
         h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5"))
-        
         for h5_file in h5_files:
             with h5py.File(h5_file, 'r') as f:
                 for key in f.keys():
@@ -59,19 +53,4 @@ class IOHandler:
                         for dset in [grp[dset_name]]
                     ]
                     tensor_group.setdefault(key, []).extend(tensors)
-        
         return tensor_group
-
-
-# 向后兼容的函数接口
-def fetch_files(directory: str) -> List[str]:
-    return IOHandler.fetch_files(directory)
-
-def fetch_folders(root_dir: str, filter_func=None) -> List[str]:
-    return IOHandler.fetch_folders(root_dir, filter_func)
-
-def save_h5(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]) -> None:
-    return IOHandler.save_h5(file_path, file_name, tensor_group)
-
-def load_h5(file_path: str, share_memory: bool = True) -> Dict[str, List[Tensor]]:
-    return IOHandler.load_h5(file_path, share_memory)

pipeline/packing.py

@@ -4,14 +4,13 @@ from torch import Tensor
 
 
 class SequencePacker:
-    """序列打包策略"""
+    """序列打包（bin-packing）"""
 
     def __init__(self, pack_size: int, pad_value: int = 0):
         self.pack_size = pack_size
         self.pad_value = pad_value
 
     def pack(self, sequences: List[Tensor]) -> List[Tensor]:
-        """打包序列到固定大小"""
         packages = []
         sequences.sort(key=lambda x: x.numel(), reverse=True)
 
@@ -34,8 +33,3 @@ class SequencePacker:
             packages.append(current_pack)
 
         return packages
-
-
-def pack_sequences(sequences: List[Tensor], pack_size: int, pad_value: int) -> List[Tensor]:
-    """向后兼容的函数接口"""
-    return SequencePacker(pack_size, pad_value).pack(sequences)

pipeline/processors.py

@@ -1,21 +1,18 @@
 from abc import ABC, abstractmethod
-from typing import Dict, List, Callable
+from typing import Dict, List
 import torch
-from torch import Tensor
 
 
 class BaseProcessor(ABC):
-    """处理器抽象基类 - 策略模式"""
+    """处理器抽象基类"""
 
     @abstractmethod
     def process(self, input_dict: dict) -> dict:
-        """处理单个数据项"""
         pass
 
     @property
     @abstractmethod
     def output_keys(self) -> List[str]:
-        """输出字段列表"""
         pass
 
 
@@ -47,11 +44,9 @@ class SFTProcessor(BaseProcessor):
             f"<|im_start|>user\n{query}<|im_end|>\n<|im_start|>assistant\n"
         )
         a = self.tokenizer.encode(f"{response}<|im_end|>\n<eos>")
-        
         tokens = torch.tensor(q + a, dtype=torch.int32)
         loss_mask = torch.zeros_like(tokens, dtype=torch.bool)
         loss_mask[len(q):] = True
-        
         return {"sequence": tokens, "loss_mask": loss_mask}
 
     @property
@@ -60,13 +55,13 @@ class SFTProcessor(BaseProcessor):
 
 
 class DPOProcessor(BaseProcessor):
-    """DPO偏好学习数据处理器"""
+    """DPO 偏好学习数据处理器"""
 
     def __init__(self, tokenizer):
         self.tokenizer = tokenizer
 
     def process(self, input_dict: dict) -> dict:
-        # TODO: 实现DPO处理逻辑
+        # TODO: 实现 DPO 处理逻辑
         return None
 
     @property
@@ -75,7 +70,7 @@ class DPOProcessor(BaseProcessor):
 
 
 class ProcessorFactory:
-    """处理器工厂 - 工厂模式"""
+    """处理器工厂"""
 
     _processors = {
         "pt": PreTrainProcessor,
@@ -85,26 +80,10 @@ class ProcessorFactory:
 
     @classmethod
     def create(cls, processor_type: str, tokenizer) -> BaseProcessor:
-        """创建处理器实例"""
         if processor_type not in cls._processors:
             raise ValueError(f"Invalid processor type: {processor_type}")
         return cls._processors[processor_type](tokenizer)
 
     @classmethod
     def register(cls, processor_type: str, processor_class: type):
-        """注册新的处理器类型"""
         cls._processors[processor_type] = processor_class
-
-
-# 向后兼容的函数接口
-def get_pt_processor(tokenizer) -> Callable:
-    processor = PreTrainProcessor(tokenizer)
-    return processor.process
-
-def get_sft_processor(tokenizer) -> Callable:
-    processor = SFTProcessor(tokenizer)
-    return processor.process
-
-def get_dpo_processor(tokenizer) -> Callable:
-    processor = DPOProcessor(tokenizer)
-    return processor.process

pipeline/text.py

@@ -3,7 +3,7 @@ from typing import Dict
 
 
 class TextNormalizer:
-    """文本规范化策略"""
+    """文本规范化"""
 
     DEFAULT_REPLACEMENTS = {
         "\\[": "$$", "\\]": "$$", "\\(": "$", "\\)": "$",
@@ -18,10 +18,4 @@ class TextNormalizer:
         self._pattern = re.compile('|'.join(re.escape(k) for k in self.replacements))
 
     def normalize(self, text: str) -> str:
-        """规范化文本"""
         return self._pattern.sub(lambda m: self.replacements[m.group()], text)
-
-
-def comprehensive_normalization(text: str) -> str:
-    """向后兼容的函数接口"""
-    return TextNormalizer().normalize(text)

pre_train/chinese-c4.py

@@ -1,11 +0,0 @@
-from datasets import load_dataset
-from modules.datapipeline import DataPipeline
-
-if __name__ == "__main__":
-    dataset = load_dataset("shjwudp/chinese-c4")
-    
-    pipeline = DataPipeline()
-    pipeline.process_dataset(
-        dataset_dict=dataset,
-        output_subdir="chinese-c4-pretrain"
-    )

pre_train/chinese-cosmopedia.py

@@ -1,17 +0,0 @@
-from datasets import load_dataset
-from modules.datapipeline import DataPipeline
-
-if __name__ == "__main__":
-    chunk_size = 1000000
-    
-    dataset = load_dataset(
-        "opencsg/chinese-cosmopedia", 
-        data_files={"train": [f"data/000{i:02d}.parquet" for i in range(25)]}
-    )
-    
-    pipeline = DataPipeline()
-    pipeline.process_dataset(
-        dataset_dict=dataset,
-        output_subdir="chinese-wiki-pretrain",
-        chunk_size=chunk_size,
-    )

pre_train/english-fineweb.py

@@ -1,11 +0,0 @@
-from datasets import load_dataset
-from modules.datapipeline import DataPipeline
-
-if __name__ == "__main__":
-    dataset = load_dataset("HuggingFaceFW/fineweb", "sample-10BT")
-    
-    pipeline = DataPipeline()
-    pipeline.process_dataset(
-        dataset_dict=dataset,
-        output_subdir="english-fineweb-pretrain",
-    )

pre_train/english-wiki.py

@@ -1,12 +0,0 @@
-from datasets import load_dataset
-from modules.datapipeline import DataPipeline
-
-if __name__ == "__main__":
-    dataset = load_dataset("Blaze7451/enwiki_structured_content")
-    
-    pipeline = DataPipeline()
-    pipeline.process_dataset(
-        dataset_dict=dataset,
-        output_subdir="english-wiki-pretrain",
-        max_chunk_num=5,
-    )

reforce_learning/dpp_chinese_dpo_pairs.py

@@ -1,20 +0,0 @@
-from datasets import load_dataset
-from modules.datapipeline import DataPipeline
-
-def process_func(input_dict: dict):
-    return {
-        "promopt": input_dict["prompt"],
-        "chosen": input_dict["chosen"],
-        "rejected": input_dict["rejected"]
-    }
-
-
-if __name__ == "__main__":
-    dataset = load_dataset("wenbopan/Chinese-dpo-pairs")
-    
-    pipeline = DataPipeline()
-    pipeline.process_dataset(
-        dataset_dict=dataset,
-        output_subdir="Chinese-dpo-pairs",
-        process_func=process_func
-    )

scripts/pre_train/chinese-c4.py

@@ -0,0 +1,10 @@
+from datasets import load_dataset
+from pipeline import export_dataset
+
+if __name__ == "__main__":
+    dataset = load_dataset("shjwudp/chinese-c4")
+    export_dataset(
+        dataset=dataset["train"],
+        output_dir="./dataset",
+        output_prefix="chinese-c4-pretrain",
+    )

scripts/pre_train/chinese-cosmopedia.py

@@ -0,0 +1,14 @@
+from datasets import load_dataset
+from pipeline import export_dataset
+
+if __name__ == "__main__":
+    dataset = load_dataset(
+        "opencsg/chinese-cosmopedia",
+        data_files={"train": [f"data/000{i:02d}.parquet" for i in range(25)]}
+    )
+    export_dataset(
+        dataset=dataset["train"],
+        output_dir="./dataset",
+        output_prefix="chinese-cosmopedia-pretrain",
+        chunk_size=1_000_000,
+    )

scripts/pre_train/english-fineweb.py

@@ -0,0 +1,10 @@
+from datasets import load_dataset
+from pipeline import export_dataset
+
+if __name__ == "__main__":
+    dataset = load_dataset("HuggingFaceFW/fineweb", "sample-10BT")
+    export_dataset(
+        dataset=dataset["train"],
+        output_dir="./dataset",
+        output_prefix="english-fineweb-pretrain",
+    )

scripts/pre_train/english-wiki.py

@@ -0,0 +1,11 @@
+from datasets import load_dataset
+from pipeline import export_dataset
+
+if __name__ == "__main__":
+    dataset = load_dataset("Blaze7451/enwiki_structured_content")
+    export_dataset(
+        dataset=dataset["train"],
+        output_dir="./dataset",
+        output_prefix="english-wiki-pretrain",
+        max_chunks=5,
+    )

scripts/reforce_learning/dpp_chinese_dpo_pairs.py

@@ -0,0 +1,20 @@
+from datasets import load_dataset
+from pipeline import export_dataset
+
+
+def process_func(input_dict: dict):
+    return {
+        "prompt": input_dict["prompt"],
+        "chosen": input_dict["chosen"],
+        "rejected": input_dict["rejected"],
+    }
+
+
+if __name__ == "__main__":
+    dataset = load_dataset("wenbopan/Chinese-dpo-pairs")
+    export_dataset(
+        dataset=dataset["train"],
+        output_dir="./dataset",
+        output_prefix="Chinese-dpo-pairs",
+        process_func=process_func,
+    )

scripts/supervised_finetuning/sft_belle.py

@@ -0,0 +1,23 @@
+from datasets import load_dataset
+from pipeline import export_dataset
+
+
+def process_func(input_dict: dict):
+    conversations = input_dict["conversations"]
+    n = len(conversations) // 2
+    examples = []
+    for i in range(n):
+        user_msg = conversations[2 * i]["value"]
+        assistant_msg = conversations[2 * i + 1]["value"]
+        examples.append({"query": user_msg, "response": assistant_msg})
+    return examples
+
+
+if __name__ == "__main__":
+    dataset = load_dataset("BelleGroup/train_3.5M_CN")
+    export_dataset(
+        dataset=dataset["train"],
+        output_dir="./dataset",
+        output_prefix="belle-sft",
+        process_func=process_func,
+    )

scripts/supervised_finetuning/sft_chinese_instruct.py

@@ -0,0 +1,31 @@
+from datasets import load_dataset, concatenate_datasets
+from pipeline import export_dataset, TextNormalizer
+
+normalizer = TextNormalizer()
+
+
+def process_func(input_dict: dict):
+    query = input_dict["prompt"] if input_dict["prompt"] else ""
+    resp = input_dict["response"] if input_dict["response"] else ""
+    return {"query": normalizer.normalize(query), "response": normalizer.normalize(resp)}
+
+
+if __name__ == "__main__":
+    all_data = [
+        'stem_zh', 'infinity-instruct', 'firefly', 'magpie', 'dpsk-r1-distil',
+        'coig-cqia', 'disc-law', 'neo_sft_phase2', 'chinese-medical', 'chinese-reasoning-distil',
+        'psycho-10k-dpsk-r1', 'sof-c-zh', 'industryinstruction', 'Chinese-QA-AFAF',
+    ]
+
+    datasets = []
+    for subset in all_data:
+        ds = load_dataset("Mxode/Chinese-Instruct", name=subset)
+        datasets.append(ds["train"])
+
+    combined_dataset = concatenate_datasets(datasets)
+    export_dataset(
+        dataset=combined_dataset,
+        output_dir="./dataset",
+        output_prefix="chinese-instruct-sft",
+        process_func=process_func,
+    )

scripts/supervised_finetuning/sft_coder.py

@@ -1,6 +1,5 @@
-# inclusionAI/Ling-Coder-SFT
 from datasets import load_dataset
-from modules.datapipeline import DataPipeline
+from pipeline import export_dataset
 
 
 def process_func(input_dict: dict) -> dict:
@@ -12,11 +11,9 @@ def process_func(input_dict: dict) -> dict:
 
 if __name__ == "__main__":
     dataset = load_dataset("inclusionAI/Ling-Coder-SFT")
-    
+    export_dataset(
-    pipeline = DataPipeline()
+        dataset=dataset["train"],
-    pipeline.process_dataset(
+        output_dir="./dataset",
-        dataset_dict=dataset,
+        output_prefix="Ling-Coder-sft",
-        output_subdir="Ling-Coder-sft",
+        process_func=process_func,
-        process_func=process_func
     )
-    

scripts/supervised_finetuning/sft_firefly-1.1m-rephrased.py

@@ -0,0 +1,16 @@
+from datasets import load_dataset
+from pipeline import export_dataset
+
+
+def process_func(input_dict: dict):
+    return {"query": input_dict["instruction"], "response": input_dict["output"]}
+
+
+if __name__ == "__main__":
+    dataset = load_dataset("Mxode/Firefly-1.1M-Rephrased")
+    export_dataset(
+        dataset=dataset["train"],
+        output_dir="./dataset",
+        output_prefix="Firefly-1.1M-Rephrased",
+        process_func=process_func,
+    )

scripts/supervised_finetuning/sft_magpie-pro-300k.py

@@ -0,0 +1,24 @@
+from datasets import load_dataset
+from pipeline import export_dataset
+
+
+def process_func(input_dict: dict):
+    conversations = input_dict["conversations"]
+    assert len(conversations) % 2 == 0
+    n = len(conversations) // 2
+    examples = []
+    for i in range(n):
+        user_msg = conversations[2 * i]["value"]
+        assistant_msg = conversations[2 * i + 1]["value"]
+        examples.append({"query": user_msg, "response": assistant_msg})
+    return examples
+
+
+if __name__ == "__main__":
+    dataset = load_dataset("HuggingFaceTB/Magpie-Pro-300K-Filtered-H4")
+    export_dataset(
+        dataset=dataset["train_sft"],
+        output_dir="./dataset",
+        output_prefix="Magpie-Pro-300K-sft",
+        process_func=process_func,
+    )

supervised_finetuning/sft_belle.py

@@ -1,30 +0,0 @@
-from datasets import load_dataset
-from modules.datapipeline import DataPipeline
-
-
-
-def process_func(input_dict: dict):
-    conversations = input_dict["conversations"]
-    n = len(conversations) // 2
-    examples = []
-    
-    for i in range(n):
-        user_msg = conversations[2*i]["value"]
-        assistant_msg = conversations[2*i+1]["value"]
-        examples.append({
-            "query": user_msg,
-            "response": assistant_msg
-        })
-    
-    return examples
-
-
-if __name__ == "__main__":
-    dataset = load_dataset("BelleGroup/train_3.5M_CN")
-    
-    pipeline = DataPipeline()
-    pipeline.process_dataset(
-        dataset_dict=dataset,
-        output_subdir="belle-sft",
-        process_func=process_func,
-    )

supervised_finetuning/sft_chinese_instruct.py

@@ -1,34 +0,0 @@
-from datasets import DatasetDict
-from datasets import load_dataset, concatenate_datasets
-from modules.datapipeline import DataPipeline, TextNormalizer
-
-
-def process_func(input_dict: dict):
-    query = input_dict["prompt"] if input_dict["prompt"] else ""
-    resp = input_dict["response"] if input_dict["response"] else ""
-    
-    normalizer = TextNormalizer()
-    query = normalizer.normalize(query)
-    resp = normalizer.normalize(resp)
-    
-    return {"query": query, "response": resp }
-
-
-if __name__ == "__main__":
-    all_data = ['stem_zh', 'infinity-instruct', 'firefly', 'magpie', 'dpsk-r1-distil', 
-                'coig-cqia', 'disc-law', 'neo_sft_phase2', 'chinese-medical', 'chinese-reasoning-distil', 
-                'psycho-10k-dpsk-r1', 'sof-c-zh', 'industryinstruction', 'Chinese-QA-AFAF']
-    
-    datasets = []
-    for subset in all_data:
-        ds = load_dataset("Mxode/Chinese-Instruct", name=subset)
-        datasets.append(ds["train"]) 
-    
-    combined_dataset = concatenate_datasets(datasets)
-    
-    pipeline = DataPipeline()
-    pipeline.process_dataset(
-        dataset_dict=DatasetDict({"train": combined_dataset}),
-        output_subdir="chinese-instruct-sft",
-        process_func=process_func,
-    )

supervised_finetuning/sft_firefly-1.1m-rephrased.py

@@ -1,19 +0,0 @@
-# Mxode/Firefly-1.1M-Rephrased
-from datasets import load_dataset
-from modules.datapipeline import DataPipeline
-
-def process_func(input_dict: dict):
-    instruction = input_dict["instruction"]
-    output = input_dict["output"]
-    return {"query": instruction, "response": output}
-
-
-if __name__ == "__main__":
-    dataset = load_dataset("Mxode/Firefly-1.1M-Rephrased")
-    
-    pipeline = DataPipeline()
-    pipeline.process_dataset(
-        dataset_dict=dataset,
-        output_subdir="Firefly-1.1M-Rephrased",
-        process_func=process_func
-    )

supervised_finetuning/sft_magpie-pro-300k.py

@@ -1,32 +0,0 @@
-# HuggingFaceTB/Magpie-Pro-300K-Filtered-H4
-from datasets import load_dataset
-from modules.datapipeline import DataPipeline
-
-
-def process_func(input_dict: dict):
-    conversations = input_dict["conversations"]
-    assert len(conversations) % 2 == 0
-    n = len(conversations) // 2
-    examples = []
-    
-    for i in range(n):
-        user_msg = conversations[2*i]["value"]
-        assistant_msg = conversations[2*i+1]["value"]
-        examples.append({
-            "query": user_msg,
-            "response": assistant_msg
-        })
-    
-    return examples
-
-
-if __name__ == "__main__":
-    dataset = load_dataset("HuggingFaceTB/Magpie-Pro-300K-Filtered-H4")
-    
-    pipeline = DataPipeline()
-    pipeline.process_dataset(
-        dataset_dict=dataset,
-        output_subdir="Magpie-Pro-300K-sft",
-        process_func=process_func,
-        split_name="train_sft",
-    )