2. 朴素贝叶斯垃圾邮件过滤器

任务说明

构建一个基于朴素贝叶斯的垃圾邮件过滤器,输入为邮件文本,输出为“垃圾邮件”或“正常邮件”的分类结果。本案例使用公开可获取的CSV格式邮件数据集(含标签列),完成训练、预测与评估的完整闭环。

数据集与预处理

使用spam.csv数据集(SMS Spam Collection,包含5572条短信,标签为spam/ham)。若本地无此文件,可通过fetch_spam_data()函数自动下载。预处理流程:去除标点与数字、统一小写、按空白分词,并过滤长度小于2的单词。

实现代码

import pandas as pd
import re
import urllib.request
import zipfile
import os
from collections import defaultdict
from math import log

# ---------- 数据获取 ----------
def fetch_spam_data(path="spam.csv"):
    if os.path.exists(path):
        return path
    url = "https://archive.ics.uci.edu/ml/machine-learning-databases/00228/smsspamcollection.zip"
    urllib.request.urlretrieve(url, "spam.zip")
    with zipfile.ZipFile("spam.zip") as z:
        z.extract("SMSSpamCollection")
    df = pd.read_csv("SMSSpamCollection", sep="\t", header=None, names=["label", "msg"])
    df.to_csv(path, index=False)
    return path

# ---------- 预处理 ----------
def clean_text(text):
    text = text.lower()
    text = re.sub(r"[^a-z\s]", "", text)  # 去除非字母字符
    words = [w for w in text.split() if len(w) > 1]
    return words

# ---------- 朴素贝叶斯模型 ----------
class NaiveBayesSpamFilter:
    def __init__(self):
        self.vocab = set()
        self.class_probs = {}   # 类别先验概率(对数)
        self.word_probs = {}    # 类别下单词条件概率(对数)
        self.classes = ["spam", "ham"]

    def train(self, X, y):
        # 统计类别文档数
        doc_count = defaultdict(int)
        word_count = defaultdict(lambda: defaultdict(int))  # word_count[cls][word]
        total_words = defaultdict(int)

        for msg, label in zip(X, y):
            doc_count[label] += 1
            words = clean_text(msg)
            for w in set(words):  # 每个文档内去重,避免重复计数
                word_count[label][w] += 1
                total_words[label] += 1
                self.vocab.add(w)

        total_docs = len(X)
        for cls in self.classes:
            self.class_probs[cls] = log(doc_count[cls] / total_docs)

        # 拉普拉斯平滑
        V = len(self.vocab)
        for cls in self.classes:
            self.word_probs[cls] = {}
            for w in self.vocab:
                count = word_count[cls].get(w, 0)
                self.word_probs[cls][w] = log((count + 1) / (total_words[cls] + V))

    def predict(self, text):
        words = clean_text(text)
        scores = {}
        for cls in self.classes:
            score = self.class_probs[cls]
            for w in words:
                if w in self.vocab:
                    score += self.word_probs[cls][w]
                # 未登录词忽略(拉普拉斯平滑已覆盖)
            scores[cls] = score
        return max(scores, key=scores.get)

# ---------- 训练与评估 ----------
def main():
    path = fetch_spam_data()
    df = pd.read_csv(path)
    X = df["msg"].tolist()
    y = df["label"].tolist()

    # 划分训练/测试(70/30)
    split = int(len(X) * 0.7)
    X_train, X_test = X[:split], X[split:]
    y_train, y_test = y[:split], y[split:]

    model = NaiveBayesSpamFilter()
    model.train(X_train, y_train)

    # 测试
    correct = 0
    for msg, true_label in zip(X_test, y_test):
        pred = model.predict(msg)
        if pred == true_label:
            correct += 1
    acc = correct / len(X_test)

    # 输出示例预测
    sample_msg = "Congratulations! You've won a free ticket to Bahamas. Call now!"
    sample_pred = model.predict(sample_msg)
    print(f"测试集准确率: {acc:.4f}")
    print(f"示例预测: '{sample_msg}' → {sample_pred}")

if __name__ == "__main__":
    main()

运行与验证

运行方式:将代码保存为spam_filter.py,在Python 3.8+环境中执行python spam_filter.py。依赖库:pandas、numpy(pandas依赖)。首次运行会自动下载约500KB的数据集。

验证方法:程序输出测试集准确率(预期在0.95~0.98之间,因随机划分略有波动)以及一条示例短信的分类结果。判定标准:程序能正常输出准确率数值且示例短信被正确判为spam,即视为通过。

朴素贝叶斯垃圾邮件过滤器 — 数据流与判定流程 数据获取 spam.csv (5572条) 预处理 小写化 · 去标点 分词 · 长度过滤 训练阶段 计算先验概率 计算条件概率(拉普拉斯平滑) 模型参数 P(c) · P(w|c) 测试邮件 30% 数据 预测判定 argmax P(c|w) 分类结果 spam / ham 评估 对比预测与真实标签 → 计算准确率

判定标准

  • 程序无报错运行完成,输出测试集准确率与示例预测结果。
  • 示例短信“Congratulations! You've won a free ticket...”被预测为spam。
  • 测试集准确率高于0.90(朴素贝叶斯在此数据集上的典型表现)。