ViperEkura/DataPipeline
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

reafactor: 重构项目

Browse Source
This commit is contained in:
ViperEkura
2026-03-30 20:58:51 +08:00
parent f67bad0d8b
commit 35963bcb08
29 changed files with 1395 additions and 1234 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+141 -64
View File
@@ -1,111 +1,188 @@
# DataPipeline
数据集处理工具，支持预训练 / SFT / DPO 三种训练范式。
语言模型训练数据集处理工具包。支持 PT / SFT / DPO 范式。
## 项目结构
```
pipeline/
├── tokenizer.py # BPE 分词器
├── text.py # 文本规范化
├── packing.py # 序列打包
├── io.py # 文件/HDF5 读写
├── processors.py # PT / SFT / DPO 处理器
├── export.py # Dataset → JSONL
└── cache.py # JSONL → Tokenize → H5
├── tokenizer.py # BPE tokenizer
├── text.py # Text normalization
├── packing.py # Sequence packing
├── io.py # File / HDF5 I/O
├── processors.py # PT / SFT / DPO processors
├── export.py # Dataset -> JSONL
├── cache.py # JSONL -> Tokenize -> H5
├── utils.py # Logging, error handling
└── strategies/ # Prompt strategy (strategy pattern)
├── base.py # PromptStrategy ABC
├── chatml.py # ChatML format (configurable tokens)
├── alpaca.py # Alpaca format (configurable tokens)
└── factory.py # StrategyFactory
```
## 设计理念
## 数据流
模块**独立可用**，通过磁盘文件解耦，按需组合：
整体分为两个阶段，通过磁盘 JSONL 文件解耦：
```
Dataset → export_dataset() → JSONL → cache_jsonl() → HDF5
↑
processors.py
packing.py
io.py
Stage 1: Export Dataset Stage 2: Tokenize & Cache
┌──────────────────────────────────────┐ ┌──────────────────────────────────────┐
│ │ │ │
│ HuggingFace Dataset │ │ JSONL files (one dict per line) │
│ │ │ │ │ │
│ ▼ process_func (optional) │ │ ▼ json.loads() │
│ export_dataset() │ │ Processor.process(input_dict) │
│ │ chunk by chunk_size │ │ │ │
│ ▼ │ │ │ ┌───────────────────────┐ │
│ ./dataset/prefix_chunk_0.jsonl ────────────────|> │ PT / SFT / DPO │ │
│ ./dataset/prefix_chunk_1.jsonl │ │ │ │ processor details │ │
│ ... │ │ │ └───────────────────────┘ │
│ │ │ ▼ │
│ modules: export.py │ │ List[Tensor] │
│ │ │ │ │
│ │ │ ▼ SequencePacker (optional) │
│ │ │ Fixed-length packed tensors │
│ │ │ │ │
│ │ │ ▼ IOHandler.save_h5() │
│ │ │ ./cached/chunk_0.h5 │
│ │ │ ./cached/chunk_1.h5 │
│ │ │ ... │
│ │ │ │
│ │ │ modules: io.py, processors.py, │
│ │ │ packing.py, tokenizer.py │
└──────────────────────────────────────┘ └──────────────────────────────────────┘
```
## 使用方法
### Processor 转换规则
### 1. 导出数据集
**PT (Pre-training)**
```
Input: {"text": "Hello world"}
Action: tokenizer.encode(text + "<eos>")
Output: {"sequence": Tensor[int32]}
```
**SFT (Supervised Fine-tuning)** — uses PromptStrategy
```
Input: {"query": "...", "response": "..."}
Action:
1. strategy.build_prompt(input_dict) -> "<|im_start|>user\n...\n<|im_start|>assistant\n"
2. concat response + response_suffix
3. tokenizer.encode full string
4. build loss_mask: query part=False, response part=True
Output: {"sequence": Tensor[int32], "loss_mask": Tensor[bool]}
```
**DPO (Direct Preference Optimization)** — uses PromptStrategy
```
Input: {"query": "...", "chosen": "...", "rejected": "..."}
Action:
1. strategy.build_prompt(input_dict) -> shared query prompt
2. encode chosen = query + response_start + chosen + response_suffix
3. encode rejected = query + response_start + rejected + response_suffix
4. build masks: query part=False, response part=True
Output: {"chosen": Tensor, "chosen_mask": Tensor[bool],
"rejected": Tensor, "rejected_mask": Tensor[bool]}
```
### 序列打包
当 `pack_size > 0` 时启用，使用 First-Fit Decreasing 算法将变长序列填充到固定长度：
- 按序列长度降序排列
- 逐个放入当前包，超长则截断
- 当前包放不下时开新包
- 不足部分用 `pad_value` 填充
## 快速开始
### 1. 导出数据集为 JSONL
```python
from datasets import load_dataset
from pipeline.export import export_dataset
from pipeline import export_dataset
dataset = load_dataset("your-dataset")
export_dataset(
dataset=dataset["train"],
output_dir="./data",
output_prefix="train",
process_func=lambda x: {"text": x["content"]}, # 可选
process_func=lambda x: {"text": x["content"]},
)
```
### 2. Tokenize 并缓存
### 2. 分词并缓存为 HDF5
```python
from pipeline import BpeTokenizer, ProcessorFactory, cache_jsonl
tokenizer = BpeTokenizer("tokenizer.json")
processor = ProcessorFactory.create("pt", tokenizer) # "pt" | "sft" | "dpo"
processor = ProcessorFactory.create("pt", tokenizer)
cache_jsonl(
files=["./data/train.jsonl"],
output_dir="./cached",
processor=processor,
pack_size=4096, # <=0 不打包
pad_value=1,
pack_size=4096,
)
```
### 3. 处理器类型
| 类型 | key | 输入 | 输出 |
|------|-----|------|------|
| 预训练 | `"pt"` | `{"text": "..."}` | `["sequence"]` |
| SFT | `"sft"` | `{"query": "...", "response": "..."}` | `["sequence", "loss_mask"]` |
| DPO | `"dpo"` | `{"query": "...", "chosen": "...", "rejected": "..."}` | `["chosen", "chosen_mask", "rejected", "rejected_mask"]` |
## 参数参考
### export_dataset()
| 参数 | 类型 | 默认 | 说明 |
|------|------|------|------|
| `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" | 默认文本列名 |
### cache_jsonl()
| 参数 | 类型 | 默认 | 说明 |
|------|------|------|------|
| `files` | List[str] | 必填 | JSONL 文件列表 |
| `output_dir` | str | 必填 | 输出目录 |
| `processor` | BaseProcessor | 必填 | 处理器实例 |
| `pack_size` | int | -1 | 打包长度，<=0 不打包 |
| `pad_value` | int | 1 | 填充值 |
### SequencePacker
### 3. 使用策略模式（默认 token）
```python
packer = SequencePacker(pack_size=4096, pad_value=0)
packed = packer.pack([tensor1, tensor2, ...]) # → List[Tensor]
from pipeline import StrategyFactory, ProcessorFactory
strategy = StrategyFactory.create("alpaca")
processor = ProcessorFactory.create_with_strategy("sft", tokenizer, strategy)
```
### IOHandler
### 4. 使用策略模式（自定义 token）
```python
# 保存
IOHandler.save_h5("./out", "name", {"key": [tensor1, tensor2]})
# 加载
data = IOHandler.load_h5("./out") # → {"key": [tensor1, ...]}
# 自定义 ChatML 的特殊 token
strategy = StrategyFactory.create("chatml",
user_start="<s>user\n",
user_end="</s>\n",
assistant_start="<s>assistant\n",
assistant_end="</s>\n<eos>",
)
processor = ProcessorFactory.create_with_strategy("sft", tokenizer, strategy)
```
## 策略格式
| Strategy | Key | Default Tokens |
|-----------|------------|--------------------------------------------------------------------------------------------------|
| ChatML | `"chatml"` | `<\|im_start\|>user`, `<\|im_end\|>`, `<\|im_start\|>assistant`, `<eos>` |
| Alpaca | `"alpaca"` | `### Instruction:`, `### Response:`, `<eos>` |
所有策略的 token 均可通过构造函数参数自定义，同时支持通过 `StrategyFactory.register()` 注册新格式。
## 命令行工具
```bash
# 缓存 JSONL 到 H5
python scripts/cache_h5.py pt ./dataset/chinese-c4-pretrain
python scripts/cache_h5.py sft ./dataset/belle-sft --pack-size 4096 --strategy alpaca
```
## 脚本示例
```
scripts/
├── cache_h5.py # Stage 2: JSONL -> H5 (CLI tool)
├── pre_train/
│ ├── chinese-c4.py # Chinese pretrain data export
│ ├── chinese-cosmopedia.py # Chinese pretrain data export
│ ├── english-fineweb.py # English pretrain data export
│ └── english-wiki.py # English pretrain data export
├── supervised_finetuning/
│ ├── sft_belle.py # Belle Chinese SFT data export
│ ├── sft_chinese_instruct.py # Chinese instruct SFT (with TextNormalizer)
│ ├── sft_coder.py # Code SFT data export
│ ├── sft_firefly-1.1m-rephrased.py # Firefly SFT data export
│ └── sft_magpie-pro-300k.py # Magpie Pro SFT data export
└── reforce_learning/
└── dpp_chinese_dpo_pairs.py # Chinese DPO preference pairs export
```