本文将会介绍如何使用keras-bert实现文本多分类任务，其中对BERT进行微调。
项目结构
其中依赖的Python第三方模块如下：
pandas==0.23.4
Keras==2.3.1
keras_bert==0.83.0
numpy==1.16.4
数据集
本文采用的多分类数据集为sougou小分类数据集和THUCNews数据集，简介如下：
- sougou小分类数据集
共有5个类别，分别为体育、健康、军事、教育、汽车。划分为训练集和测试集，其中训练集每个分类800条样本，测试集每个分类100条样本。
- THUCNews数据集
共有10个分类，类别为：体育, 财经, 房产, 家居, 教育, 科技, 时尚, 时政, 游戏, 娱乐。数据集划分为：训练集: 5000 * 10，测试集: 1000 * 10。
模型训练
模型训练脚本model_train.py的完整代码如下：
# -*- coding: utf-8 -*-
import json
import codecs
import pandas as pd
import numpy as np
from keras_bert import load_trained_model_from_checkpoint, Tokenizer
from keras.layers import *
from keras.models import Model
from keras.optimizers import Adam
# 建议长度<=510
maxlen = 300
BATCH_SIZE = 8
config_path = './chinese_L-12_H-768_A-12/bert_config.json'
checkpoint_path = './chinese_L-12_H-768_A-12/bert_model.ckpt'
dict_path = './chinese_L-12_H-768_A-12/vocab.txt'
token_dict = {
}
with codecs.open(dict_path, 'r', 'utf-8') as reader:
for line in reader:
token = line.strip()
token_dict[token] = len(token_dict)
class OurTokenizer(Tokenizer):
def _tokenize(self, text):
R = []
for c in text:
if c in self._token_dict:
R.append(c)
else:
R.append('[UNK]') # 剩余的字符是[UNK]
return R
tokenizer = OurTokenizer(token_dict)
def seq_padding(X, padding=0):
L = [len(x) for x in X]
ML = max(L)
return np.array([
np.concatenate([x, [padding] * (ML - len(x))]) if len(x) < ML else x for x in X
])
class DataGenerator:
def __init__(self, data, batch_size=BATCH_SIZE):
self.data = data
self.batch_size = batch_size
self.steps = len(self.data) // self.batch_size
if len(self.data) % self.batch_size != 0:
self.steps += 1
def __len__(self):
return self.steps
def __iter__(self):
while True:
idxs = list(range(len(self.data)))
np.random.shuffle(idxs)
X1, X2, Y = [], [], []
for i in idxs:
d = self.data[i]
text = d[0][:maxlen]
x1, x2 = tokenizer.encode(first=text)
y = d[1]
X1.append(x1)
X2.append(x2)
Y.append(y)
if len(X1) == self.batch_size or i == idxs[-1]:
X1 = seq_padding(X1)
X2 = seq_padding(X2)
Y = seq_padding(Y)
yield [X1, X2], Y
[X1, X2, Y] = [], [], []
# 构建模型
def create_cls_model(num_labels):
bert_model = load_trained_model_from_checkpoint(config_path, checkpoint_path, seq_len=None)
for layer in bert_model.layers:
layer.trainable = True
x1_in = Input(shape=(None,))
x2_in = Input(shape=(None,))
x = bert_model([x1_in, x2_in])
cls_layer = Lambda(lambda x: x[:, 0])(x) # 取出[CLS]对应的向量用来做分类
p = Dense(num_labels, activation='softmax')(cls_layer) # 多分类
model = Model([x1_in, x2_in], p)
model.compile(
loss='categorical_crossentropy',
optimizer=Adam(1e-5), # 用足够小的学习率
metrics=['accuracy']
)
# model.summary()
return model
if __name__ == '__main__':
# 数据处理, 读取训练集和测试集
print("begin data processing...")
train_df = pd.read_csv("data/cnews/cnews_train.csv").fillna(value="")
test_df = pd.read_csv("data/cnews/cnews_test.csv").fillna(value="")
labels = train_df["label"].unique()
with open("label.json", "w", encoding="utf-8") as f:
f.write(json.dumps(dict(zip(range(len(labels)), labels)), ensure_ascii=False, indent=2))
train_data = []
test_data = []
for i in range(train_df.shape[0]):
label, content = train_df.iloc[i, :]
label_id = [0] * len(labels)
for j, _ in enumerate(labels):
if _ == label:
label_id[j] = 1
train_data.append((content, label_id))
for i in range(test_df.shape[0]):
label, content = test_df.iloc[i, :]
label_id = [0] * len(labels)
for j, _ in enumerate(labels):
if _ == label:
label_id[j] = 1
test_data.append((content, label_id))
print("finish data processing!")
# 模型训练
model = create_cls_model(len(labels))
train_D = DataGenerator(train_data)
test_D = DataGenerator(test_data)
print("begin model training...")
model.fit_generator(
train_D.__iter__(),
steps_per_epoch=len(train_D),
epochs=3,
validation_data=test_D.__iter__(),
validation_steps=len(test_D)
)
print("finish model training!")
# 模型保存
model.save('cls_cnews.h5')
print("Model saved!")
result = model.evaluate_generator(test_D.__iter__(), steps=len(test_D))
print("模型评估结果:", result)
其中模型结构如下：
__________________________________________________________________________________________________
Layer (type) Output Shape Param # Connected to
==================================================================================================
input_1 (InputLayer) (None, None) 0
__________________________________________________________________________________________________
input_2 (InputLayer) (None, None) 0
__________________________________________________________________________________________________
model_2 (Model) (None, None, 768) 101677056 input_1[0][0]
input_2[0][0]
__________________________________________________________________________________________________
lambda_1 (Lambda) (None, 768) 0 model_2[1][0]
__________________________________________________________________________________________________
dense_1 (Dense) (None, 10) 7690 lambda_1[0][0]
==================================================================================================
Total params: 101,684,746
Trainable params: 101,684,746
Non-trainable params: 0
在上述模型中，我们取取出[CLS]对应的向量，后接全连接层，激活函数采用Softmax函数，就完成多分类模型的搭建了，非常简单方便。
模型评估
模型评估脚本model_evaluate.py的完整代码如下：
# -*- coding: utf-8 -*-
# 模型评估脚本
import json
import numpy as np
import pandas as pd
from keras.models import load_model
from keras_bert import get_custom_objects
from sklearn.metrics import classification_report
from model_train import token_dict, OurTokenizer
maxlen = 300
# 加载训练好的模型
model = load_model("cls_cnews.h5", custom_objects=get_custom_objects())
tokenizer = OurTokenizer(token_dict)
with open("label.json", "r", encoding="utf-8") as f:
label_dict = json.loads(f.read())
# 对单句话进行预测
def predict_single_text(text):
# 利用BERT进行tokenize
text = text[:maxlen]
x1, x2 = tokenizer.encode(first=text)
X1 = x1 + [0] * (maxlen - len(x1)) if len(x1) < maxlen else x1
X2 = x2 + [0] * (maxlen - len(x2)) if len(x2) < maxlen else x2
# 模型预测并输出预测结果
predicted = model.predict([[X1], [X2]])
y = np.argmax(predicted[0])
return label_dict[str(y)]
# 模型评估
def evaluate():
test_df = pd.read_csv("data/cnews/cnews_test.csv").fillna(value="")
true_y_list, pred_y_list = [], []
for i in range(test_df.shape[0]):
print("predict %d samples" % (i+1))
true_y, content = test_df.iloc[i, :]
pred_y = predict_single_text(content)
true_y_list.append(true_y)
pred_y_list.append(pred_y)
return classification_report(true_y_list, pred_y_list, digits=4)
output_data = evaluate()
print("model evaluate result:\n")
print(output_data)
运行上述代码，对两个数据集进行评估，结果如下：
- sougou数据集
模型参数: batch_size = 8, maxlen = 256, epoch=10
评估结果:
precision recall f1-score support
体育 0.9802 1.0000 0.9900 99
健康 0.9495 0.9495 0.9495 99
军事 1.0000 1.0000 1.0000 99
教育 0.9307 0.9495 0.9400 99
汽车 0.9895 0.9495 0.9691 99
accuracy 0.9697 495
macro avg 0.9700 0.9697 0.9697 495
weighted avg 0.9700 0.9697 0.9697 495
- THUCNews数据集
模型参数: batch_size = 8, maxlen = 300, epoch=3
评估结果:
precision recall f1-score support
体育 0.9970 0.9990 0.9980 1000
娱乐 0.9890 0.9890 0.9890 1000
家居 0.9949 0.7820 0.8757 1000
房产 0.8006 0.8710 0.8343 1000
教育 0.9753 0.9480 0.9615 1000
时尚 0.9708 0.9980 0.9842 1000
时政 0.9318 0.9560 0.9437 1000
游戏 0.9851 0.9950 0.9900 1000
科技 0.9689 0.9970 0.9828 1000
财经 0.9377 0.9930 0.9645 1000
accuracy 0.9528 10000
macro avg 0.9551 0.9528 0.9524 10000
weighted avg 0.9551 0.9528 0.9524 10000
模型预测
模型预测脚本model_predict.py的完整代码如下：
# -*- coding: utf-8 -*-
# @Time : 2020/12/23 15:28
# @Author : Jclian91
# @File : model_predict.py
# @Place : Yangpu, Shanghai
# 模型预测脚本
import time
import json
import numpy as np
from model_train import token_dict, OurTokenizer
from keras.models import load_model
from keras_bert import get_custom_objects
maxlen = 256
# 加载训练好的模型
model = load_model("cls_sougou.h5", custom_objects=get_custom_objects())
tokenizer = OurTokenizer(token_dict)
with open("label.json", "r", encoding="utf-8") as f:
label_dict = json.loads(f.read())
s_time = time.time()
# 预测示例语句
text = "说到硬派越野SUV，你会想起哪些车型？是被称为“霸道”的丰田 普拉多 (配置 | 询价) ，还是被叫做“山猫”的帕杰罗，亦或者是“渣男专车”奔驰大G、" \
"“沙漠王子”途乐。总之，随着世界各国越来越重视对环境的保护，那些大排量的越野SUV在不久的将来也会渐渐消失在我们的视线之中，所以与其错过，" \
"不如趁着还年轻，在有生之年里赶紧去入手一台能让你心仪的硬派越野SUV。而今天我想要来跟你们聊的，正是全球公认的十大硬派越野SUV，" \
"越野迷们看完之后也不妨思考一下，到底哪款才是你的菜，下面话不多说，赶紧开始吧。"
# 利用BERT进行tokenize
text = text[:maxlen]
x1, x2 = tokenizer.encode(first=text)
X1 = x1 + [0] * (maxlen-len(x1)) if len(x1) < maxlen else x1
X2 = x2 + [0] * (maxlen-len(x2)) if len(x2) < maxlen else x2
# 模型预测并输出预测结果
predicted = model.predict([[X1], [X2]])
y = np.argmax(predicted[0])
print("原文: %s" % text)
print("预测标签: %s" % label_dict[str(y)])
e_time = time.time()
print("cost time:", e_time-s_time)
我们在新的样本上进行模型预测。
- sougou数据集
原文: 说到硬派越野SUV，你会想起哪些车型？是被称为“霸道”的丰田 普拉多 (配置 | 询价) ，还是被叫做“山猫”的帕杰罗，亦或者是“渣男专车”奔驰大G、“沙漠王子”途乐。总之，随着世界各国越来越重视对环境的保护，那些大排量的越野SUV在不久的将来也会渐渐消失在我们的视线之中，所以与其错过，不如趁着还年轻，在有生之年里赶紧去入手一台能让你心仪的硬派越野SUV。而今天我想要来跟你们聊的，正是全球公认的十大硬派越野SUV，越野迷们看完之后也不妨思考一下，到底哪款才是你的菜，下面话不多说，赶紧开始吧。
预测标签: 汽车
原文: 【#美30架战机在阿拉斯加海岸大象漫步#】据美国艾尔森空军基地网站消息称，近日美国空军30架战斗机和2架空中加油机自艾尔森空军基地起飞，在阿拉斯加海岸完成了”大象漫步“式的演习。
预测标签: 军事
原文: “十三五”期间，我国义务教育三科统编教材实现所有年级全覆盖；普通高中三科统编教材已覆盖20个省份，预计2022年前实现所有省份全覆盖，2025年实现所有年级全覆盖。昨日，在教育部新闻发布会上，教育部教材局局长田慧生透露，义务教育课程方案和各学科课程标准修订明年完成。
预测标签: 教育
- THUCNews数据集
原文: 北京时间12月26日,2020-21赛季NBA圣诞大战如约上演。在一场焦点对决中，洛杉矶湖人在主场与达拉斯独行侠遭遇。全场打完，湖人138-115轻取独行侠，拿到赛季首胜，同时也送给对手2连败。
预测标签: 体育
原文: 近两年来，手机屏幕就开始不断升级，高刷新率也成为一种趋势，就算性能做的再好，手机屏幕不能流畅真实的展现出来，也会很大程度上影响使用感受，所以屏幕就是手机硬件的窗户，可以预见未来的高端手机在冲击性能的同时，必然会对提高对屏幕的要求。
预测标签: 科技
原文: 松江佘山板块已太久没有豪宅入市，令区域内不少高端置业客寂寞难耐。不过好在，不久前火爆登场的国贸佘山原墅一下子满足了这类客户的需求。 无论是社区规划、产品户型、装修配置，还是设计手法，相较于周边千万级别墅，国贸佘山原墅也是不逞多让，更令人惊喜的是别墅的品质却仅需公寓的价格！
预测标签: 房产
总结
本项目已经开源，Github地址为： https://github.com/percent4/keras_bert_text_classification 。
后续将会介绍如何使用keras-bert实现文本多标签分类任务。
2020年12月26日于上海浦东