朴素贝叶斯进行新闻分类

数据来源

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

以上是大体的数据，三列分别为title、content、class；由于这里讲的的不是爬虫，爬虫部分省略

项目最终结构

 其中主要逻辑在native_bayes.py文件中实现，utils.py为部分工具函数，tr_model.m为tf_idf模型，train_model为我们用朴素贝叶斯训练出来的分类模型。110.txt为预测时的文章内容

使用

通过run方法，来训练一个模型，predict方法去预测输入新闻的类别，代码的大部分解释都在注释里

代码实现

native_bayes.py

import os
import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import train_test_split
from sklearn.naive_bayes import MultinomialNB
from sklearn.externals import joblib

from two_naive_bayes.utils import cut_word, rep_invalid_char


class NaiveBayes(object):
    """朴素贝叶斯进行新闻分类，现阶段只支持(科技、娱乐、时尚)三个类别的数据分类"""
    def __init__(self, title, content):
        print('__init__')
        self.excel_reader = os.path.join(os.path.join(os.path.pardir, 'one_spider'), 'tx.xlsx')
        self.data = cut_word([rep_invalid_char(str(title)+str(content))])

    def read(self):
        """读取数据，返回特征值与目标值的列表"""
        print('read.....')
        df = pd.read_excel(self.excel_reader, 'Sheet')
        title_list = df['title'].tolist()
        content_list = df['content'].tolist()
        target = df['class'].tolist()
        data = [(str(title_list[i])+str(content_list[i])) for i in range(len(title_list))]
        data = cut_word(data)
        return data, target

    def data_split(self, data, target):
        """
        将数据分割成数据集与测试集
        :param data: 特征值
        :param target: 目标值
        :return: 数据集的特征值，测试集的特征值，数据集的目标值，测试集的目标值
        """
        print('data_split.....')
        x_train, x_test, y_train, y_test = train_test_split(data, target)
        return x_train, x_test, y_train, y_test

    def feature_extraction(self, x_train):
        """对训练集进行特征抽取，返回TfIdf处理后的数据和tf"""
        print('feature_extraction.....')
        tf = TfidfVectorizer()
        x_train = tf.fit_transform(x_train)
        self.tf = tf  # 保存此tf，后面进行特征抽取都要用到
        joblib.dump(tf, 'tf_model.m')
        return x_train

    def bayes(self, x_train, y_train):
        """贝叶斯建模"""
        print('bayes.....')
        mlt = MultinomialNB()
        mlt.fit(x_train, y_train)
        return mlt

    def test(self, mlt, x_test, y_test):
        """
        预测，返回准确率
        :param mlt: 朴素贝叶斯训练的模型
        :param x_test: 测试集的原始数据
        :param y_test: 测试集的目标值
        :return: 模型的准确率，每次运行的结果是不一样的
        """
        print('test.....')
        x_test = self.tf.transform(x_test)  # 这里用fit_transform数据集的tf，保证维度等的统一
        y_predict = mlt.predict(x_test)
        same_num = 0
        for i in range(len(y_test)):
            if y_predict[i] == y_test[i]:
                same_num += 1
        return float(same_num/len(y_test))

    def predict(self):
        """根据传入的文章标题和内容，预测该文章属于哪一个类别"""
        print('predict.....')
        # 如果本地模型存在，则从本地模型读取，否则重新进入训练过程，训练并保存模型，得出准确率，并输出文章类别
        condition_one = os.path.exists(os.path.join(os.path.curdir, 'train_model.m'))
        condition_two = os.path.exists(os.path.join(os.path.curdir, 'tf_model.m'))
        if all([condition_one, condition_two]) :
            model = joblib.load('train_model.m')
            self.tf = joblib.load('tf_model.m')
        else:
            data, target = self.read()
            x_train, x_test, y_train, y_test = self.data_split(data, target)
            x_train = self.feature_extraction(x_train)
            model = self.bayes(x_train, y_train)
            joblib.dump(model, 'train_model.m')  # 保存训练出来的模型
            number = self.test(model, x_test, y_test)
            print('模型预测的准确率为：', number)

        data = self.tf.transform(self.data)
        predict = model.predict(data)
        print('预测该文章的类别为：', predict[0])

    def run(self):
        """
        run方法用来训练模型并保存至本地，和评估准确率，predict方法虽然也有这个功能，
        但是为了训练处一个正确率较大的模型，就需要用run方法反复训练比较后得出最佳的
        模型保存至本地，在之后的预测中就不用重新训练模型而浪费时间了
        训练对比模型时，可传入空的title和content即可
        """
        # todo: =======================================================
        # 1.读取爬虫爬取的excel表数据
        # 2.使用jieba分词进行中文分词，得到总的特征值与目标值
        # 3.进行数据分割，分为数据集与测试集
        # 4.对数据集进行特征抽取 TfIdf，测试集作为测试数据也需要特征抽取
        # 5.以训练集中的列表进行重要性统计
        # 6.朴素贝叶斯算法训练模型
        # 7.模型对测试集进行预测，得出准确率
        print('run.....')
        data, target = self.read()
        x_train, x_test, y_train, y_test = self.data_split(data, target)
        x_train = self.feature_extraction(x_train)
        model = self.bayes(x_train, y_train)
        joblib.dump(model, 'train_model.m')  # 保存训练出来的模型
        number = self.test(model, x_test, y_test)
        print('模型预测的准确率为：', number)



if __name__ == '__main__':
    # 需要预测的文章标题和文章内容(可更改)
    title = '蔚来与Mobileye结盟，合作产品将于2022年面世'
    with open('110.txt', 'r') as f:
        content = f.read()
    # title: 文章标题， content: 文章正文内容
    nb = NaiveBayes(title, content)
    # nb.run()
    nb.predict()

utils.py

import re
import jieba


def cut_word(data:list):
    """
    进行中文分词
    :param data: 待分词的数据，列表类型
    :return: 分词后的数据，列表类型
    """
    res = []
    for s in data:
        word_cut = ' '.join(list(jieba.cut(s)))
        res.append(word_cut)
    return res

def rep_invalid_char(old:str):
    """只保留字符串中的大小写字母，数字，汉字"""
    invalid_char_re = r"[^0-9A-Za-zu4e00-u9fa5]"
    return re.sub(invalid_char_re, "", old)