112 lines
2.9 KiB
Markdown
112 lines
2.9 KiB
Markdown
# DataPipeline
|
|
|
|
数据集处理工具，支持预训练 / 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
|
|
```
|
|
|
|
## 设计理念
|
|
|
|
模块**独立可用**，通过磁盘文件解耦，按需组合：
|
|
|
|
```
|
|
Dataset → export_dataset() → JSONL → cache_jsonl() → HDF5
|
|
↑
|
|
processors.py
|
|
packing.py
|
|
io.py
|
|
```
|
|
|
|
## 使用方法
|
|
|
|
### 1. 导出数据集
|
|
|
|
```python
|
|
from datasets import load_dataset
|
|
from pipeline.export 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"]}, # 可选
|
|
)
|
|
```
|
|
|
|
### 2. Tokenize 并缓存
|
|
|
|
```python
|
|
from pipeline import BpeTokenizer, ProcessorFactory, cache_jsonl
|
|
|
|
tokenizer = BpeTokenizer("tokenizer.json")
|
|
processor = ProcessorFactory.create("pt", tokenizer) # "pt" | "sft" | "dpo"
|
|
|
|
cache_jsonl(
|
|
files=["./data/train.jsonl"],
|
|
output_dir="./cached",
|
|
processor=processor,
|
|
pack_size=4096, # <=0 不打包
|
|
pad_value=1,
|
|
)
|
|
```
|
|
|
|
### 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
|
|
|
|
```python
|
|
packer = SequencePacker(pack_size=4096, pad_value=0)
|
|
packed = packer.pack([tensor1, tensor2, ...]) # → List[Tensor]
|
|
```
|
|
|
|
### IOHandler
|
|
|
|
```python
|
|
# 保存
|
|
IOHandler.save_h5("./out", "name", {"key": [tensor1, tensor2]})
|
|
|
|
# 加载
|
|
data = IOHandler.load_h5("./out") # → {"key": [tensor1, ...]}
|
|
```
|