垃圾邮件分类是一种具有广泛应用场景的二分类问题，可以利用机器学习进行解决，市场上已经有很成熟的垃圾邮件分类软件或集成在 Outlook上 的垃圾邮件过滤插件。接下来，我们使用朴素贝叶斯算法训练一个《垃圾邮件分类器》。

1. 邮件数据介绍
2. 邮件数据读取
3. 邮件数据清洗

import time
from sklearn.model_selection import train_test_split
from collections import Counter
import codecs
from tqdm import tqdm
import pickle
import multiprocessing
from joblib import Parallel
from joblib import delayed
import os
import re

1. 邮件数据介绍

数据集链接: https://plg.uwaterloo.ca/cgi-bin/cgiwrap/gvcormac/foo06

邮件数据存放在 trec06c 目录下，该目录下有 data、delay、full 三个子目录，其中 full 目录下的 index 文件中存储了所有邮件的路径，每一个路径为一个垃圾邮件，如下图所示:

第一项为标签值，第二项为邮件路径(注意: 这里使用的是相对路径)。其中的垃圾邮件内容如下:

    Received: from coozo.com ([219.133.254.230])
	by spam-gw.ccert.edu.cn (MIMEDefang) with ESMTP id j8L2Zoqi028766
	for <li@ccert.edu.cn>; Fri, 23 Sep 2005 13:01:45 +0800 (CST)
Message-ID: <200509211035.j8L2Zoqi028766@spam-gw.ccert.edu.cn>
From: "you" <you@coozo.com>
Subject: =?gb2312?B?us/X9w==?=
To: li@ccert.edu.cn
Content-Type: text/plain;charset="GB2312"
Content-Transfer-Encoding: 8bit
Date: Sun, 23 Oct 2005 23:44:32 +0800
X-Priority: 3
X-Mailer: Microsoft Outlook Express 6.00.2800.1106

 您好！ 
       我公司有多余的发票可以向外代开！（国税、地税、运输、广告、海关缴款书）。 
  
    如果贵公司（厂）有需要请来电洽谈、咨询！ 
              
               联系电话: 013510251389  陈先生
                 

                                                           谢谢


顺祝商祺!

所需要用到的程序包:

from datasets import load_dataset
import pandas as pd
import os
import codecs
import re
import zhconv
import jieba
import jieba.posseg as psg
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.model_selection import train_test_split
import pickle
import time
from tqdm import tqdm


# 将当前目录设置为工作目录
os.chdir(os.path.dirname(os.path.abspath(__file__)))

# 结巴不输出日志
jieba.setLogLevel(jieba.logging.INFO)

2. 邮件数据读取

邮件数据是以单个文件的方式存储，我们将所有的邮件数据从文件读取出来，并转存到一个文件中，方便后面的数据处理。并且，我们也会对数据划分为训练集、测试集，训练集大概有 5万+左右，测试集 1万+左右。

完整处理代码如下:

def read_mail_data():

    # 读取目录
    fnames, labels = [], []
    for line in open('../data/full/index'):
        label, path = line.strip().split()
        # 读取的路径: ../data/215/104
        # 转换为路径: data/data/215/104
        path = path.replace('..', '../data')
        fnames.append(path)
        labels.append(label)

    # 读取文件
    emails = [open(fname, encoding='gbk', errors='ignore').read() for fname in fnames]
    # 数据分布
    print('数据集分布:', dict(Counter(labels)))
    # 分割数据
    x_train, x_test, y_train, y_test = train_test_split(emails, labels, test_size=0.1, random_state=22, stratify=labels)
    print('训练集分布:', dict(Counter(y_train)))
    print('测试集分布:', dict(Counter(y_test)))

    pickle.dump({'email': x_train, 'label': y_train}, open('../temp/原始训练集.pkl', 'wb'), protocol=pickle.HIGHEST_PROTOCOL)
    pickle.dump({'email': x_test,  'label': y_test}, open('../temp/原始测试集.pkl', 'wb'), protocol=pickle.HIGHEST_PROTOCOL)

上面的代码执行成功之后，会在 temp 目录下创建两个文件:

1. 原始训练集.csv
2. 原始测试集.csv

3. 邮件数据处理

邮件中数据很杂乱，也包含了很多我们不需要的信息。这一步我们主要对每一封邮件内容进行一些处理工作。我们这里只进行简单的邮件内容分词处理.

完整处理代码如下:

# 邮件内容清洗
def clean_email(content):

    import jieba
    jieba.setLogLevel(0)
    # 保留中文
    # content = re.sub(r'[^\u4e00-\u9fa5]', '', content)
    # 内容分词
    content = jieba.lcut(content)

    return ' '.join(content)


def clean_train_data(input_emails, input_labels, worker_count=None):

    # 数据较多，清洗需要10分钟左右
    data_count = len(input_labels)
    # 设置 cpu 数量
    worker_count = worker_count if worker_count is not None else multiprocessing.cpu_count()
    # 设置子任务数据
    task_range = list(range(0, data_count + 1, int(data_count / worker_count)))

    def task(s, e):
        emails = input_emails[s:e]
        labels = input_labels[s:e]
        result_emails = []
        result_labels = []
        progress = tqdm(range(len(labels)), desc='进程 %6d 数据区间 (%5d, %5d)' % (os.getpid(), s, e))
        filter_number = 0
        for email, label in zip(emails, labels):
            email = clean_email(email)
            progress.update(1)
            if len(email) == 0:
                filter_number += 1
                continue
            result_emails.append(email)
            result_labels.append(label)
        progress.close()

        return {'email': result_emails, 'label': result_labels, 'filter': filter_number}

    delayed_tasks = []
    for index in range(1, len(task_range)):
        s = task_range[index - 1]
        e = task_range[index]
        delayed_tasks.append(delayed(task)(s, e))

    # 多任务运行任务
    # 16核心需要81秒，6核需要140秒
    results = Parallel(n_jobs=worker_count)(delayed_tasks)
    # 合并计算结果
    clean_emails = []
    clean_labels = []
    clean_number = 0
    for result in results:
        clean_emails.extend(result['email'])
        clean_labels.extend(result['label'])
        clean_number += result['filter']

    print('训练集数量:', data_count)
    print('清理后数量:', len(clean_labels), len(clean_emails))
    print('被清理数量:', clean_number)

    # 预览数据样式
    print('预览数据:', clean_emails[0])
    print('预览标签:', clean_labels[0])

    return {'email': clean_emails, 'label': clean_labels}


def prepare_mail_data():

    train_data = pickle.load(open('../temp/原始训练集.pkl', 'rb'))
    test_data = pickle.load(open('../temp/原始测试集.pkl', 'rb'))

    train_result = clean_train_data(train_data['email'], train_data['label'])
    test_result =  clean_train_data(test_data['email'],  test_data['label'])

    # 存储清洗结果
    pickle.dump(train_result, open('../temp/清洗训练集.pkl', 'wb'), protocol=pickle.HIGHEST_PROTOCOL)
    pickle.dump(test_result, open('../temp/清洗测试集.pkl', 'wb'), protocol=pickle.HIGHEST_PROTOCOL)

当执行 prepare_mail_data 函数之后，会在 temp 目录下创建两个文件：

1. 清洗训练集.pkl
2. 清洗测试集.pkl