bert 使用（3）

之前 分享了如何原生的使用 bert 和 albert ，但 在预测的时候 只能预测单条，无法进行批量预测。bert模型 现在已经出现两年，各种使用bert 的方式越来越方便，今天介绍一下 如何使用 bert-for-tf2 调用 bert 和albert 等

相较于使用原生bert ，下面方法 代码简单调用 方便，和常规模型搭建 相同 ，便于 使用 keras 来搭建更复杂的模型 

环境要求： tensorflow 1.14 以上 或 tensorflow 2.0

pip install bert-for-tf2

完整代码 参考码云地址：https://gitee.com/heras/albert_keras

核心代码：

import tensorflow as tf
import bert
from bert.tokenization.albert_tokenization import FullTokenizer


def build_transformer(max_seq_length,bert_dir):
    lower_case = True
    # 调用 bert-for-tf2
    vocab = glob.glob(os.path.join(bert_dir, '*vocab*.txt'))[0]
    tokenizer = FullTokenizer(vocab_file=vocab, do_lower_case=lower_case)
    bert_params = bert.params_from_pretrained_ckpt(bert_dir)
    l_bert = bert.BertModelLayer.from_params(bert_params, name='albert')
    l_input_ids = tf.keras.layers.Input(shape=(max_seq_length,), dtype='int32', name="input_ids")
    l_mask_ids = tf.keras.layers.Input(shape=(max_seq_length,), dtype='int32', name="mask_ids")
    l_token_type_ids = tf.keras.layers.Input(shape=(max_seq_length,), dtype='int32', name="token_type_ids")
    output = l_bert([l_input_ids, l_token_type_ids], mask=l_mask_ids)
    model = tf.keras.Model(inputs=[l_input_ids, l_mask_ids, l_token_type_ids], outputs=output)
    model.build(input_shape=(None, max_seq_length))
    # 加载 bert 参数
    ckpt = glob.glob(os.path.join(bert_dir, '*.index'))
    assert ckpt, f'No checkpoint found under {bert_dir}'
    ckpt, _ = os.path.splitext(ckpt[0])
    bert.load_bert_weights(l_bert, ckpt)
    return model,tokenizer

特别说明：

核心代码 是copy Hanlp 开源项目：https://github.com/hankcs/HanLP.git