如何改造HuggingFace Datasets预处理函数适配自定义Pandas数据集

原有为HuggingFace Datasets设计的文本摘要预处理函数如下：

from datasets import load_dataset, load_metric
from transformers import AutoTokenizer

raw_datasets = load_dataset("xsum")
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)

max_input_length = 1024
max_target_length = 128

if model_checkpoint in ["t5-small", "t5-base", "t5-larg", "t5-3b", "t5-11b"]:
    prefix = "summarize: "
else:
    prefix = ""

def preprocess_function(examples):
    inputs = [prefix + doc for doc in examples["document"]]
    model_inputs = tokenizer(inputs, max_length=max_input_length, truncation=True)

    # Setup the tokenizer for targets
    with tokenizer.as_target_tokenizer():
        labels = tokenizer(examples["summary"], max_length=max_target_length, truncation=True)

    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

tokenized_datasets = raw_datasets.map(preprocess_function, batched=True)

当前使用自定义Pandas DataFrame格式的train数据集，无法调用HF Dataset的.map()方法，将最后一行替换为tokenized_datasets = train.apply(preprocess_function)后，触发如下KeyError报错：

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-18-ad0e3caaca6d> in <module>()
----> 1 tokenized_datasets = train.apply(preprocess_function)

7 frames
/usr/local/lib/python3.7/dist-packages/pandas/core/indexes/range.py in get_loc(self, key, method, tolerance)
    386                 except ValueError as err:
    387                     raise KeyError(key) from err
--> 388             raise KeyError(key)
    389         return super().get_loc(key, method=method, tolerance=tolerance)
    390 

KeyError: 'input'

需要找到正确的改造方式，将原始训练集处理为符合模型输入要求的分词后数据集。

• 原preprocess_function适配HuggingFace Datasets的批量处理逻辑：当.map()设置batched=True时，传入函数的examples是按列组织的字典，键为列名，值为该列所有样本组成的列表，因此函数中直接取examples["document"]、examples["summary"]可以正常拿到全量文本。
• Pandas的apply默认传参逻辑和上述格式完全不匹配：默认axis=0时按列遍历传入整列Series，axis=1时按行遍历传入单行Series，都不符合原函数预期的输入格式，因此触发键不存在的报错。

方案1：转换为HuggingFace Dataset格式（推荐，改动最小） #

不需要修改原有预处理函数，直接将Pandas DataFrame转为HF Dataset对象，即可复用原有.map()逻辑，且批量处理效率更高：

from datasets import Dataset

# Pandas DataFrame 转 HuggingFace Dataset
hf_train = Dataset.from_pandas(train)
# 直接复用原预处理函数，batched=True开启批量处理
tokenized_train = hf_train.map(
    preprocess_function, 
    batched=True,
    remove_columns=train.columns.tolist() # 处理完移除原文本列，减少内存占用
)

方案2：纯Pandas批量处理（不依赖datasets库） #

如果不想引入HuggingFace Datasets依赖，不要用逐行apply（效率极低），直接对整列做批量分词处理：

# 批量处理输入文本
input_texts = [prefix + doc for doc in train["document"]]
model_inputs = tokenizer(
    input_texts,
    max_length=max_input_length,
    truncation=True,
    padding="max_length" # 根据训练需求选择是否补全到最大长度
)

# 批量处理目标摘要
with tokenizer.as_target_tokenizer():
    label_encodings = tokenizer(
        train["summary"].tolist(),
        max_length=max_target_length,
        truncation=True,
        padding="max_length"
    )

# 组装为模型可接收的字典格式
tokenized_train = {
    "input_ids": model_inputs["input_ids"],
    "attention_mask": model_inputs["attention_mask"],
    "labels": label_encodings["input_ids"]
}

注意：无论用哪种方案，都要优先选择批量分词，逐行调用tokenizer的处理速度会比批量处理慢1~2个数量级，数据量较大时会严重拖慢预处理效率。

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