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

发表于 更新于 分类于 ， 阅读次数： Disqus：

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

数据处理

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

  • 下载数据链接：link
  • 下载 BERT 模型：link
  • 安装 Kashgari: pip install "kashgari>=1.1,<2.0"
  • 安装 TensorFlow：pip install "tensorflow>=1.14,<2.0"

如果想使用 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
19
20
21
22
23
24
25
import pandas as pd
from kashgari.embeddings import BERTEmbedding

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

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

tokenizer = embed.tokenizer

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
import kashgari
from kashgari.tasks.classification import BiLSTM_Model

# 使用 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)