python bert预训练模型加载，基于tensorflow-hub

前言

我的任务需要发现超短文本的语义（10个字左右的文本），需要对文本进行向量化处理，传统的词频或者tf-idf其实都是以词语的出现频率进行计算的，对于长文本来说还好，毕竟文本越长所包含的信息就越多，但对于短文本来说，传统的方法简直是灾难性的。所以我需要用深度学习的方法来实现这一任务。
但深度学习模型首先是需要大量的数据，其次需要硬件支持。我手上的文本虽然多，但大多数都是未标注的超短文本，而且我的硬件也比较一般，毕竟像bert这种300多M个参数的模型，也不是一台正常的电脑所能跑出来的。
所以我的目光投向了bert的预训练模型，我将在这里记录下bert预训练的加载过程，原理啥的我就不多说了，毕竟网上一抓一大把。

本次部署基于
anaconda
python 3.6
tensorflow 1.15.0 pip install tensorflow==1.15.0
tensorflow_hub 0.12.0 pip install tensorflow_hub
bert-tensorflow 1.0.1 pip install bert-tensorflow==1.0.1

python虚拟环境

首先我电脑上的python版本是3.8，并不支持tf1，因此需要一个py3.6虚拟环境，命令如下：

conda create -n py36 python=3.6

注意 -n 后面是你这个虚拟环境的名字，可以随便取，记住就行，之后在命令行下激活这个虚拟环境就行：

activate py36

如果需要在jupyter-notebook中使用该虚拟环境，需要先回到主环境，安装以下两个包：

conda install nb_conda
 
conda install ipykernel

之后再激活虚拟环境，在虚拟环境下安装：

conda install ipykernel

之后在notebook上就能看到我们的虚拟环境了：

前期准备

这一部分主要是安装各种包，我用的版本在前面已经说过了，这里直接放代码：

pip install tensorflow==1.15.0
pip install tensorflow_hub
pip install bert-tensorflow==1.0.1

然后就是我们的预训练模型，这里放一些中文的预训练模型：
Google原版bert
brightmart版roberta
哈工大版roberta
Google原版albert
brightmart版albert
转换后的albert
华为的NEZHA
自研语言模型
T5模型
GPT2_ML
Google原版ELECTRA
哈工大版ELECTRA
CLUE版ELECTRA
我下的是google原版

模型加载

#导入相关包
import tensorflow as tf
import tensorflow_hub as hub
from bert.tokenization import FullTokenizer #pip install bert-tensorflow==1.0.1
from tensorflow.keras.models import Model

#你存放预训练模型的路径
BERT_PATH  = r'bert_zh_L-12_H-768_A-12_1'

# 创建模型
#最大长度，我这里是32，意味着我每个短文本标记化后的的长度最多为32
#这里我用的是FullTokenizer，即文本最多30个字（因为还要包含两个占位符）
max_seq_length = 32 
input_word_ids = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32,
                                       name="input_word_ids")
input_mask = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32,
                                   name="input_mask")
segment_ids = tf.keras.layers.Input(shape=(max_seq_length,), dtype=tf.int32,
                                    name="segment_ids")
bert_layer = hub.KerasLayer(BERT_PATH,
                            trainable=True)

pooled_output, sequence_output = bert_layer([input_word_ids, input_mask, segment_ids])
 
model = Model(inputs=[input_word_ids, input_mask, segment_ids], outputs=[pooled_output, sequence_output])

tokenizer = FullTokenizer(r'bert_zh_L-12_H-768_A-12_1\assets\vocab.txt')

# 获取bert的输入
def get_masks(tokens, max_seq_length):
    """Mask for padding"""
    if len(tokens)>max_seq_length:
        raise IndexError("Token length more than max seq length!")
    return [1]*len(tokens) + [0] * (max_seq_length - len(tokens))
 
def get_segments(tokens, max_seq_length):
    """Segments: 0 for the first sequence, 1 for the second"""
    if len(tokens)>max_seq_length:
        raise IndexError("Token length more than max seq length!")
    segments = []
    current_segment_id = 0
    for token in tokens:
        segments.append(current_segment_id)
        if token == "[SEP]":
            current_segment_id = 1
    return segments + [0] * (max_seq_length - len(tokens))
 
def get_ids(tokens, tokenizer, max_seq_length):
    """Token ids from Tokenizer vocab"""
    token_ids = tokenizer.convert_tokens_to_ids(tokens)
    input_ids = token_ids + [0] * (max_seq_length-len(token_ids))
    return input_ids

def get_embed(s, tokenizer):
    stokens = tokenizer.tokenize(s)
    stokens = ["[CLS]"] + stokens + ["[SEP]"]
    input_ids = get_ids(stokens, tokenizer, max_seq_length)
    input_masks = get_masks(stokens, max_seq_length)
    input_segments = get_segments(stokens, max_seq_length)
 
    pool_embs, all_embs = model.predict([[input_ids],[input_masks],[input_segments]])
 
    # 查看结果
    return pool_embs, all_embs

调用这个get_embed函数即可获得短文本的向量化表示。

pool_embs, _ = get_embd("你好", tokenizer)
print(pool_embs)

到这里，整个流程就走完了，可以开始下一步的工作了。