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
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
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）
from pipeline import StrategyFactory, ProcessorFactory
strategy = StrategyFactory.create("alpaca")
processor = ProcessorFactory.create_with_strategy("sft", tokenizer, strategy)
4. 使用策略模式（自定义 token）
# 自定义 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() 注册新格式。
命令行工具
# 缓存 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