使用Huggingface Tokenizer时如何保留数据集原始列？

当使用设置了return_overflowing_tokens=True的Huggingface Tokenizer时，单个输入字符串会生成多个token序列，导致原始数据集列与分词后的token序列不再是1:1对应。如果需要在分词后的数据集里保留原始列（如示例中的txt、src），可以利用tokenizer返回的overflow_to_sample_mapping字段实现，具体步骤如下：

核心思路 #

overflow_to_sample_mapping是tokenizer返回的数组，其中每个元素对应当前生成的token序列所属的原始样本索引。通过这个映射关系，我们可以把原始样本的对应列复制到每个分词后的样本中。

具体实现代码 #

1. 构造示例数据集 #

from datasets import Dataset

data = {
    "txt": ["The quick brown fox", "jumped over the lazy hens"],
    "src": ["Nursery rhyme 1", "Nursery rhyme 2"]
}
ds = Dataset.from_dict(data)

2. 定义带列复制的分词函数 #

from transformers import AutoTokenizer

# 加载示例tokenizer，可替换为你需要的模型
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

def tokenize_batch(examples):
    # 执行分词，开启overflow返回，并设置截断/填充参数
    tokenized_result = tokenizer(
        examples["txt"],
        return_overflowing_tokens=True,
        truncation=True,
        padding="max_length",
        max_length=5  # 故意设置短长度触发overflow，方便演示
    )
    
    # 获取token序列到原始样本的映射关系
    sample_mapping = tokenized_result.pop("overflow_to_sample_mapping")
    
    # 遍历需要保留的原始列，复制到分词结果中
    for col_name in ["txt", "src"]:
        tokenized_result[col_name] = [examples[col_name][idx] for idx in sample_mapping]
    
    return tokenized_result

3. 执行数据集映射 #

# 批量处理数据集，保留所有需要的列
tokenized_ds = ds.map(tokenize_batch, batched=True)

验证结果 #

执行后查看分词后的数据集：

# 第一个分词样本（对应原始第一个文本的第一部分）
print(tokenized_ds[0])
# 输出示例：
# {'input_ids': [101, 1996, 4248, 2829, 102], 'token_type_ids': [0,0,0,0,0], 'attention_mask': [1,1,1,1,1], 'txt': 'The quick brown fox', 'src': 'Nursery rhyme 1'}

# 第二个分词样本（对应原始第一个文本的第二部分）
print(tokenized_ds[1])
# 输出示例：
# {'input_ids': [101, 4419, 0, 0, 0], 'token_type_ids': [0,0,0,0,0], 'attention_mask': [1,1,0,0,0], 'txt': 'The quick brown fox', 'src': 'Nursery rhyme 1'}

注意事项 #

• 无需手动设置remove_columns参数，分词函数会自动将原始列复制到结果中，避免丢失
• 字符串类型的原始列会被直接保留，不会被转换为Tensor，符合需求
• 批量处理时，sample_mapping会自动对应当前batch内的样本索引，map函数会正确处理全局数据集的映射关系

内容的提问来源于stack exchange，提问作者SRobertJames