朴素贝叶斯(MultinomialNB)-垃圾邮件分类

一、概述

•  版本：python:2.7.15
• 第三方库：scikit-learn
• 数据样本：https://gitee.com/chenfenli/document/blob/master/learning/sms_spam.txt

 二、代码

# coding=UTF-8
import codecs
from sklearn.naive_bayes import MultinomialNB
from sklearn.feature_extraction.text import CountVectorizer
    
if __name__ == "__main__":
    # 1、构建数据
    labels = [] # 训练集label
    corpus = [] # 训练集文本
    labels_test = []    # 测试集label
    corpus_test = []    # 测试集文本
    file = codecs.open("/Users/chenfenli/Desktop/sms_spam.txt")
    count = 0
    while True:
        # 从文件中按行读取
        line = file.readline()
        if line:
            count = count + 1
            line = line.split(",")
            label = line[0]
            sentence = line[1]
            
            # 训练集
            corpus.append(sentence)
            if "ham" == label:
                # 正常邮件
                labels.append(0)
            elif "spam" == label:
                # 垃圾邮件
                labels.append(1)
            
            # 测试集
            if count > 5550:
                corpus_test.append(sentence)
                if "ham" == label:
                    labels_test.append(0)
                elif "spam" == label:
                    labels_test.append(1)
        else:
            break
    
    # 2、文本特征提取
    #   将文本数据转化成特征向量的过程
    #   比较常用的文本特征表示法为词带法
    #
    # 词带法
    #   不考虑词语出现的顺序，每个出现过的词汇单独作为一列特征
    #   这些不重复的特征词汇集合为词表
    #   每一个文本都可以在很长的词表上统计出一个很多列的特征向量
    #   CountVectorizer是将文本向量转换成稀疏表示数值向量（字符频率向量），vectorizer 将文档词块化，只考虑词汇在文本中出现的频率
    # 词带
    vectorizer = CountVectorizer()
    fea_train = vectorizer.fit_transform(corpus) 
    #print(vectorizer.get_feature_names())
    #print(fea_train.toarray())

    # 3、训练模型：MultinomialNB(基于贝叶斯定理的朴素贝叶斯分类器)
    clf = MultinomialNB(alpha = 1) 
    clf.fit(fea_train,labels)

    # 4、准备测试数据
    vectorizer2 = CountVectorizer(vocabulary=vectorizer.vocabulary_)
    fea_test =  vectorizer2.fit_transform(corpus_test)

    # 5、预测
    pred = clf.predict(fea_test)
    for p in pred:
        if p == 0:
            print("正常邮件")
        elif p == 1:
            print("垃圾邮件")