调用sklearn的NB模型来实现文章分类

在上一篇文章里面,我自己手写了一个朴素贝叶斯算法(NB),进行了训练和验证,并且通过在网上随便找一篇文章进行测试,能够正确得到文章类型结果。这篇文章中,我打算采用sklearn库,直接调用NB模型进行训练和预测。

1、首先加载sklearn的库;train_test_split--对数据进行分割,分成训练集和测试集。TfidfVectorizer--生成文章单词的TFIDF矩阵。MultinomialNB--sklearn的NB模型。

from sklearn.model_selection import train_test_split
from sklearn.naive_bayes import MultinomialNB
from sklearn.feature_extraction.text import TfidfVectorizer

 其次,准备数据集,读取分词后的文章数据集,加载列表words和label中,通过train_test_split按照4:1的比例进行切分。

调用sklearn的NB模型来实现文章分类_第1张图片

接下来,对切分好的训练集进行向量化,使用TfidfVectorizer,求出文章单词的tfidf,构建单词向量。然后通过MultinomialNB模型对训练集的向量数据进行训练,并且保存TfidfVectorizer和MultinomialNB模型。

最后,读取测试集中的文章数据,通过TfidfVectorizer模型,进行测试集向量化,放入MultinomialNB模型中进行预测。

调用sklearn的NB模型来实现文章分类_第2张图片

2、同样,拿出上篇文章中测试用例,通过建立好的模型进行预测,看下这篇新闻文章归属于哪类。

调用sklearn的NB模型来实现文章分类_第3张图片

在pycharm远程运行下得到预测结果属于“娱乐”类。

你可能感兴趣的:(文章分类,机器学习,自然语言处理)