基于BERT对中文邮件内容分类

项目背景与意义

本文是《用BERT做中文邮件内容分类》系列的第二篇，该系列项目持续更新中。系列的起源是《使用PaddleNLP识别垃圾邮件》项目，旨在解决企业面临的垃圾邮件问题，通过深度学习方法探索多语言垃圾邮件的内容、标题提取与分类识别。

在本篇文章中，我们使用PaddleNLP的BERT预训练模型，根据提取的中文邮件内容判断邮件是否为垃圾邮件。该项目的思路在于基于前一篇项目的中文邮件内容提取，在98.5%的垃圾邮件分类器基线上，通过BERT的finetune进一步提升性能。

项目思路

在《使用PaddleNLP识别垃圾邮件（一）》项目的基础上，我们使用BERT进行finetune，力求在LSTM的98.5%的基线上进一步提升准确率。同时，文章中详细介绍了BERT模型的原理和PaddleNLP对BERT模型的应用，读者可以参考项目PaddleNLP2.0：BERT模型的应用进行更深入的了解。

本项目参考了陆平老师的项目应用BERT模型做短文本情绪分类（PaddleNLP 2.0），但由于PaddleNLP版本迭代的原因，进行了相应的调整和说明。

数据集介绍

我们使用了TREC 2006 Spam Track Public Corpora，这是一个公开的垃圾邮件语料库，包括英文数据集（trec06p）和中文数据集（trec06c）。在本项目中，我们仅使用了TREC 2006提供的中文数据集进行演示。数据集来源于真实邮件，保留了邮件的原有格式和内容。

除了TREC 2006外，还有TREC 2005和TREC 2007的英文垃圾邮件数据集，但本项目仅使用了TREC 2006提供的中文数据集。数据集文件目录形式如下：

trec06c
│
├── data
│   │   000
│   │   001
│   │   ...
│   └───215
├── delay
│   │   index
└── full
    │   index

邮件内容样本示例：

负责人您好我是深圳金海实业有限公司...
GG非常好的朋友H在计划马上的西藏自助游...

环境配置

本项目基于Paddle 2.0编写，如果你的环境不是本版本，请先参考官网安装Paddle 2.0。以下是环境配置代码：

# 导入相关的模块
import re
import jieba
import os
import random
import paddle
import paddlenlp as ppnlp
from paddlenlp.data import Stack, Pad, Tuple
import paddle.nn.functional as F
import paddle.nn as nn
from visualdl import LogWriter
import numpy as np
from functools import partial

数据加载与预处理

项目中使用了PaddleNLP的BertTokenizer进行数据处理，该tokenizer可以将原始输入文本转化成模型可接受的输入数据格式。以下是数据加载与预处理的代码：

# 解压数据集
!tar xvf data/data89631/trec06c.tgz

# 去掉非中文字符
def clean_str(string):
    string = re.sub(r"[^\u4e00-\u9fff]", " ", string)
    string = re.sub(r"\s{2,}", " ", string)
    return string.strip()

# 从指定路径读取邮件文件内容信息
def get_data_in_a_file(original_path, save_path='all_email.txt'):
    email = ''
    f = open(original_path, 'r', encoding='gb2312', errors='ignore')
    for line in f:
        line = line.strip().strip('\n')
        line = clean_str(line)
        email += line
    f.close()
    return email[-200:]

# 读取标签文件信息
f = open('trec06c/full/index', 'r')
for line in f:
    str_list = line.split(" ")
    if str_list[0] == 'spam':
        label = '0'
    elif str_list[0] == 'ham':
        label = '1'
    text = get_data_in_a_file('trec06c/full/' + str(str_list[1].split("\n")[0]))
    with open("all_email.txt", "a+") as f:
        f.write(text + '\t' + label + '\n')

自定义数据集

在项目中，我们需要自定义数据集，并使其数据格式与使用ppnlp.datasets.ChnSentiCorp.get_datasets加载后完全一致。以下是自定义数据集的代码：

class SelfDefinedDataset(paddle.io.Dataset):
    def __init__(self, data):
        super(SelfDefinedDataset, self).__init__()
        self.data = data

    def __getitem__(self, idx):
        return self.data[idx]

    def __len__(self):
        return len(self.data)

    def get_labels(self):
        return ["0", "1"]

def txt_to_list(file_name):
    res_list = []
    for line in open(file_name):
        res_list.append(line.strip().split('\t'))
    return res_list

trainlst = txt_to_list('train_list.txt')
devlst = txt_to_list('eval_list.txt')
testlst = txt_to_list('test_list.txt')

train_ds, dev_ds, test_ds = SelfDefinedDataset.get_datasets([trainlst, devlst, testlst])

模型训练

加载BERT预训练模型

项目中使用了PaddleNLP提供的BertForSequenceClassification模型进行文本分类的Fine-tune。由于垃圾邮件识别是二分类问题，所以设置num_classes为2。

以下是加载BERT预训练模型的代码：

# 加载预训练模型
model = ppnlp.transformers.BertForSequenceClassification.from_pretrained("bert-base-chinese", num_classes=2)

开始训练

为了监控训练过程，引入了VisualDL记录训练log信息。以下是开始训练的代码：

# 设置训练超参数
learning_rate = 1e-5
epochs = 10
warmup_proption = 0.1
weight_decay = 0.01

num_training_steps = len(train_loader) * epochs
num_warmup_steps = int(warmup_proption * num_training_steps)

def get_lr_factor(current_step):
    if current_step < num_warmup_steps:
        return float(current_step) / float(max(1, num_warmup_steps))
    else:
        return max(0.0,
                    float(num_training_steps - current_step) /
                    float(max(1, num_training_steps - num_warmup_steps)))

# 学习率调度器
lr_scheduler = paddle.optimizer.lr.LambdaDecay(learning_rate, lr_lambda=lambda current_step: get_lr_factor(current_step))

# 优化器
optimizer = paddle.optimizer.AdamW(
    learning_rate=lr_scheduler,
    parameters=model.parameters(),
    weight_decay=weight_decay,
    apply_decay_param_fun=lambda x: x in [
        p.name for n, p in model.named_parameters()
        if not any(nd in n for nd in ["bias", "norm"])
    ])

# 损失函数
criterion = paddle.nn.loss.CrossEntropyLoss()

# 评估函数
metric = paddle.metric.Accuracy()

# 训练过程
global_step = 0
with LogWriter(logdir="./log") as writer:
    for epoch in range(1, epochs + 1):    
        for step, batch in enumerate(train_loader, start=1):
            input_ids, segment_ids, labels = batch
            logits = model(input_ids, segment_ids)
            loss = criterion(logits, labels)
            probs = F.softmax(logits, axis=1)
            correct = metric.compute(probs, labels)
            metric.update(correct)
            acc = metric.accumulate()

            global_step += 1
            if global_step % 50 == 0:
                print("global step %d, epoch: %d, batch: %d, loss: %.5f, acc: %.5f" % (global_step, epoch, step, loss, acc))
                writer.add_scalar(tag="train/loss", step=global_step, value=loss)
                writer.add_scalar(tag="train/acc", step=global_step, value=acc)
            
            loss.backward()
            optimizer.step()
            lr_scheduler.step()
            optimizer.clear_gradients()
        
        eval_loss, eval_acc = evaluate(model, criterion, metric, dev_loader)
        writer.add_scalar(tag="eval/loss", step=epoch, value=eval_loss)
        writer.add_scalar(tag="eval/acc", step=epoch, value=eval_acc)

可以看到，在第2个epoch后验证集准确率已经达到99.4%以上，在第3个epoch就能达到99.6%以上。

预测效果

完成模型训练后，我们可以使用训练好的模型对测试集进行预测。以下是预测效果的代码：

data = ['您好我公司有多余的发票可以向外代开,国税,地税,运输,广告,海关缴款书如果贵公司,厂,有需要请来电洽谈,咨询联系电话,罗先生谢谢顺祝商祺']
label_map = {0: '垃圾邮件', 1: '正常邮件'}

predictions = predict(model, data, tokenizer, label_map, batch_size=32)
for idx, text in enumerate(data):
    print('预测内容: {} \n邮件标签: {}'.format(text, predictions[idx]))

预测效果良好，一个验证集准确率高达99.6%以上、基于BERT的中文邮件内容分类顺利完成！

以上是本文的全部内容，希望对读者理解如何使用BERT进行中文邮件内容分类有所帮助。欢迎交流指导！