数据来源
通过爬虫,爬取腾讯新闻三个分类每个分类大约1000条左右数据,存入excel

以上是大体的数据,三列分别为title、content、class;由于这里讲的的不是爬虫,爬虫部分省略
项目最终结构

其中主要逻辑在native_bayes.py文件中实现,utils.py为部分工具函数,tr_model.m为tf_idf模型,train_model为我们用朴素贝叶斯训练出来的分类模型。110.txt为预测时的文章内容
使用
通过run方法,来训练一个模型,predict方法去预测输入新闻的类别,代码的大部分解释都在注释里
代码实现
native_bayes.py
1 import os
2 import pandas as pd
3 from sklearn.feature_extraction.text import TfidfVectorizer
4 from sklearn.model_selection import train_test_split
5 from sklearn.naive_bayes import MultinomialNB
6 from sklearn.externals import joblib
7
8 from two_naive_bayes.utils import cut_word, rep_invalid_char
9
10
11 class NaiveBayes(object):
12 """朴素贝叶斯进行新闻分类,现阶段只支持(科技、娱乐、时尚)三个类别的数据分类"""
13 def __init__(self, title, content):
14 print('__init__')
15 self.excel_reader = os.path.join(os.path.join(os.path.pardir, 'one_spider'), 'tx.xlsx')
16 self.data = cut_word([rep_invalid_char(str(title)+str(content))])
17
18 def read(self):
19 """读取数据,返回特征值与目标值的列表"""
20 print('read.....')
21 df = pd.read_excel(self.excel_reader, 'Sheet')
22 title_list = df['title'].tolist()
23 content_list = df['content'].tolist()
24 target = df['class'].tolist()
25 data = [(str(title_list[i])+str(content_list[i])) for i in range(len(title_list))]
26 data = cut_word(data)
27 return data, target
28
29 def data_split(self, data, target):
30 """
31 将数据分割成数据集与测试集
32 :param data: 特征值
33 :param target: 目标值
34 :return: 数据集的特征值,测试集的特征值,数据集的目标值,测试集的目标值
35 """
36 print('data_split.....')
37 x_train, x_test, y_train, y_test = train_test_split(data, target)
38 return x_train, x_test, y_train, y_test
39
40 def feature_extraction(self, x_train):
41 """对训练集进行特征抽取,返回TfIdf处理后的数据和tf"""
42 print('feature_extraction.....')
43 tf = TfidfVectorizer()
44 x_train = tf.fit_transform(x_train)
45 self.tf = tf # 保存此tf,后面进行特征抽取都要用到
46 joblib.dump(tf, 'tf_model.m')
47 return x_train
48
49 def bayes(self, x_train, y_train):
50 """贝叶斯建模"""
51 print('bayes.....')
52 mlt = MultinomialNB()
53 mlt.fit(x_train, y_train)
54 return mlt
55
56 def test(self, mlt, x_test, y_test):
57 """
58 预测,返回准确率
59 :param mlt: 朴素贝叶斯训练的模型
60 :param x_test: 测试集的原始数据
61 :param y_test: 测试集的目标值
62 :return: 模型的准确率,每次运行的结果是不一样的
63 """
64 print('test.....')
65 x_test = self.tf.transform(x_test) # 这里用fit_transform数据集的tf,保证维度等的统一
66 y_predict = mlt.predict(x_test)
67 same_num = 0
68 for i in range(len(y_test)):
69 if y_predict[i] == y_test[i]:
70 same_num += 1
71 return float(same_num/len(y_test))
72
73 def predict(self):
74 """根据传入的文章标题和内容,预测该文章属于哪一个类别"""
75 print('predict.....')
76 # 如果本地模型存在,则从本地模型读取,否则重新进入训练过程,训练并保存模型,得出准确率,并输出文章类别
77 condition_one = os.path.exists(os.path.join(os.path.curdir, 'train_model.m'))
78 condition_two = os.path.exists(os.path.join(os.path.curdir, 'tf_model.m'))
79 if all([condition_one, condition_two]) :
80 model = joblib.load('train_model.m')
81 self.tf = joblib.load('tf_model.m')
82 else:
83 data, target = self.read()
84 x_train, x_test, y_train, y_test = self.data_split(data, target)
85 x_train = self.feature_extraction(x_train)
86 model = self.bayes(x_train, y_train)
87 joblib.dump(model, 'train_model.m') # 保存训练出来的模型
88 number = self.test(model, x_test, y_test)
89 print('模型预测的准确率为:', number)
90
91 data = self.tf.transform(self.data)
92 predict = model.predict(data)
93 print('预测该文章的类别为:', predict[0])
94
95 def run(self):
96 """
97 run方法用来训练模型并保存至本地,和评估准确率,predict方法虽然也有这个功能,
98 但是为了训练处一个正确率较大的模型,就需要用run方法反复训练比较后得出最佳的
99 模型保存至本地,在之后的预测中就不用重新训练模型而浪费时间了
100 训练对比模型时,可传入空的title和content即可
101 """
102 # todo: =======================================================
103 # 1.读取爬虫爬取的excel表数据
104 # 2.使用jieba分词进行中文分词,得到总的特征值与目标值
105 # 3.进行数据分割,分为数据集与测试集
106 # 4.对数据集进行特征抽取 TfIdf,测试集作为测试数据也需要特征抽取
107 # 5.以训练集中的列表进行重要性统计
108 # 6.朴素贝叶斯算法训练模型
109 # 7.模型对测试集进行预测,得出准确率
110 print('run.....')
111 data, target = self.read()
112 x_train, x_test, y_train, y_test = self.data_split(data, target)
113 x_train = self.feature_extraction(x_train)
114 model = self.bayes(x_train, y_train)
115 joblib.dump(model, 'train_model.m') # 保存训练出来的模型
116 number = self.test(model, x_test, y_test)
117 print('模型预测的准确率为:', number)
118
119
120
121 if __name__ == '__main__':
122 # 需要预测的文章标题和文章内容(可更改)
123 title = '蔚来与Mobileye结盟,合作产品将于2022年面世'
124 with open('110.txt', 'r') as f:
125 content = f.read()
126 # title: 文章标题, content: 文章正文内容
127 nb = NaiveBayes(title, content)
128 # nb.run()
129 nb.predict()
utils.py
1 import re 2 import jieba 3 4 5 def cut_word(data:list): 6 """ 7 进行中文分词 8 :param data: 待分词的数据,列表类型 9 :return: 分词后的数据,列表类型 10 """ 11 res = [] 12 for s in data: 13 word_cut = ' '.join(list(jieba.cut(s))) 14 res.append(word_cut) 15 return res 16 17 def rep_invalid_char(old:str): 18 """只保留字符串中的大小写字母,数字,汉字""" 19 invalid_char_re = r"[^0-9A-Za-z\u4e00-\u9fa5]" 20 return re.sub(invalid_char_re, "", old)