TensorFlow中如何为BERT分词器指定自定义输入序列长度

首先确保你已安装和TensorFlow版本匹配的tensorflow-text依赖，该依赖是BERT预处理器正常运行的必要条件。

你可以参考如下代码修改现有逻辑，直接自定义你需要的序列长度：

import tensorflow as tf
import tensorflow_hub as hub
import tensorflow_text

# 自定义需要的序列长度，最大不超过BERT支持的512
SEQ_LENGTH = 256

# 加载预处理器并拆分步骤，自定义序列长度
preprocessor = hub.load("你原有使用的bert预处理器hub地址")
tokenize_layer = hub.KerasLayer(preprocessor.tokenize)
pack_layer = hub.KerasLayer(
    preprocessor.bert_pack_inputs,
    arguments={"seq_length": SEQ_LENGTH}
)

# 原有Keras函数式API逻辑适配
text_input = tf.keras.layers.Input(shape=(), dtype=tf.string)
tokenized_text = tokenize_layer(text_input)
encoder_inputs = pack_layer([tokenized_text])

# 后续编码器逻辑和你原有代码完全一致
encoder = hub.KerasLayer(
    "你原有使用的bert编码器hub地址",
    trainable=True
)
outputs = encoder(encoder_inputs)
pooled_output = outputs["pooled_output"]
sequence_output = outputs["sequence_output"]

# 构建模型并测试
embedding_model = tf.keras.Model(text_input, pooled_output)
sentences = tf.constant(["(你的输入文本)"])
print(embedding_model(sentences))

注意事项 #

• 序列长度最大可设置为512，超过该值BERT基础模型无法正常处理
• 如果你的任务需要处理句子对输入，只需将两段文本分别分词后传入pack_layer的列表即可：encoder_inputs = pack_layer([tokenized_text1, tokenized_text2])
• 该写法完全兼容Keras函数式API的所有特性，你可以直接在输出层后拼接Dropout、全连接层等完成情感分类任务的构建

内容的提问来源于stack exchange，提问作者Jane Sully