如何获取每个类别的SHAP值并生成对应条形图？

你的代码报错核心原因是SHAP返回的shap_values对象结构需要调整索引方式，且要适配每个类别的维度处理。以下是修正后的实现方案：

修正要点 #

• 从模型配置中直接获取类别标签，避免手动输入出错
• 对每个类别，计算所有样本中各token的平均SHAP值，再调用条形图绘制函数
• 动态token长度（即shap_values中的None维度）会被SHAP内置方法自动处理，无需手动对齐

修正后的完整代码 #

import datasets
import pandas as pd
import transformers
import shap

# 加载数据集
dataset = datasets.load_dataset("emotion", split="train")
data = pd.DataFrame({"text": dataset["text"], "emotion": dataset["label"]})

# 加载模型和分词器
tokenizer = transformers.AutoTokenizer.from_pretrained(
    "nateraw/bert-base-uncased-emotion", use_fast=True
)
model = transformers.AutoModelForSequenceClassification.from_pretrained(
    "nateraw/bert-base-uncased-emotion"
).cuda()

# 构建预测pipeline
pred = transformers.pipeline(
    "text-classification",
    model=model,
    tokenizer=tokenizer,
    device=0,
    return_all_scores=True,
)

# 初始化SHAP解释器并计算SHAP值
explainer = shap.Explainer(pred)
shap_values = explainer(data["text"][:10])  # 取10个样本示例

# 获取模型的类别标签
class_labels = model.config.id2label.values()

# 循环绘制每个类别的SHAP条形图
for label in class_labels:
    # 计算当前类别下，所有样本各token的平均SHAP值
    avg_shap = shap_values[:, :, label].mean(0)
    # 绘制条形图，设置标题区分类别
    shap.plots.bar(avg_shap, title=f"SHAP特征重要性：{label}")

关键说明 #

• model.config.id2label.values()能准确获取模型对应的6个情绪类别（sadness、joy、love、anger、fear、surprise）
• shap_values[:, :, label]会自动匹配对应类别的SHAP值维度，mean(0)是对所有样本的同位置token取平均
• 不同样本的token长度差异（即None维度）无需额外处理，SHAP会自动完成维度对齐和平均计算

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