朴素贝叶斯垃圾邮件分类器Python实现

import math
import os
import re
import string
from collections import defaultdict

from sklearn.model_selection import train_test_split

DATA_DIR = 'enron'  # 数据集地址
target_names = ['ham', 'spam']  # 正常、 垃圾
stopwords = set(open('stopwords/stopwords.txt', 'r').read().splitlines())  # .splitlines() 按换行符分割


def get_data(DATA_DIR):
    subfolders = ['enron%d' % i for i in range(1, 7)]  # 返回enron1-enron7列表
    data = []
    target = []
    for subfolder in subfolders:  # 对enron1-enron7依次遍历
        # spam
        spam_files = os.listdir(os.path.join(DATA_DIR, subfolder, 'spam'))  # 获文件夹下的所有文件 返回一个列表
        for spam_file in spam_files:  # 依次打开返回的文件
            with open(os.path.join(DATA_DIR, subfolder, 'spam', spam_file), encoding="latin-1") as f:
                data.append(f.read())
                target.append(1)  # 垃圾邮件标签 为1
        # ham
        ham_files = os.listdir(os.path.join(DATA_DIR, subfolder, 'ham'))
        for ham_file in ham_files:
            with open(os.path.join(DATA_DIR, subfolder, 'ham', ham_file), encoding="latin-1") as f:
                data.append(f.read())
                target.append(0)  # 正常邮件标签 为0
    return data, target


def preprocess(text):
    '''
    对文本进行处理，包括去除标点、转为小写、去除停用词等操作。
    '''
    text = text.lower()  # 转为小写
    text = re.sub(f'[{string.punctuation}]', ' ', text)  # 去除标点符号
    text = [word for word in text.split() if word not in stopwords]  # 去除停用词
    return text


class NaiveBayesClassifier():
    def __init__(self):
        self.vocabulary = set()  # 词汇表
        self.class_total = defaultdict(int)  # 每个类别的文档数
        self.word_total = defaultdict(int)  # 每个类别中所有单词出现次数之和
        self.word_given_class = defaultdict(lambda: defaultdict(int))  # 每个类别中每个单词出现次数

    def fit(self, X, y):
        '''
        训练分类器，其中 X 为训练集数据，y 为训练集标签。
        '''
        for text, label in zip(X, y):
            words = preprocess(text)  # 对文本进行处理，包括去除标点、转为小写、去除停用词等操作。
            self.class_total[label] += 1  # 该类别文档数加1
            for word in words:  # 遍历每一个单词
                self.vocabulary.add(word)  # 加入词汇表
                self.word_given_class[label][word] += 1  # 该类别中该单词出现次数加1
                self.word_total[label] += 1  # 该类别所有单词出现次数之和加1

    def predict(self, X):
        '''
        对新邮件进行分类，其中 X 为待分类邮件。
        '''
        log_priors = {}  # 存储每个类别的先验概率的对数（用于避免下溢）
        for c in self.class_total.keys():
            #  log（每个类别的文档数/文档数）  也就是训练集正常邮件和垃圾邮件占总邮件的比例
            log_priors[c] = math.log(self.class_total[c] / sum(self.class_total.values()))

        predictions = []
        for text in X:  # 遍历测试集的每一个邮件
            words = preprocess(text)  # 对文本进行处理，包括去除标点、转为小写、去除停用词等操作。

            log_probs = {}
            for c in self.class_total.keys():  # 遍历两个类别  正常和垃圾邮件类别
                log_probs[c] = log_priors[c]
                for word in words:  # 遍历每个单词
                    if word in self.vocabulary:  # 如果测试邮件的这个单词在对应类别词汇表中
                        # 计算条件概率的对数
                        # 这里需要说明一下： +1是为了防止有些单词在测试样本中出现而垃圾邮件或正常邮件中没有导致概率为0 的情况
                        # 取log是为了防止概率太小 导致最后结果约等于0
                        log_probs[c] += math.log((self.word_given_class[c][word] + 1) / (self.word_total[c] + len(self.vocabulary)))
            predictions.append(max(log_probs, key=log_probs.get))  # 取概率最大值对应的类别为预测结果

        return predictions


def get_classifier():
    # 读取数据
    X, y = get_data(DATA_DIR)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)  # 划分测试集和训练集

    # 训练分类器
    clf = NaiveBayesClassifier()  # 实例化分类器
    clf.fit(X_train, y_train)  # 送入数据训练
    return clf
解释说明以上代码内容：该代码实现了朴素贝叶斯分类器对垃圾邮件的分类。具体实现步骤如下：

1.定义了数据集路径、目标类别（正常邮件和垃圾邮件）、停用词（即无意义的词汇，如“的”、“是”等）。

2.定义了读取数据的函数get_data()，该函数遍历每个邮件文件夹，读取每个邮件的内容和标签，并将它们存储在两个列表中。

3.定义了文本预处理函数preprocess()，该函数将邮件转换为小写字母、去除标点符号和停用词。

4.实现了NaiveBayesClassifier类，该类包含fit()和predict()方法。fit()方法用于训练分类器，它遍历训练数据集中的每个邮件并计算每个类别的先验概率、每个单词在每个类别中出现的次数以及每个类别中所有单词出现次数之和。predict()方法用于对新邮件进行分类。对于每个测试邮件，该方法计算每个类别的条件概率，并选择具有最大条件概率的类别作为预测结果。

5.定义了get_classifier()函数，该函数使用get_data()函数读取数据集并使用train_test_split()函数划分训练集和测试集。然后实例化NaiveBayesClassifier类并将训练集送入fit()方法进行训练，最后返回训练好的分类器。