如何将HuggingFace Tokenizer转换为TensorFlow Keras Layer实现推理？

我正在尝试理解如何对以TensorFlow Keras模型形式加载的预训练HuggingFace模型进行推理。

背景 #

我正在微调一个预训练的DistilBert分类器，预处理数据和加载/训练模型的代码如下：

from transformers import TFAutoModelForSequenceClassification

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

model = TFAutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased", num_labels=2, id2label=id2label, label2id=label2id
)

# 添加额外层

tf_train = model.prepare_tf_dataset(question_train_test_split['train'], batch_size=16, shuffle=True, tokenizer=tokenizer)

model.compile(optimizer=tf.keras.optimizers.Adam(2e-5))

# 冻结模型的第一个Transformer层

model.layers[0].trainable=False

print('Model Architecture:')
print(model.summary())

model.fit(tf_train, epochs=3)

其中question_train_test_split是HuggingFace Dataset对象实例。

这段代码运行完全符合预期，能将HuggingFace模型加载为tf.keras层，并且可以通过.fit方法正常训练。

但我在进行预测时遇到了问题。我知道需要对字符串输入进行分词，但希望将tokenizer作为tf.keras层加载。多方查找方法却一无所获。

理想情况下，我希望实现如下效果：

user_input = 'When were the Beatles formed?'

model_input = tokenizer(user_input) # 此处HF Tokenizer应为tf.keras层

model = model(model_input)

这样我就可以将整个模型（包含tokenizer、Transformer层以及分类器层）保存为TensorFlow SavedModel。希望得到将HuggingFace tokenizer转换为TensorFlow Keras层的相关指引。

1. 封装HuggingFace分词器为Keras层 #

通过自定义tf.keras.layers.Layer，把HuggingFace分词器包装成可嵌入Keras模型的层，直接处理字符串输入并输出模型所需的张量：

import tensorflow as tf
from transformers import AutoTokenizer

class HuggingFaceTokenizerLayer(tf.keras.layers.Layer):
    def __init__(self, tokenizer_name, max_length=128, **kwargs):
        super().__init__(**kwargs)
        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
        self.max_length = max_length

    def call(self, inputs):
        # 兼容单个或批量字符串输入
        tokenized = self.tokenizer(
            inputs.numpy().tolist(),
            padding='max_length',
            truncation=True,
            max_length=self.max_length,
            return_tensors='tf'
        )
        return {
            'input_ids': tokenized['input_ids'],
            'attention_mask': tokenized['attention_mask']
        }

    # 必须实现该方法，保证模型可保存和加载
    def get_config(self):
        config = super().get_config()
        config.update({
            'tokenizer_name': self.tokenizer.name_or_path,
            'max_length': self.max_length
        })
        return config

2. 构建端到端完整模型 #

将分词层与训练好的DistilBert模型拼接，形成一个直接接受原始字符串输入的Keras模型：

# 若已保存微调后的模型，用以下方式加载；否则直接使用训练阶段的model对象
# from transformers import TFAutoModelForSequenceClassification
# model = TFAutoModelForSequenceClassification.from_pretrained('./fine_tuned_distilbert')

# 创建分词层实例
tokenizer_layer = HuggingFaceTokenizerLayer("distilbert-base-uncased")

# 定义字符串输入层
input_layer = tf.keras.layers.Input(shape=(), dtype=tf.string)

# 分词层处理输入
tokenized_output = tokenizer_layer(input_layer)

# 传入预训练模型得到分类结果
model_output = model(tokenized_output)

# 组装完整模型
end_to_end_model = tf.keras.Model(inputs=input_layer, outputs=model_output.logits)

3. 推理与模型保存 #

现在可以直接用字符串输入推理，并且能将整个模型保存为TensorFlow SavedModel格式：

# 单样本推理示例
user_input = 'When were the Beatles formed?'
predictions = end_to_end_model.predict([user_input])

# 保存完整模型
end_to_end_model.save('./end_to_end_distilbert_classifier')

# 加载保存的模型并推理
loaded_model = tf.keras.models.load_model(
    './end_to_end_distilbert_classifier',
    custom_objects={'HuggingFaceTokenizerLayer': HuggingFaceTokenizerLayer}
)
loaded_predictions = loaded_model.predict([user_input])

关键注意事项 #

• 自定义层必须实现get_config方法，否则模型无法正常序列化和加载。
• 如果模型需要token_type_ids（如BERT-base），可在call方法中添加对应的返回键值对。
• 若微调时添加了额外自定义层，需确保在构建端到端模型时正确衔接输入输出。

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