基于Python与scikit-learn的SVM文本分类：获取垃圾/正常邮件高频词

你可以通过统计训练集中两类邮件的词频来获取高频词，以下是基于你现有代码的修改方案：

步骤说明 #

• 从训练数据中分离垃圾邮件（spam）和正常邮件（ham）的文本
• 利用CountVectorizer的词汇映射，结合词频矩阵统计每个类别的词频
• 提取并展示每个类别中频率最高的词汇

修改后的完整代码 #

import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import CountVectorizer
from sklearn import svm
import numpy as np

# 依赖安装
# pip install pandas
# pip install -U scikit-learn

spam = pd.read_csv('Cartel1.csv')
z = spam['v2']
y = spam["v1"]

# 拆分训练集和测试集
z_train, z_test, y_train, y_test = train_test_split(z, y, test_size=0.2)

# 文本转特征向量
cv = CountVectorizer()
features = cv.fit_transform(z_train)
svm_model = svm.SVC()
svm_model.fit(features, y_train)
features_test = cv.transform(z_test)

# ------------------- 新增：获取高频词部分 -------------------
# 获取词汇表（索引到词的映射）
vocab = cv.get_feature_names_out()

# 拆分训练集中的spam和ham对应的特征矩阵
spam_features = features[y_train == 'spam']
ham_features = features[y_train == 'ham']

# 统计每个词在spam中的总出现次数
spam_word_counts = np.array(spam_features.sum(axis=0)).flatten()
# 统计每个词在ham中的总出现次数
ham_word_counts = np.array(ham_features.sum(axis=0)).flatten()

# 生成词-频率的字典并按频率降序排序
spam_word_freq = sorted(zip(vocab, spam_word_counts), key=lambda x: x[1], reverse=True)
ham_word_freq = sorted(zip(vocab, ham_word_counts), key=lambda x: x[1], reverse=True)

# 输出前20个高频词
print("=== 垃圾邮件高频词（Top 20） ===")
for word, count in spam_word_freq[:20]:
    print(f"{word}: {count}")

print("\n=== 正常邮件高频词（Top 20） ===")
for word, count in ham_word_freq[:20]:
    print(f"{word}: {count}")

# ------------------- 原有预测代码保留 -------------------
comment = ["Sexy free Call and text messages on 08002986030"]
vect = cv.transform(comment)
print("\nThis comment: ", comment, " is: ", svm_model.predict(vect))  # spam

comment2 = ["Hi there, I am emailing you today to let you know we have created a new task for you."]
vect2 = cv.transform(comment2)
print("This comment: ", comment2, " is: ", svm_model.predict(vect2))  # ham --no spam
# print(svm_model.score(features_test, y_test))

代码解释 #

1. 词汇表获取：cv.get_feature_names_out()获取所有被统计的词汇，对应特征矩阵的每一列
2. 类别特征拆分：根据训练集的标签，分别筛选出垃圾邮件和正常邮件的特征矩阵
3. 词频统计：对每个类别的特征矩阵按列求和，得到每个词在该类别中的总出现次数
4. 排序展示：将词和对应的频率配对后按降序排序，取前N个展示高频词

注意事项 #

• 如果你想过滤掉无意义的停用词（比如"the"、"and"），可以在初始化CountVectorizer时添加参数stop_words='english'，即cv = CountVectorizer(stop_words='english')，这样统计出的高频词会更有实际意义
• 这里的统计基于训练集数据，确保了结果和你的SVM模型训练数据一致，避免数据泄露问题

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