基于Spring Boot的ALBERT词向量服务(2)

ALBERT模型转换与准备

模型载入与简单推理

最基本的，我们需要有一个ALBERT的中文预训练模型。这里选择是Google推出的Tiny版本ALBERT，下载地址为：https://storage.googleapis.com/albert_zh/albert_tiny_zh_google.zip，下载完成后解压，会得到TensorFlow的checkpoint：

众所周知，TensorFlow的checkpoint保存了模型的各种权重和其他参数，可直接载入。而无论是BERT还是ALBERT，都有Google官方的TensorFlow开源代码，理论上只要参考其开源代码就可以实现checkpoint载入与推理。但是，官方的TensorFlow代码非常冗长，晦涩难懂，加上各种Estimator的使用，大大增加了理解成本，不适合快速模型搭建与使用。不过借助某个大神创建的bert4keras工具(https://github.com/bojone/bert4keras)，可以方便的使用Keras载入各种BERT模型的预训练权重，实现BERT模型的应用。而Keras代码就要简明清晰得多，而且可以做到和原生的TensorFlow无缝切换，是一个快速模型搭建和使用的好选择。在Python环境中可以方便地使用pip安装，这里我的版本是0.5.8（其他依赖包参照开源github项目中的requirements）：

pip install bert4keras==0.5.8

使用bert4keras工具载入上述ALBERT的checkpoint：

from bert4keras.models import build_transformer_model
from bert4keras.tokenizers import Tokenizer
import numpy as np
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'  # 屏蔽GPU

# config_path: albert模型结构描述文件路径
config_path = 'model/albert_tiny_zh_google/albert_config_tiny_g.json'
# checkpoint_path: TensorFlow checkpoint文件路径
checkpoint_path = 'model/albert_tiny_zh_google/albert_model.ckpt'
# dict_path: albert词汇表路径
dict_path = 'model/albert_tiny_zh_google/vocab.txt'

# 载入预训练ALBERT模型
model = build_transformer_model(config_path, checkpoint_path, model='albert',
                                with_pool=True)

一个build_transferomer_model函数就可以搞定，比官方的TensorFlow代码简单了不知多少倍。当然，要使用BERT系列模型，分词器Tokenizer是必不可少的，上面的代码中也导入了，现在需要创建一个Tokenizer对象以完成分词：

tokenizer = Tokenizer(dict_path, do_lower_case=True)
token_ids, segment_ids = tokenizer.encode(u'你好 世界')
print("Token ID: " + str(token_ids))
print("Segment ID:" + str(segment_ids))

上述输出结果是：

Token ID: [101, 872, 1962, 686, 4518, 102]
Segment ID:[0, 0, 0, 0, 0, 0]

稍微解释一下，除了首尾，Token ID中的每个数字表示ALBERT词汇表（第一张图的 vocab.txt）中各个输入汉字的位置 - 1，例如“你”在ALBERT词汇表中是第873个字符，则其Token ID是872；首尾的101和102分别表示[CLS]和[SEP]标记，是BERT系列模型中用来标记一段文本开头与结束的特殊符号，它们分别是第102和第103个字符。而Segment ID主要是用在输入两个不同的句子时起作用，比如判定两个文本的相似度时，第一个文本的Segment ID都是0，而第二个都是1，不过这里我们只输入单文本生成向量化表示，因此Segment ID都是0。而我们在“你好”和“世界”中输入的空格明显被忽略掉了。另外这里只考虑中文文本，使用的也是中文预训练模型，对于英文文本，其tokenizer更复杂一些，并非简单的根据空格切分单词。
接下来使用Token ID和Segment ID输入ALBERT模型，生成其向量表示：

print('\n ===== predicting =====\n')
vec1 = model.predict([np.array([token_ids]), np.array([segment_ids])])
print(vec1.shape)

输出的vec1是一个1行312列的numpy array，即为“你好 世界”的ALBERT嵌入式向量表示。

模型Freeze生成pb文件

我们已经成功地载入了ALBERT的预训练模型，也能在Python环境下非常方便地实现向量的推理。但是，我们的目的是部署到Java中能实现类似的功能，而TensorFlow提供了一个Java api接口(https://blog.csdn.net/mandagod/article/details/89446473)，可以方便地调用在Python中训练好的模型。这需要我们事先把模型freeze，生成TensorFlow的pb文件，借助以下的函数，可以将Keras模型对象输出为TensorFlow pb文件：

from bert4keras.models import build_transformer_model
import tensorflow as tf
from keras import backend as K
from tensorflow.python.framework import graph_util, graph_io
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '-1'


def export_graph(model, export_path, output_name):
    '''

    :param model: keras模型对象
    :param export_path: 导出路径
    :param output_name: 导出pb文件的名称
    :return: 模型的输入、输出Tensor名称——input_names, input_nodes

    '''
    input_names = model.input_names

    if not tf.gfile.Exists(export_path):
        tf.gfile.MakeDirs(export_path)

    with K.get_session() as sess:
        init_graph = sess.graph
        with init_graph.as_default():
            out_nodes = []

            for i in range(len(model.outputs)):
                out_nodes.append("output_" + str(i + 1))
                tf.identity(model.output[i], "output_" + str(i + 1))

            init_graph = sess.graph.as_graph_def()
            main_graph = graph_util.convert_variables_to_constants(sess, init_graph, out_nodes)
            graph_io.write_graph(main_graph, export_path, name='%s.pb' % output_name, as_text=False)

    return input_names, out_nodes

然后调用流程如下：

if __name__ == '__main__':
    config_path = 'model/albert_tiny_zh_google/albert_config_tiny_g.json'
    checkpoint_path = 'model/albert_tiny_zh_google/albert_model.ckpt'
    dict_path = 'model/albert_tiny_zh_google/vocab.txt'
    output_path = "output/"
    model = build_transformer_model(config_path, checkpoint_path, model='albert', with_pool=True)  # 建立模型，加载权重
    inputs, outputs = export_graph(model, output_path, "albert_tiny_zh_google")
    print(inputs)
    print(outputs)

以上会输出整个模型输入与输出的Tensor的名称，这个很重要，因为我们后面在Java中调用模型的时候，需要知道输入输出Tensor的名称才能正确地喂数据，取结果，输出如下：

input_names:['Input-Token', 'Input-Segment']
output_names:['output_1']

"Input-Token"对应Token ids输入，"Input-Segment"对应Segment ids输入，而"output_1"是输出的ALBERT向量表示。最后我们会得到一个名为"albert_tiny_zh_google.pb"的pb文件，这是我们在Java工程中推理需要用到的模型文件，大小仅有15.7 MB：

小结

至此我们所需要的模型文件已经准备完毕，另外还需要ALBERT对应的词汇表，这个是现成的vocab.txt，这两个文件是模型推理必不可少的，将会在下一篇中说明其在Java工程中的加载和使用。

Python支持工程开源代码：https://github.com/Aiwiscal/albert-vec-support
Java主工程开源代码：https://github.com/Aiwiscal/albert-vec

喜欢请给star哦~