基于BERT文本分类模型为38个类别提取Top N特征词

实现思路 #

通过CountVectorizer从训练语料构建高频词汇表，将每个词汇单独输入训练好的BERT分类模型，获取分类层的激活值（对应每个类别的输出得分），最终基于激活值排序得到每个类别下的Top N特征词。

完整代码实现 #

from sklearn.feature_extraction.text import TfidfVectorizer, CountVectorizer
import pandas as pd 
import torch, os
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# 设备配置
if torch.cuda.is_available():
    device = torch.device("cuda")
    print('GPU:', torch.cuda.get_device_name(0))
else:
    device = torch.device("cpu")

# 获取指定层的激活值
def get_hidden_state_activations(input_ids, model, layer_idx):
    with torch.no_grad():
        outputs = model(input_ids)
        # AutoModelForSequenceClassification的outputs[0]为分类层logits
        activations = outputs[layer_idx]
        return activations

# 1. 从训练语料构建高频词汇表
vectorizer = CountVectorizer(ngram_range=(1,1), max_features=10000, min_df=50)
vecs = vectorizer.fit_transform(corpus)  # corpus为你的训练语料列表
print('统计词汇频率')
dense = vecs.todense()
lst1 = dense.tolist()
feature_names = vectorizer.get_feature_names_out()

# 整理词汇及频率，过滤纯数字词汇
df = pd.DataFrame(lst1, columns=feature_names)
df = df.T.sum(axis=1).sort_values(ascending=False).reset_index()
df.columns = ['word', 'frequency']
df = df[~df.word.str.isdigit()].reset_index(drop=True)

# 2. 加载训练好的BERT分类模型与分词器
TOKENIZER_MODEL_NAME = "你的分词器模型名称"  # 例如'bert-base-chinese'
BERT_MODEL_NAME = "你的BERT分类模型路径/名称"
tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(BERT_MODEL_NAME).to(device)

# 准备词汇数据与类别标签（直接从模型配置读取）
data = df.word.tolist()
labels = list(model.config.id2label.values())

# 3. 将词汇转换为模型输入格式
input_ids = [tokenizer.encode(example, add_special_tokens=True, return_tensors='pt') for example in data]

# 4. 获取每个词汇对应的分类层激活值
layer_idx = 0  # 对应分类层logits输出
activations = [get_hidden_state_activations(input_id.to(device), model, layer_idx) for input_id in tqdm(input_ids)]

# 5. 整理激活值到DataFrame
activations_df = pd.DataFrame([i.cpu().detach().numpy().reshape(-1) for i in activations])
activations_df.columns = labels
activations_df['word'] = data
activations_df = activations_df[['word'] + labels]

# 6. 提取每个类别的Top N特征词
N = 10  # 定义每个类别要提取的特征词数量
top_n_words_per_category = {}

for category in labels:
    # 按当前类别的激活值降序排序，取前N个词汇
    sorted_words = activations_df.sort_values(by=category, ascending=False)[['word', category]].head(N)
    top_n_words_per_category[category] = sorted_words['word'].tolist()

# 打印结果示例
for cat, words in top_n_words_per_category.items():
    print(f"类别 {cat} 的Top {N} 特征词: {', '.join(words)}")

关键说明 #

• 词汇表构建：用CountVectorizer过滤低频词汇（min_df=50），仅保留高频且非纯数字的词汇，减少噪声干扰。
• 激活值获取：AutoModelForSequenceClassification的outputs[0]直接对应分类层的logits输出，无需使用layer_idx=-1。
• Top N提取：遍历每个类别列，按激活值降序排序后提取前N个词汇，即可得到对应类别的特征词。

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