NLP - BERT_ERNIE 文本分类和部署

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

数据处理

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

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

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

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

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 和模型，并且训练。

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，使用非常方便。

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

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

              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

保存模型

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

model.save('bert_model')

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

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

Tensorflor Serving 部署

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

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 版本启动。

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 中提供了工具方法简化处理过程。

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)