数据集处理

模型的效果，一半取决于数据。这一页讲如何用 Datasets 库把原始语料变成模型能吃的张量批次，尤其是语言模型（LM）数据集特有的「切块」处理。

1. 加载数据的几种来源

from datasets import load_dataset, Dataset

ds = load_dataset("imdb")                 # 官方数据集(电影评论情感)
ds = load_dataset("csv", data_files="a.csv")   # 自己的 CSV/JSON/JSONL/文本
ds = Dataset.from_dict({"text": ["句子1", "句子2"]})  # 内存里直接构造

print(ds["train"][0])   # {'text': '...', 'label': 1}

• 一个 DatasetDict 通常含 train / validation / test 拆分。
• 数据量大时用流式 streaming=True，边下边处理，不占满内存。

2. 预处理：map + 分词

map 对每条样本做变换（常是分词），并设 batched=True 批量加速：

from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained("bert-base-uncased")

def encode(batch):
    return tok(batch["text"], truncation=True, max_length=256)

ds = ds.map(encode, batched=True, num_proc=4)   # 并行分词
ds.set_format("torch")                           # 输出转成 PyTorch 张量

对分类任务，分词后 (input_ids, attention_mask, label) 就够训练了。

3. 语言模型数据的特殊之处：切块（Chunking）

预训练/续写一个 LM 时，没有 label——输入本身就是标签（右移一位）。我们要做的是把长短不一的文本首尾相连拼起来，再切成固定长度的块，避免浪费与碎片：

多条文本拼接:  "...文章A[EOS]文章B[EOS]文章C..."
                 │  平铺成一条长 token 流, 按 block_size 切
切成固定块:    [块1: 512 tokens][块2: 512 tokens][块3: 512 tokens] ...

训练时: input  = 块[:-1]      # 前 511 个
        target = 块[1:]       # 后 511 个(右移一位) → 预测下一词

block_size = 512
tokenized = ds["train"].map(lambda x: tok(x["text"]), batched=True)
concat = tokenized["input_ids"]            # 所有 token 连成一条(示意)
import numpy as np
import torch
from torch.utils.data import Dataset

class LMDataset(Dataset):
    def __init__(self, ids, block_size):
        self.ids = ids
        self.bs = block_size
    def __len__(self):
        return (len(self.ids) - 1) // self.bs
    def __getitem__(self, i):
        chunk = self.ids[i*self.bs : (i+1)*self.bs + 1]
        x = torch.tensor(chunk[:-1])
        y = torch.tensor(chunk[1:])         # 右移一位作为目标
        return {"input_ids": x, "labels": y}

这正是 搭建NanoGPT 里数据加载的核心思路，只是我们用字符级词表。

4. 分类数据的 DataLoader 打包

不定长样本要 padding 到 batch 内统一长度，用 DataCollator：

from transformers import DataCollatorWithPadding
from torch.utils.data import DataLoader

collator = DataCollatorWithPadding(tokenizer=tok)   # 动态补齐到本批最长
loader = DataLoader(ds["train"], batch_size=16, shuffle=True, collate_fn=collator)

• 动态 padding（按当前 batch 最长补）比无脑补到 512 省显存、更快。
• LM 固定块数据因为长度已对齐，collator 更简单。

5. 数据集常做操作速查

ds = ds.train_test_split(test_size=0.1, seed=42)     # 切分
ds = ds.shuffle(seed=42)                             # 打乱
ds = ds.select(range(1000))                          # 取子集(快速试跑)
small = ds["train"].train_test_split(test_size=0.01)["train"]  # 抽样调试
ds.save_to_disk("./my_ds")                           # 缓存到磁盘, 下次秒开

6. 数据质量的经验法则

• 去重:   重复样本会让模型"死记"而非"学会"
• 清洗:   去掉乱码/广告/截断的半句话/HTML残留
• 配比:   领域微调里, 领域数据 + 少量通用数据 防止遗忘
• 长度:   block_size/最大长度要与目标分布匹配, 别把大量 padding 当正文学
• 划分:   训练/验证严格不泄漏(同一文档别同时出现在两边)

小结

• 分类数据：map 分词 + DataCollatorWithPadding 动态补齐即可训练。
• 语言模型数据：拼接成 token 流 → 切成固定块 → 输入/目标错位一位，是自监督 LM 的关键预处理。
• 善用 set_format("torch")、save_to_disk、抽样 select 提升迭代速度。
• 数据质量与去重、清洗、不泄漏，往往比调参更决定上限。

下一步

数据备好了，用 Trainer 把训练循环一站式跑起来 → 训练流程Trainer。