如何从Transformers Pipeline获取输入文本与NER结果配对？

当然可以实现，以下是几种高效的解决办法，支持批量处理：

方法1：手动批量配对输入与结果 #

直接提取数据集中的文本列表，传递给pipeline后用zip将输入文本和结果一一对应，适合快速实现：

from transformers import pipeline
from datasets import load_dataset

model = ...
tokenizer = ...
# 设置batch_size开启批量处理，提升效率
pipe = pipeline("ner", model=model, tokenizer=tokenizer, batch_size=8)
dataset = load_dataset("my_ner_dataset", split="test")

# 提取所有输入文本
texts = [item["text"] for item in dataset]
# 批量获取NER结果
ner_results = pipe(texts)

# 配对输出
for text, entities in zip(texts, ner_results):
    print(text, entities)

如果数据集较小，也可以逐个遍历数据项处理：

for item in dataset:
    text = item["text"]
    entities = pipe(text)
    print(text, entities)

方法2：用Dataset.map实现高效批量处理 #

这是Hugging Face官方推荐的大数据集处理方式，会自动处理批量逻辑，还能直接把结果存入数据集：

from transformers import pipeline
from datasets import load_dataset

model = ...
tokenizer = ...
pipe = pipeline("ner", model=model, tokenizer=tokenizer, batch_size=8)
dataset = load_dataset("my_ner_dataset", split="test")

# 定义批量处理函数
def add_ner_annotations(examples):
    # examples["text"]是批量文本列表，pipe返回对应批量结果
    examples["ner_entities"] = pipe(examples["text"])
    return examples

# 应用批量处理
dataset_with_ner = dataset.map(add_ner_annotations, batched=True)

# 遍历获取配对数据
for item in dataset_with_ner:
    print(item["text"], item["ner_entities"])

说明 #

你原来用的KeyDataset只负责提取文本字段输入给pipeline，丢失了原始数据的关联，所以无法直接拿到输入文本。上面的两种方法都保留了输入与结果的对应关系，且支持批量处理，可根据数据集大小选择合适的方案。

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