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. 导出数据集

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 并缓存

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

packer = SequencePacker(pack_size=4096, pad_value=0)
packed = packer.pack([tensor1, tensor2, ...])  # → List[Tensor]

IOHandler

# 保存
IOHandler.save_h5("./out", "name", {"key": [tensor1, tensor2]})

# 加载
data = IOHandler.load_h5("./out")  # → {"key": [tensor1, ...]}