如何可视化PyTorch构建的BERT类NER模型预测的实体？

核心逻辑 #

你的NER输出是逐token对应的BIO标签列表，要转成spaCy的Doc对象用于Displacy可视化，关键要做三件事：

• 确保分词后的token列表和标签列表长度严格匹配
• 将BIO标签解析为spaCy实体的(起始token索引, 结束token索引, 实体类型)三元组
• 用spaCy空语言模型构造Doc并添加实体

代码实现 #

1. 导入依赖 #

import spacy
from spacy.tokens import Doc, Span
from spacy import displacy

2. 标签转Doc工具函数 #

def convert_ner_to_spacy_doc(tokens, tags):
    # 创建空的英文语言模型对象（无需加载预训练权重）
    nlp = spacy.blank("en")
    doc = Doc(nlp.vocab, words=tokens)
    
    # 解析BIO标签，提取实体区间
    entities = []
    current_start = None
    current_label = None
    
    for idx, tag in enumerate(tags):
        if tag.startswith("B-"):
            # 结束上一个未闭合的实体
            if current_start is not None:
                entities.append((current_start, idx, current_label))
            # 初始化新实体
            current_start = idx
            current_label = tag.split("-")[1]
        elif tag.startswith("I-"):
            # 跳过无前置B-的无效I-标签
            if current_start is None:
                continue
            # 标签不一致时，结束上一个实体并启动新实体
            if tag.split("-")[1] != current_label:
                entities.append((current_start, idx, current_label))
                current_start = idx
                current_label = tag.split("-")[1]
        else:
            # 遇到O标签，闭合当前实体
            if current_start is not None:
                entities.append((current_start, idx, current_label))
                current_start = None
                current_label = None
    # 处理最后一个未闭合的实体
    if current_start is not None:
        entities.append((current_start, len(tokens), current_label))
    
    # 将实体添加到Doc对象
    spans = [Span(doc, start, end, label=label) for start, end, label in entities]
    doc.ents = spans
    return doc

3. 结合你的示例使用 #

# 你的输入文本和预测标签
text = 'functional scores for subjects, as assessed by Knee Injury and Osteoarthritis Outcome Score PRO (KOOS)'
tags = ['O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'B-PRO', 'I-PRO', 'I-PRO', 'I-PRO', 'I-PRO', 'I-PRO', 'I-PRO', 'O', 'I-PRO', 'O']
# 分词得到与标签长度匹配的token列表（需保证和BERT分词结果一致）
tokens = text.split()
assert len(tokens) == len(tags), "Token数量必须与标签数量完全匹配"

# 转换为spaCy Doc对象
doc = convert_ner_to_spacy_doc(tokens, tags)

# 可视化（Jupyter环境直接渲染）
displacy.render(doc, style="ent", jupyter=True)

# 本地脚本可导出为HTML文件
# displacy.render(doc, style="ent", page=True)

关键注意事项 #

• 分词一致性：必须确保传入的tokens列表和BERT模型的分词结果完全一致，否则标签对应关系会错乱。如果BERT用了WordPiece子词分词，需要先将子词合并为原始token再匹配标签。
• 无效标签处理：函数会自动跳过孤立的I-标签（无前置B-的情况），如果你的数据集存在这类无效标签，可根据需求调整逻辑。

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