如何在DataFrame指定列应用NER并解决spaCy AttributeError报错

现有如下DataFrame：

import pandas as pd
df = pd.DataFrame({'source': ['Paul', 'Paul'],
                   'target': ['GOOGLE', 'Ferrari'],
                   'edge': ['works at', 'drive']
                   })

原始数据输出：

source  target  edge
0   Paul    GOOGLE  works at
1   Paul    Ferrari drive

需求是对source和target列应用命名实体识别（NER），将原列值替换为对应的实体类型，预期结果：

source  target        edge
0   PERSON  ORGANIZATION  works at
1   PERSON  CAR           drive

尝试了以下代码，但调用ner(df)时抛出AttributeError: 'spacy.tokens.doc.Doc' object has no attribute 'label_'错误：

!python -m spacy download en_core_web_sm

import spacy
nlp = spacy.load('en_core_web_sm')

def ner(df):
    df['source_entities'] = df['source'].apply(lambda x: nlp(x).label_)
    df['target_entities'] = df['target'].apply(lambda x: nlp(x).label_)
    return df

nlp(x)返回的是Doc对象，该对象本身没有label_属性——label_是Doc中实体Span对象的专属属性，必须先从Doc中提取实体，再获取对应的标签。

基础适配版本（基于默认模型的实体识别） #

import spacy
import pandas as pd

nlp = spacy.load('en_core_web_sm')

def get_entity_label(text):
    doc = nlp(text)
    # 提取文本中第一个实体的标签，若无识别到实体则返回'UNKNOWN'
    return doc.ents[0].label_ if doc.ents else 'UNKNOWN'

def ner(df):
    # 直接替换原source列的值为实体类型
    df['source'] = df['source'].apply(get_entity_label)
    # 直接替换原target列的值为实体类型
    df['target'] = df['target'].apply(get_entity_label)
    return df

# 测试执行
df = pd.DataFrame({'source': ['Paul', 'Paul'],
                   'target': ['GOOGLE', 'Ferrari'],
                   'edge': ['works at', 'drive']
                   })

result_df = ner(df)
print(result_df)

补充：适配Ferrari的CAR实体标签 #

默认的en_core_web_sm模型不会将Ferrari识别为CAR（通常会识别为ORG或无实体），如果需要得到预期的CAR标签，可以通过自定义规则实现：

def get_entity_label(text):
    # 自定义特定词汇的实体映射
    custom_entity_map = {'Ferrari': 'CAR'}
    if text in custom_entity_map:
        return custom_entity_map[text]
    doc = nlp(text)
    return doc.ents[0].label_ if doc.ents else 'UNKNOWN'

source  target        edge
0   PERSON  ORGANIZATION  works at
1   PERSON  CAR           drive

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