用BERT做中文邮件内容分类
- 项目背景与意义
- 项目思路
- 数据集介绍
- 环境配置
- 数据加载与预处理
- 自定义数据集
- 模型训练
- 加载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进行中文邮件内容分类有所帮助。欢迎交流指导！
