perf: 优化 processors、cache、packing 模块性能并简化 README
This commit is contained in:
@@ -1,6 +1,6 @@
|
||||
# DataPipeline
|
||||
|
||||
用于训练 KHAOSZ 模型的数据集处理工具。提供文本导出、Tokenize、序列打包、H5 存储等独立工具，支持预训练 / SFT / DPO 三种训练范式。
|
||||
数据集处理工具，支持预训练 / SFT / DPO 三种训练范式。
|
||||
|
||||
## 项目结构
|
||||
|
||||
@@ -8,231 +8,104 @@
|
||||
pipeline/
|
||||
├── tokenizer.py # BPE 分词器
|
||||
├── text.py # 文本规范化
|
||||
├── packing.py # 序列打包（bin-packing）
|
||||
├── packing.py # 序列打包
|
||||
├── 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 数据处理脚本
|
||||
├── export.py # Dataset → JSONL
|
||||
└── cache.py # JSONL → Tokenize → H5
|
||||
```
|
||||
|
||||
## 设计理念
|
||||
|
||||
每个模块**独立可用、零互相依赖**，调用者按需组合：
|
||||
模块**独立可用**，通过磁盘文件解耦，按需组合：
|
||||
|
||||
```
|
||||
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 张量文件
|
||||
Dataset → export_dataset() → JSONL → cache_jsonl() → HDF5
|
||||
↑
|
||||
processors.py
|
||||
packing.py
|
||||
io.py
|
||||
```
|
||||
|
||||
> 各阶段之间通过磁盘文件解耦。你可以只执行阶段 1（导出 JSONL），也可以继续执行阶段 2（tokenize 并缓存为 H5），按需选择。
|
||||
## 使用方法
|
||||
|
||||
## 快速开始
|
||||
|
||||
### 安装依赖
|
||||
|
||||
```bash
|
||||
pip install -e .
|
||||
```
|
||||
|
||||
### 阶段 1：导出数据集为 JSONL
|
||||
### 1. 导出数据集
|
||||
|
||||
```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,
|
||||
output_dir="./data",
|
||||
output_prefix="train",
|
||||
process_func=lambda x: {"text": x["content"]}, # 可选
|
||||
)
|
||||
```
|
||||
|
||||
> `process_func` 返回单个 `dict` 或 `list[dict]`（一条样本可展开为多条）。
|
||||
|
||||
**使用文本规范化：**
|
||||
### 2. Tokenize 并缓存
|
||||
|
||||
```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
|
||||
from pipeline import BpeTokenizer, ProcessorFactory, cache_jsonl
|
||||
|
||||
tokenizer = BpeTokenizer("tokenizer.json")
|
||||
processor = ProcessorFactory.create("pt", tokenizer)
|
||||
processor = ProcessorFactory.create("pt", tokenizer) # "pt" | "sft" | "dpo"
|
||||
|
||||
cache_jsonl(
|
||||
files=["./dataset/my-pretrain_chunk_0.jsonl"],
|
||||
files=["./data/train.jsonl"],
|
||||
output_dir="./cached",
|
||||
processor=processor,
|
||||
pack_size=4096, # 可选，打包长度；<=0 不打包
|
||||
pack_size=4096, # <=0 不打包
|
||||
pad_value=1,
|
||||
)
|
||||
```
|
||||
|
||||
**处理器类型：**
|
||||
### 3. 处理器类型
|
||||
|
||||
| 类型 | 工厂 key | 输入格式 | 输出 keys |
|
||||
|------|----------|---------|-----------|
|
||||
| 类型 | key | 输入 | 输出 |
|
||||
|------|-----|------|------|
|
||||
| 预训练 | `"pt"` | `{"text": "..."}` | `["sequence"]` |
|
||||
| SFT | `"sft"` | `{"query": "...", "response": "..."}` | `["sequence", "loss_mask"]` |
|
||||
| DPO | `"dpo"` | 待定 | `["chosen", "chosen_mask", "rejected", "rejected_mask"]` |
|
||||
| DPO | `"dpo"` | `{"query": "...", "chosen": "...", "rejected": "..."}` | `["chosen", "chosen_mask", "rejected", "rejected_mask"]` |
|
||||
|
||||
## 独立工具参考
|
||||
## 参数参考
|
||||
|
||||
### BpeTokenizer
|
||||
### export_dataset()
|
||||
|
||||
```python
|
||||
from pipeline.tokenizer import BpeTokenizer
|
||||
| 参数 | 类型 | 默认 | 说明 |
|
||||
|------|------|------|------|
|
||||
| `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" | 默认文本列名 |
|
||||
|
||||
tokenizer = BpeTokenizer("tokenizer.json")
|
||||
ids = tokenizer.encode("hello world") # → [1234, 5678, 1]
|
||||
text = tokenizer.decode(ids) # → "hello world"
|
||||
len(tokenizer) # → 词表大小
|
||||
```
|
||||
### cache_jsonl()
|
||||
|
||||
### TextNormalizer
|
||||
|
||||
```python
|
||||
from pipeline.text import TextNormalizer
|
||||
|
||||
normalizer = TextNormalizer()
|
||||
text = normalizer.normalize(raw_text)
|
||||
```
|
||||
|
||||
替换规则包括：全角引号 → 半角、各种短横线统一、不间断空格 → 普通空格等。支持自定义规则：
|
||||
|
||||
```python
|
||||
normalizer = TextNormalizer(custom_rules={"旧词": "新词"})
|
||||
```
|
||||
| 参数 | 类型 | 默认 | 说明 |
|
||||
|------|------|------|------|
|
||||
| `files` | List[str] | 必填 | JSONL 文件列表 |
|
||||
| `output_dir` | str | 必填 | 输出目录 |
|
||||
| `processor` | BaseProcessor | 必填 | 处理器实例 |
|
||||
| `pack_size` | int | -1 | 打包长度，<=0 不打包 |
|
||||
| `pad_value` | int | 1 | 填充值 |
|
||||
|
||||
### 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
|
||||
packed = packer.pack([tensor1, tensor2, ...]) # → List[Tensor]
|
||||
```
|
||||
|
||||
### IOHandler
|
||||
|
||||
```python
|
||||
from pipeline.io import IOHandler
|
||||
|
||||
# 保存
|
||||
IOHandler.save_h5("./output", "my_data", {"sequence": [tensor1, tensor2]})
|
||||
IOHandler.save_h5("./out", "name", {"key": [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
|
||||
└── ...
|
||||
data = IOHandler.load_h5("./out") # → {"key": [tensor1, ...]}
|
||||
```
|
||||
|
||||
+19
-12
@@ -2,7 +2,7 @@
|
||||
import json
|
||||
import os
|
||||
import logging
|
||||
from typing import List
|
||||
from typing import List, Dict
|
||||
from pathlib import Path
|
||||
|
||||
from tqdm import tqdm
|
||||
@@ -39,33 +39,40 @@ def cache_jsonl(
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
output_files: List[str] = []
|
||||
# Cache output_keys to avoid repeated attribute access
|
||||
output_keys = processor.output_keys
|
||||
|
||||
for file_path in files:
|
||||
file_name = Path(file_path).stem
|
||||
|
||||
arrows = []
|
||||
# Pre-allocate lists for each output key
|
||||
arrows: Dict[str, List] = {key: [] for key in output_keys}
|
||||
|
||||
# Read and process all lines
|
||||
with open(file_path, "r", encoding="utf-8") as f:
|
||||
for line_num, line in enumerate(tqdm(f, desc=f"Processing {file_name}", leave=False), start=1):
|
||||
try:
|
||||
arrow = processor.process(json.loads(line))
|
||||
result = processor.process(json.loads(line))
|
||||
if result is not None:
|
||||
# Batch append: add each key's tensor to corresponding list
|
||||
for key in output_keys:
|
||||
arrows[key].append(result[key])
|
||||
except json.JSONDecodeError as e:
|
||||
logger.warning(f"JSON decode error in {file_path} line {line_num}: {e}. Skipping line.")
|
||||
continue
|
||||
except Exception as e:
|
||||
logger.warning(f"Unexpected error processing line {line_num} in {file_path}: {e}. Skipping line.")
|
||||
continue
|
||||
if arrow is not None:
|
||||
arrows.append(arrow)
|
||||
|
||||
package = {key: [a[key] for a in arrows] for key in processor.output_keys}
|
||||
|
||||
output = {}
|
||||
for key in processor.output_keys:
|
||||
# Convert lists to tensors once per key
|
||||
if pack_size > 0:
|
||||
packer = SequencePacker(pack_size, pad_value) # independent instance per key
|
||||
output[key] = packer.pack(package[key])
|
||||
output = {}
|
||||
for key in output_keys:
|
||||
packer = SequencePacker(pack_size, pad_value)
|
||||
output[key] = packer.pack(arrows[key])
|
||||
else:
|
||||
output[key] = package[key]
|
||||
# No packing: directly use the arrow tensors
|
||||
output = arrows
|
||||
|
||||
IOHandler.save_h5(output_dir, file_name, output)
|
||||
h5_path = os.path.join(output_dir, f"{file_name}.h5")
|
||||
|
||||
@@ -34,6 +34,7 @@ class IOHandler:
|
||||
def save_h5(output_dir: str, file_name: str, tensor_group: Dict[str, List[Tensor]]) -> None:
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
full_path = os.path.join(output_dir, f"{file_name}.h5")
|
||||
|
||||
with h5py.File(full_path, 'w') as f:
|
||||
for key, tensors in tensor_group.items():
|
||||
grp = f.create_group(key)
|
||||
|
||||
+41
-22
@@ -1,5 +1,5 @@
|
||||
import logging
|
||||
from typing import List
|
||||
from typing import List, Optional
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
@@ -14,14 +14,23 @@ class SequencePacker:
|
||||
self.pack_size = pack_size
|
||||
self.pad_value = pad_value
|
||||
self.dtype = dtype
|
||||
# Pre-allocate buffer for better performance
|
||||
self._buffer: Optional[Tensor] = None
|
||||
self._reset()
|
||||
|
||||
def _reset(self) -> None:
|
||||
"""Reset internal state for instance reuse."""
|
||||
self._current_pack = torch.full(
|
||||
# Reuse buffer instead of creating new tensors
|
||||
if self._buffer is None or self._buffer.shape[0] != self.pack_size:
|
||||
self._buffer = torch.full(
|
||||
(self.pack_size,), self.pad_value, dtype=self.dtype
|
||||
)
|
||||
else:
|
||||
self._buffer.fill_(self.pad_value)
|
||||
self._current_pos = 0
|
||||
self._packages: List[Tensor] = []
|
||||
# Backward compatibility: maintain _current_pack reference
|
||||
self._current_pack = self._buffer
|
||||
|
||||
@error_handler()
|
||||
def pack(self, sequences: List[Tensor]) -> List[Tensor]:
|
||||
@@ -37,53 +46,63 @@ class SequencePacker:
|
||||
# Input validation
|
||||
if not sequences:
|
||||
return []
|
||||
|
||||
# Validate and cache tensor sizes in one pass
|
||||
tensor_sizes = []
|
||||
for i, seq in enumerate(sequences):
|
||||
if seq.dim() != 1:
|
||||
raise ValueError(
|
||||
f"Expected 1D tensor at index {i}, got {seq.dim()}D tensor with shape {seq.shape}"
|
||||
)
|
||||
# Check dtype compatibility and warn if mismatched
|
||||
tensor_sizes.append(seq.numel())
|
||||
if seq.dtype != self.dtype:
|
||||
logger.warning(
|
||||
f"Input tensor dtype {seq.dtype} does not match packer dtype {self.dtype}, "
|
||||
f"will be converted. This may affect packing efficiency."
|
||||
)
|
||||
|
||||
packages = []
|
||||
# Sort by length in descending order to improve packing efficiency
|
||||
# Use sorted() to avoid modifying the input list
|
||||
sorted_sequences = sorted(sequences, key=lambda x: x.numel(), reverse=True)
|
||||
# Reset state for new packing
|
||||
self._packages = []
|
||||
self._reset()
|
||||
|
||||
for tensor in sorted_sequences:
|
||||
# Combine sequences with their sizes for sorting
|
||||
indexed_seqs = list(zip(sequences, tensor_sizes))
|
||||
# Sort by size descending (First-Fit Decreasing algorithm)
|
||||
indexed_seqs.sort(key=lambda x: x[1], reverse=True)
|
||||
|
||||
for tensor, tensor_size in indexed_seqs:
|
||||
# Truncate sequences that exceed pack_size
|
||||
if tensor.numel() > self.pack_size:
|
||||
if tensor_size > self.pack_size:
|
||||
logger.warning(
|
||||
f"Sequence length {tensor.numel()} exceeds pack_size {self.pack_size}, truncating"
|
||||
f"Sequence length {tensor_size} exceeds pack_size {self.pack_size}, truncating"
|
||||
)
|
||||
tensor_size = self.pack_size
|
||||
tensor = tensor[: self.pack_size]
|
||||
tensor_size = tensor.numel()
|
||||
|
||||
# Current package is full, create a new one
|
||||
if self._current_pos + tensor_size > self.pack_size:
|
||||
packages.append(self._current_pack)
|
||||
self._current_pack = torch.full(
|
||||
(self.pack_size,), self.pad_value, dtype=self.dtype
|
||||
)
|
||||
# Finish current package (pad to pack_size)
|
||||
package = self._buffer.clone()
|
||||
self._packages.append(package)
|
||||
# Reset buffer for reuse
|
||||
self._buffer.fill_(self.pad_value)
|
||||
self._current_pos = 0
|
||||
|
||||
# Place tensor in current package (remaining positions stay as pad_value)
|
||||
self._current_pack[self._current_pos : self._current_pos + tensor_size] = (
|
||||
tensor
|
||||
)
|
||||
# Place tensor in current package
|
||||
self._buffer[self._current_pos : self._current_pos + tensor_size] = tensor
|
||||
self._current_pos += tensor_size
|
||||
|
||||
# Handle the last package
|
||||
# Handle the last package (pad to pack_size)
|
||||
if self._current_pos > 0:
|
||||
packages.append(self._current_pack)
|
||||
package = self._buffer.clone()
|
||||
self._packages.append(package)
|
||||
|
||||
# Clear buffer and reset state for backward compatibility
|
||||
self._buffer = None
|
||||
self._current_pack = None
|
||||
self._current_pos = 0
|
||||
|
||||
return packages
|
||||
return self._packages
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset packer state for reuse. More efficient than creating a new instance."""
|
||||
|
||||
+16
-7
@@ -41,14 +41,18 @@ class SFTProcessor(BaseProcessor):
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
|
||||
query, response = input_dict["query"], input_dict["response"]
|
||||
query = input_dict["query"]
|
||||
response = input_dict["response"]
|
||||
|
||||
q = self.tokenizer.encode(
|
||||
f"<|im_start|>user\n{query}<|im_end|>\n<|im_start|>assistant\n"
|
||||
)
|
||||
a = self.tokenizer.encode(f"{response}<|im_end|>\n<eos>")
|
||||
|
||||
q_len = len(q)
|
||||
tokens = torch.tensor(q + a, dtype=torch.int32)
|
||||
loss_mask = torch.zeros_like(tokens, dtype=torch.bool)
|
||||
loss_mask[len(q):] = True
|
||||
loss_mask = torch.zeros(q_len + len(a), dtype=torch.bool)
|
||||
loss_mask[q_len:] = True
|
||||
return {"sequence": tokens, "loss_mask": loss_mask}
|
||||
|
||||
@property
|
||||
@@ -72,14 +76,19 @@ class DPOProcessor(BaseProcessor):
|
||||
)
|
||||
|
||||
chosen = self.tokenizer.encode(f"{chosen_response}<|im_end|>\n<eos>")
|
||||
q_len = len(q)
|
||||
chosen_len = len(chosen)
|
||||
|
||||
chosen_tokens = torch.tensor(q + chosen, dtype=torch.int32)
|
||||
chosen_mask = torch.zeros_like(chosen_tokens, dtype=torch.bool)
|
||||
chosen_mask[len(q):] = True
|
||||
chosen_mask = torch.zeros(q_len + chosen_len, dtype=torch.bool)
|
||||
chosen_mask[q_len:] = True
|
||||
|
||||
rejected = self.tokenizer.encode(f"{rejected_response}<|im_end|>\n<eos>")
|
||||
rejected_len = len(rejected)
|
||||
|
||||
rejected_tokens = torch.tensor(q + rejected, dtype=torch.int32)
|
||||
rejected_mask = torch.zeros_like(rejected_tokens, dtype=torch.bool)
|
||||
rejected_mask[len(q):] = True
|
||||
rejected_mask = torch.zeros(q_len + rejected_len, dtype=torch.bool)
|
||||
rejected_mask[q_len:] = True
|
||||
|
||||
return {
|
||||
"chosen": chosen_tokens,
|
||||
|
||||
Reference in New Issue
Block a user