NLP - BERT/ERNIE 文本分类和部署

本文主要是介绍使用 Kashgari 框架实现 BERT 文本分类及使用 Tensorflow Serving 方式部署模型。

数据处理

我们使用新浪微博情感分类数据集,数据集很简单,只有二分类。此外还需要下载 BERT 模型和对应的分词器。

  • 下载数据链接:link
  • 下载 BERT 模型:link
  • 下载 BERT 分词器:link
  • 安装 Kashgari: pip install kashgari-tf

如果想使用 ERNIE 模型,则需要下载 Tensorflow_ERNIE.

1
2
3
4
5
6
7
8
id      label   review
62050 0 太过分了@Rexzhenghao //@Janie_Zhang:招行最近负面新闻越来越多呀...
68263 0 希望你?得好?我本"?肥血?史"[晕][哈哈]@Pete三姑父
81472 0 有点想参加????[偷?]想安排下时间再决定[抓狂]//@黑晶晶crystal: @细腿大羽...
42021 1 [给力]感谢所有支持雯婕的芝麻![爱你]
7777 1 2013最后一天,在新加坡开心度过,向所有的朋友们问声:新年快乐!2014年,我们会更好[调...
100399 0 大中午出门办事找错路,曝晒中。要多杯具有多杯具。[泪][泪][汗]
82398 0 马航还会否认吗?到底在隐瞒啥呢?[抓狂]//@头条新闻: 转发微博
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
import pandas as pd
from tokenization import BasicTokenizer

tokenizer = BasicTokenizer()

df = pd.read_csv('weibo_senti_100k.csv')
# 进行分词处理
df['cutted'] = df['review'].apply(lambda x: tokenizer.tokenize(x))

# 准备训练测试数据集
train_x = list(df['cutted'][:int(len(df)*0.7)])
train_y = list(df['label'][:int(len(df)*0.7)])

valid_x = list(df['cutted'][int(len(df)*0.7):int(len(df)*0.85)])
valid_y = list(df['label'][int(len(df)*0.7):int(len(df)*0.85)])

test_x = list(df['cutted'][int(len(df)*0.85):])
test_y = list(df['label'][int(len(df)*0.85):])

训练验证模型

接下来我们初始化 Embedding 和模型,并且训练。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
import kashgari
from kashgari.embeddings import BERTEmbedding
from kashgari.tasks.classification import BiLSTM_Model

BERT_PATH = '<BERT/ERNIE 模型路径>'

# 初始化 Embedding
embed = BERTEmbedding(BERT_PATH,
task=kashgari.CLASSIFICATION,
sequence_length=64)

# 使用 embedding 初始化模型
model = BiLSTM_Model(embed)
# 先只训练一轮
model.fit(train_x, train_y, valid_x, valid_y, batch_size=1024, epochs=1)

Kashgari 内置了验证方法 evaluate,使用非常方便。

1
model.evaluate(test_x, test_y, batch_size=512)

评估结果如下,已经可以达到 92% 还不错。接下来保存模型并且部署。

1
2
3
4
5
6
7
8
              precision    recall  f1-score   support

0 0.9351 0.9187 0.9269 9069
1 0.9189 0.9353 0.9270 8930

accuracy 0.9269 17999
macro avg 0.9270 0.9270 0.9269 17999
weighted avg 0.9271 0.9269 0.9269 17999

保存模型

首先使用内置方法完整保存模型,以便后续使用。

1
model.save('bert_model')

然后把模型转换为 SavedModel 格式。

1
2
import kashgari
kashgari.utils.convert_to_saved_model(model, 'tf_bert_model', version=1)

Tensorflor Serving 部署

使用 docker 启动支持 GPU 的 Tensorflor Serving,如果没有接触过 Tensorflow Serving,可以查阅 TensorFlow Serving入门 一文。

1
2
3
docker run --runtime=nvidia -p 8501:8501 \
-v `pwd`/tf_bert_model:/models/tf_bert_model \
-e MODEL_NAME=tf_bert_model -t tensorflow/serving:1.14.0-gpu

如果没有配置 Docker GPU 环境,则可以使用 tensorflow/serving CPU 版本启动。

1
2
3
docker run -p 8501:8501 \
-v `pwd`/tf_bert_model:/models/tf_bert_model \
-e MODEL_NAME=tf_bert_model -t tensorflow/serving:1.14.0

此时我们已经部署好了 Tensorflow Serving 模型,但是模型的输出输出均为张量,我们还需要对输入进行预处理才能使用。Kashgari 中提供了工具方法简化处理过程。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
import requests
from kashgari import utils
import numpy as np

x = ['这', '个', '价', '不', '算', '高', ',', '和', '一', '天', '内', '训', '相', '比', '相', '差', '无', '几']

# 从保存模型加载词表
processor = utils.load_processor(model_path='tf_bert_model/1')
# 输入转换成张量
tensor = processor.process_x_dataset([x])

# 格式化为 BERT 格式
tensor = [{
"Input-Token:0": i.tolist(),
"Input-Segment:0": np.zeros(i.shape).tolist()
} for i in tensor]

# 进行预测
r = requests.post("http://localhost:8501/v1/models/tf_bert_model:predict",
json={"instances": tensor.tolist()})
preds = r.json()['predictions']
label_index = np.array(preds).argmax(-1)

# 把预测结果转换成具体的标签
labels = processor.reverse_numerize_label_sequences(label_index)