如何使用Keras创建BERT层？

pip install keras
pip install tensorflow
pip install transformers

from transformers import TFBertModel

bert_model = TFBertModel.from_pretrained("bert-base-uncased")

from tensorflow import keras
import tensorflow as tf

class BERTLayer(keras.layers.Layer):
    def __init__(self, bert_model, **kwargs):
        super(BERTLayer, self).__init__(**kwargs)
        self.bert = bert_model
        
    def call(self, inputs):
        input_ids, attention_mask = inputs
        outputs = self.bert(input_ids, attention_mask=attention_mask)[0]
        return outputs

input_ids = keras.Input(shape=(max_seq_length,), dtype=tf.int32)
attention_mask = keras.Input(shape=(max_seq_length,), dtype=tf.int32)

bert_output = BERTLayer(bert_model)([input_ids, attention_mask])
# 在此处添加其他层以完成自定义模型的构建

model = keras.Model(inputs=[input_ids, attention_mask], outputs=bert_output)