吴沂钊
70 lines
2.8 KiB
Python
70 lines
2.8 KiB
Python
import pandas as pd
|
||
import re
|
||
import jieba
|
||
import pickle
|
||
from wordcloud import WordCloud#安装wordcloud要输入后面的那个，而不是前面的，就像安装opencv，要装cv2
|
||
from sklearn.naive_bayes import GaussianNB
|
||
from sklearn.model_selection import train_test_split
|
||
from sklearn.feature_extraction.text import CountVectorizer, TfidfTransformer
|
||
from sklearn.metrics import confusion_matrix
|
||
import seaborn as sns
|
||
import matplotlib.pyplot as plt
|
||
plt.rcParams['font.sans-serif'] = 'SimHei'
|
||
plt.rcParams['axes.unicode_minus'] = False
|
||
data = pd.read_csv('message80W1.csv', header=None, index_col=0)
|
||
data.columns = ['label', 'message']
|
||
n = 5000
|
||
a = data[data['label'] == 0].sample(n)
|
||
b = data[data['label'] == 1].sample(n)
|
||
newdata = pd.concat([a, b], axis=0)
|
||
dupdata = newdata['message'].drop_duplicates()
|
||
qumindata = dupdata.apply(lambda x: re.sub('x', '', x))
|
||
# 读取自定义的分词词典
|
||
jieba.load_userdict('newdic1.txt')
|
||
data_cut = qumindata.apply(lambda x: jieba.lcut(x))
|
||
# 读取停用词表去停用词
|
||
stopWords = pd.read_csv('stopword.txt', encoding='GB18030', sep='hahaha', header=None)
|
||
stopWords = ['≮', '≯', '≠', '≮', ' ', '会', '月', '日', '–'] + list(stopWords.iloc[:, 0])
|
||
dataafterstop = data_cut.apply(lambda x: [i for i in x if i not in stopWords])
|
||
labels = newdata.loc[dataafterstop.index, 'label']
|
||
# join函数将分词后的数据转换为字符串
|
||
adata = dataafterstop.apply(lambda x: ' '.join(x))
|
||
# 划分数据集
|
||
data_tr, data_te, labels_tr, labels_te = train_test_split(adata, labels, test_size=0.15)
|
||
# 转词向量
|
||
countVectorizer = CountVectorizer()
|
||
# 训练集
|
||
data_tr = countVectorizer.fit_transform(data_tr)
|
||
X_tr = TfidfTransformer().fit_transform(data_tr.toarray()).toarray()
|
||
# 测试集
|
||
data_te = CountVectorizer(vocabulary=countVectorizer.vocabulary_).fit_transform(data_te)
|
||
X_te = TfidfTransformer().fit_transform(data_te.toarray()).toarray()
|
||
model = GaussianNB()
|
||
model.fit(X_tr, labels_tr)
|
||
pred = model.predict(X_te)
|
||
score = model.score(X_te, labels_te)
|
||
# 使用pickle模块保存模型
|
||
with open('model.pkl', 'wb') as f:
|
||
pickle.dump(model, f)
|
||
# 保存词向量模型
|
||
with open('countVectorizer.pkl', 'wb') as f:
|
||
pickle.dump(countVectorizer, f)
|
||
# 绘制混淆矩阵
|
||
sns.heatmap(confusion_matrix(labels_te,pred),annot=True)
|
||
plt.xlabel('Predicted')
|
||
plt.ylabel('True')
|
||
plt.title("混淆矩阵")
|
||
plt.show()
|
||
## 绘制词云
|
||
word_fre = {}
|
||
for i in dataafterstop[labels == 0]:
|
||
for j in i:
|
||
if j not in word_fre.keys():
|
||
word_fre[j] = 1
|
||
else:
|
||
word_fre[j] += 1
|
||
wc = WordCloud( background_color='white', font_path=r'C:/Windows/Fonts/STKAITI.TTF')#换一个长方形的边框，同时变成华文楷体
|
||
wc.fit_words(word_fre)
|
||
plt.imshow(wc)
|
||
plt.show()
|