188 lines
8.2 KiB
Markdown
188 lines
8.2 KiB
Markdown
# DataPipeline
|
|
|
|
语言模型训练数据集处理工具包。支持 PT / SFT / DPO 范式。
|
|
|
|
## 项目结构
|
|
|
|
```
|
|
pipeline/
|
|
├── 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 文件解耦：
|
|
|
|
```
|
|
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 转换规则
|
|
|
|
**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 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. 分词并缓存为 HDF5
|
|
|
|
```python
|
|
from pipeline import BpeTokenizer, ProcessorFactory, cache_jsonl
|
|
|
|
tokenizer = BpeTokenizer("tokenizer.json")
|
|
processor = ProcessorFactory.create("pt", tokenizer)
|
|
|
|
cache_jsonl(
|
|
files=["./data/train.jsonl"],
|
|
output_dir="./cached",
|
|
processor=processor,
|
|
pack_size=4096,
|
|
)
|
|
```
|
|
|
|
### 3. 使用策略模式（默认 token）
|
|
|
|
```python
|
|
from pipeline import StrategyFactory, ProcessorFactory
|
|
|
|
strategy = StrategyFactory.create("alpaca")
|
|
processor = ProcessorFactory.create_with_strategy("sft", tokenizer, strategy)
|
|
```
|
|
|
|
### 4. 使用策略模式（自定义 token）
|
|
|
|
```python
|
|
# 自定义 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
|
|
``` |