如何为spaCy NER模型创建基于Plotly的评估混淆矩阵

混淆矩阵生成的数据集选择 #

• 必做：测试集混淆矩阵。这是评估模型泛化能力的核心依据，能直接反映模型对未见过数据的实体识别偏差，是效果评估的核心产出。
• 选做：训练集混淆矩阵。仅当你需要排查过拟合问题、校验训练集标注质量时才需要生成，对比两个数据集的混淆矩阵差异，可以快速定位泛化性差的实体类别、或者过拟合问题。常规评估不需要每轮都生成训练集混淆矩阵，会拖慢训练速度。

具体实现路径 #

核心逻辑是基于spaCy的Example结构对齐真值标注和模型预测结果，逐token提取BIO格式的实体标签，统计混淆矩阵后用Plotly做可视化，所有结果可以直接同步到你现有的MLflow实验跟踪流程里。

前置依赖 #

提前导入用到的库即可：

from sklearn.metrics import confusion_matrix
import plotly.express as px
import numpy as np

代码修改（兼容现有训练/测试逻辑） #

你现有代码里的testing函数已经构造了Example列表、调用了nlp.evaluate()，只需要在这个函数里追加标签收集、混淆矩阵计算、可视化的逻辑即可，不需要改动训练循环的核心结构。
修改后的testing函数代码：

def testing(testing_data, nlp, iteration_index):
    testing_examples = []
    for text, annotations in testing_data:
        doc = nlp.make_doc(text)
        testing_examples.append(Example.from_dict(doc, annotations))

    # 原有指标计算逻辑
    scorer_example = nlp.evaluate(testing_examples)
    del scorer_example["ents_per_type"]
    mlflow.log_metrics(scorer_example, step=iteration_index)

    # 以下为新增的混淆矩阵计算逻辑
    # 1. 定义BIO格式标签列表，直接复用你已有的CLASSES实体列表即可
    label_list = CLASSES
    bio_labels = ["O"]
    for label in label_list:
        bio_labels.extend([f"B-{label}", f"I-{label}"])
    
    # 2. 逐样本逐token收集真值和预测标签
    y_true = []
    y_pred = []
    for example in testing_examples:
        for ref_token, pred_token in zip(example.reference, example.predicted):
            # 拼接真值标签
            true_tag = "O" if ref_token.ent_iob_ == "O" else f"{ref_token.ent_iob_}-{ref_token.ent_type_}"
            # 拼接预测标签
            pred_tag = "O" if pred_token.ent_iob_ == "O" else f"{pred_token.ent_iob_}-{pred_token.ent_type_}"
            y_true.append(true_tag)
            y_pred.append(pred_tag)
    
    # 3. 计算混淆矩阵，按行归一化展示每个真实类别的预测分布
    cm = confusion_matrix(y_true, y_pred, labels=bio_labels)
    cm_normalized = cm.astype("float") / cm.sum(axis=1)[:, np.newaxis]
    cm_normalized = np.nan_to_num(cm_normalized) # 处理测试集中未出现的类别导致的除0问题

    # 4. Plotly可视化并记录到MLflow
    fig = px.imshow(
        cm_normalized,
        x=bio_labels,
        y=bio_labels,
        color_continuous_scale="Blues",
        labels=dict(x="预测标签", y="真实标签", color="样本占比"),
        text_auto=".2f"
    )
    fig.update_layout(title=f"NER测试集混淆矩阵（迭代轮次{iteration_index}）")
    mlflow.log_figure(fig, f"test_confusion_matrix_iter_{iteration_index}.html")

如果需要生成训练集混淆矩阵，直接把训练数据传入上述逻辑即可，代码不需要做任何改动。

可选调整 #

• 如果你不需要BIO细粒度的混淆矩阵，想直接按实体类别统计（忽略B/I边界前缀），只需要把标签拼接逻辑改成直接取ent_type_即可，标签列表替换为["O"] + CLASSES。
• 不需要每轮迭代都生成混淆矩阵，可以加个判断，比如每5轮、或者最后一轮训练的时候再触发计算，减少训练耗时。
• 如果混淆矩阵中某一行全为0，说明对应实体类别在当前测试集中没有标注样本，需要检查数据集拆分的类别分布是否均衡。

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