Hugging Face Datasets set_format 怎样只把训练字段转换为张量

Hugging Face Datasets 里同时保留原始文本、标签和模型输入时，不必把整行数据手工转换成张量。把 input_ids、attention_mask、labels 明确列出来，再用 with_format("torch") 创建训练视图即可；text 仍可保持普通 Python 对象。官方文档地址：https://huggingface.co/docs/datasets/en/process

核心做法是“先限定训练列，再设置返回格式”。需要原地改变 Dataset 用 set_format，需要保留原始对象并派生训练视图用 with_format。

先把字段边界分清，再选择格式化方式

训练数据常见四类字段：原始 text 用于回溯，input_ids 和 attention_mask 进入模型，labels 参与损失计算。若直接对所有列做格式化，字符串列并不会自动变成有意义的训练张量，反而容易让批处理接口变得含糊。

更稳妥的边界是只选择模型真正需要的列。下面的示意代码没有执行外部数据集加载，重点是字段和返回类型的组织方式：

from datasets import Dataset

ds = Dataset.from_dict({
    "text": ["第一条文本", "第二条文本"],
    "input_ids": [[101, 10, 102], [101, 20, 102]],
    "attention_mask": [[1, 1, 1], [1, 1, 1]],
    "labels": [0, 1],
})

# 只把模型输入和标签交给张量格式化，原始 text 不参与转换。
train_view = ds.with_format(
    type="torch",
    columns=["input_ids", "attention_mask", "labels"],
    output_all_columns=False,
)

# 取样时只关注训练字段，避免把原始列误传给模型。
batch = train_view[:2]
print(batch.keys())

set_format 和 with_format 的差别在对象生命周期

set_format 会原地修改当前 Dataset 的取数格式，适合已经确定后续都按同一方式训练的对象；with_format 返回一个新的 Dataset，更适合同时保留原始视图、评估视图和训练视图。两者都支持 type、columns 与 output_all_columns。

这里的“按需”指访问数据时执行格式化，而不是提前把所有列重写一遍。列选择越窄，批次输出越容易和模型的 forward 参数对齐。

接入 DataLoader 时检查返回对象和边界

接入 PyTorch 时，先用一个小批次确认键名和类型，再交给完整训练循环。变长的 input_ids 不能直接堆叠成规则张量，应在 tokenizer 的批处理或 collator 中完成 padding；固定形状的字段才适合直接组成批次。

from torch.utils.data import DataLoader

# Dataset 已限定训练列，DataLoader 只负责按批次取出它们。
loader = DataLoader(train_view, batch_size=2, shuffle=False)
first_batch = next(iter(loader))

# 先检查接口契约，再把 batch 传入模型，避免列名错位。
required = {"input_ids", "attention_mask", "labels"}
missing = required.difference(first_batch)
if missing:
    raise KeyError(f"缺少训练字段: {sorted(missing)}")

# 训练结束后如需查看原始 Python 对象，可恢复返回格式。
train_view.reset_format()

如果希望未格式化的 text 也随批次保留，把 output_all_columns 改为 True，但模型调用前要显式取出训练字段。不要把“字段存在”误认为“字段已经是模型需要的 dtype”。

常见问题

set_format 会复制整个数据集吗？

它主要改变取数时的返回格式，不能把它理解成把所有行提前转换并保存成另一份完整数据。真正的批次内存仍取决于访问范围、字段形状和 DataLoader 配置。

为什么 labels 没有变成想要的 dtype？

格式化只负责把字段转换到目标框架，具体 dtype 还受原始特征和转换参数影响。应在第一批数据上打印键名、形状和 dtype，再决定是否在模型入口显式处理。

什么时候应该用 with_format 而不是 set_format？

需要同时保留原始 Dataset、评估视图和训练视图时优先用 with_format；只剩一个固定训练出口时再考虑原地使用 set_format。

把格式化范围限定在训练字段，通常比“所有列一起转型”更容易排查。最后保留一个小批次检查点：列名、张量类型、形状和 padding 结果都符合模型接口，再扩大到完整训练集。