CookLLM文档数据流水线
会员专享理解 Parquet shard 如何变成 input_ids、labels 和 attention_mask
这一节把前面几章串起来：预训练数据已经被切成 Parquet shard，tokenizer 已经能够把文本变成 token id，模型也已经定义好了。中间缺的就是数据流水线。
在 cookllm-bento 里，预训练 batch 的生成过程可以概括成：
Pretrain Data Pipeline
How a Parquet row becomes a training batch for next-token prediction.
1
Parquet shard
read text column from row groups
list[str]
2
Batch tokenization
prepend BOS, encode texts in parallel
token ids
3
Truncate to max_length
keep the first 512 tokens by default
sample
4
input_ids / labels
same tokens; shifting happens inside the model
tensors
5
Batch padding
pad input_ids, mask labels with -100
attention_mask
6
PretrainModule.training_step()
forward batch and log train/loss
loss
这条链路由两个文件负责：
| 层级 | 文件 | 作用 |
|---|---|---|
| DataModule | src/datamodule/pretrain_datamodule.py | 解析数据路径、加载 tokenizer、划分 train/val、创建 DataLoader |
| Dataset | src/dataset/pretrain.py | 流式读取 Parquet、批量 tokenization、生成单条训练样本 |
配置入口
预训练样例使用的数据配置核心字段如下：
登录以继续阅读
这是一篇付费内容，请登录您的账户以访问完整内容。