sklearn实现基于TF-IDF的KNN新闻标题文本分类
数据集下载
数据集已经分好了训练集和测试集，一共有10个类别。训练集每个类别有180000条数据，测试集每个类别有1000条数据。
读取数据集
import pandas as pd
# 读取训练集和测试集，并进行分词处理
# train_text[0]是新闻文本 train_text[1]是类别标签
train_text = pd.read_table("data/新闻标题文本分类/4分类_72000/train_4_72000.txt", header=None, sep=" ")
test_text = pd.read_table("data/新闻标题文本分类/4分类_72000/test_4_4000.txt", header=None, sep=" ")
print(train_text.head(5))
print("*"*50)
print(test_text.head(5))本次演示使用的是4分类的数据集。运行上面代码，得到下面的输出：
0 1
0 中华女子学院：本科层次仅1专业招男生 3
1 卡佩罗：告诉你德国脚生猛的原因 不希望英德战踢点球 7
2 《赤壁OL》攻城战诸侯战硝烟又起 8
3 上海2010上半年四六级考试报名4月8日前完成 3
4 李永波称李宗伟难阻林丹取胜 透露谢杏芳有望出战 7
**************************************************
0 1
0 词汇阅读是关键 08年考研暑期英语复习全指南 3
1 中国人民公安大学2012年硕士研究生目录及书目 3
2 日本地震：金吉列关注在日学子系列报道 3
3 名师辅导：2012考研英语虚拟语气三种用法 3
4 自考经验谈：自考生毕业论文选题技巧 3可以看到数据集主要由两部分组成，前面部分是新闻标题文本，后面部分是类别标签。
使用pandas加载后train_text[0]和test_text[0]代表的就是新闻标题文本，train_text[1]和test_text[1]代表的是类别标签。
中文分词
使用jieba对新闻标题文本分词。jieba一共提供了四种分词模式，默认使用的是精确模式，效果还不错。此外，jieba还提供了paddle模式的分词，我自己使用下来的效果是和精确模式差不多。（不过paddle模式的分词感觉确实要比精确模式好一点的。）
import jieba
# 使用jieba分词
def new_cut(text):
return " ".join(list(jieba.cut(train_text)))
train_text["分词结果"] = train_text[0].map(new_cut)
test_text["分词结果"] = test_text[0].map(new_cut)
print(train_text["分词结果"].head(5))分词结果：
0 中华 女子 学院 ： 本科 层次 仅 1 专业 招 男生
1 卡佩罗 ： 告诉 你 德国 脚 生猛 的 原因 不 希望 英德 战 踢 点球
2 《 赤壁 OL 》 攻城战 诸侯 战 硝烟 又 起
3 上海 2010 上半年 四六级 考试 报名 4 月 8 日前 完成
4 李永波 称 李宗伟 难 阻林丹 取胜 透露 谢杏芳 有望 出战去除停用词
本次使用的停用词表是哈工大停用词表，点我下载
# 加载停用词
def get_custom_stopwords(stop_words_file):
with open(stop_words_file, encoding="gbk") as f:
stopwords = f.read()
stopwords_list = stopwords.split('\n')
custom_stopwords_list = [i for i in stopwords_list]
return custom_stopwords_list
stop_words_file = "stopwords.txt"
stopwords = get_custom_stopwords(stop_words_file)TF-IDF算法提取文本特征
具体TfidfVectorizer的配置参数可以去参考[官方文档](sklearn.feature_extraction.text.TfidfVectorizer — scikit-learn 1.2.2 documentation)，或者其他的博客。
from sklearn.feature_extraction.text import TfidfVectorizer
# 使用TF-IDF将文本处理成向量格式
tfidf_vectorizer = TfidfVectorizer(max_df=0.8,
min_df=3,
token_pattern=u'(?u)\\b[^\\d\\W]\\w+\\b',
stop_words=frozenset(stopwords),
# max_features=1000) # 最大提取特征，默认就是没有限制
)
train_weight = tfidf_vectorizer.fit_transform(train_text["分词结果"])
test_weight = tfidf_vectorizer.transform(test_text["分词结果"])
print("train:\n", f"数据集大小：{train_weight.toarray().shape[0]} 提取特征数量：{train_weight.toarray().shape[1]}")
print("test:\n", f"数据集大小：{test_weight.toarray().shape[0]} 提取特征数量：{test_weight.toarray().shape[1]}")输出结果：
train:
数据集大小：72000 提取特征数量：20677
test:
数据集大小：4000 提取特征数量：20677KNN分类器的设计
KNN分类器的K值（n_neighbors参数）取3，距离计算方式（metric参数）选择余弦距离，其他的参数默认就行。
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import accuracy_score, recall_score, f1_score, precision_score
import pickle
save_model = False
k = 3
knn = KNeighborsClassifier(n_neighbors=k, metric="cosine")
# 训练
knn.fit(train_weight, train_text[1])
# 预测
test_pred = knn.predict(test_weight)
# 计算准确率、召回率、精确率和F1值
accuracy = accuracy_score(test_pred, test_text[1])
recall = recall_score(test_pred, test_text[1], average='macro')
precision = precision_score(test_pred, test_text[1], average='macro')
f1 = f1_score(test_pred, test_text[1], average='macro')
# 输出训练日志
with open("output_logs/knn_cls.txt", "a+") as log_txt:
log_txt.write(f"k = {k}\n")
log_txt.write(f"accuracy = {accuracy}\n")
log_txt.write(f"recall = {recall}\n")
log_txt.write(f"precision = {precision}\n")
log_txt.write(f"f1 = {f1}\n")
log_txt.close()
# 保存训练模型
if save_model:
with open(f"output_models/knn_{k}.pkl", "wb") as f:
pickle.dump(knn, f)
print("Save model sucess!")结果日志输出：
k = 3
accuracy = 0.8945
recall = 0.9015750424563225
precision = 0.8945
f1 = 0.8939859931730288现在的模型分类效果可以在调调参数，比如把k值调大一点等。
