之前 分享了如何原生的使用 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