239 lines
5.9 KiB
Markdown
239 lines
5.9 KiB
Markdown
# DataPipeline
|
||
|
||
用于训练 KHAOSZ 模型的数据集处理工具。提供文本导出、Tokenize、序列打包、H5 存储等独立工具，支持预训练 / SFT / DPO 三种训练范式。
|
||
|
||
## 项目结构
|
||
|
||
```
|
||
pipeline/
|
||
├── tokenizer.py # BPE 分词器
|
||
├── text.py # 文本规范化
|
||
├── packing.py # 序列打包（bin-packing）
|
||
├── io.py # 文件/HDF5 读写
|
||
├── processors.py # PT / SFT / DPO 处理器
|
||
├── export.py # Dataset → JSONL 导出
|
||
└── cache.py # JSONL → Tokenize → H5 缓存
|
||
|
||
pre_train/ # 预训练数据处理脚本
|
||
supervised_finetuning/ # SFT 数据处理脚本
|
||
reforce_learning/ # DPO 数据处理脚本
|
||
```
|
||
|
||
## 设计理念
|
||
|
||
每个模块**独立可用、零互相依赖**，调用者按需组合：
|
||
|
||
```
|
||
HuggingFace Hub
|
||
│
|
||
▼ load_dataset()
|
||
DatasetDict
|
||
│
|
||
▼ 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），按需选择。
|
||
|
||
## 快速开始
|
||
|
||
### 安装依赖
|
||
|
||
```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_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):
|
||
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
|
||
└── ...
|
||
```
|