In the following code, we expect a connection from input_token and masked_token to the TFDistilBertModel.
from transformers import TFAutoModel
import tensorflow as tf
def build_classifier_model():
bert = TFAutoModel.from_pretrained('distilbert-base-uncased')
input_ids_in = tf.keras.layers.Input(shape=(128,), name='input_token', dtype='int32')
input_masks_in = tf.keras.layers.Input(shape=(128,), name='masked_token', dtype='int32')
embedding_layer = bert(input_ids_in, attention_mask=input_masks_in)[0]
model = tf.keras.models.Model(inputs=[input_ids_in, input_masks_in], outputs=embedding_layer)
for layer in model.layers[:3]:
layer.trainable = False
return model
model = build_classifier_model()
tf.keras.utils.plot_model(model)
However, I get the following plot. Is there something wrong in my code or it is an error with keras?
