How to get the hidden states of a fine-tuned TFBertModel?

Question

I first fine-tune the Bert model on a text classification task and but then I want to get the embeddings of the fine-tuned model in TensorFlow. Unfortunately, I can only say output_hidden_states=True, in the first line where I download the pre-trained Bert model but not in the second stage where I create a tf.Keras.Model. Here is the code for how I make and train the model:

max_len = 55

from transformers import BertConfig, BertTokenizer, TFBertModel

def build_custome_model():
    bert_encoder = TFBertModel.from_pretrained(Base_BERT_Path)

    input_word_ids = tf.keras.Input(shape=(max_len,), dtype=tf.int32, name="input_word_ids")
    input_mask = tf.keras.Input(shape=(max_len,), dtype=tf.int32, name="input_mask")
    input_type_ids = tf.keras.Input(shape=(max_len,), dtype=tf.int32, name="input_type_ids")
    
    embedding = bert_encoder([input_word_ids, input_mask, input_type_ids])[0]
    clf_output = embedding[:,0,:]

    net = tf.keras.layers.Dropout(0.4)(clf_output)
    output = tf.keras.layers.Dense(5, activation='softmax')(net)
    
    model = tf.keras.Model(inputs=[input_word_ids, input_mask, input_type_ids], outputs=output)

    model.compile(tf.keras.optimizers.Adam(lr=1e-5), loss='sparse_categorical_crossentropy', metrics=['accuracy'])
    
    return model

Then I train the model on a dataset with 2 sentences and as score for their similarity

#------Training with stratifiedkfold-------


k = 5
kfold = StratifiedKFold(n_splits = k, shuffle = True)

for i, (train_idx, val_idx) in enumerate(kfold.split(first_sentences, labels.score), 1):

    epoch_evaluation = {}

    train_input = create_input(np.array(first_sentences)[train_idx], np.array(second_sentences)[train_idx], tokenizer, max_len=max_seq_length)
    validation_input = create_input(np.array(first_sentences)[val_idx], np.array(second_sentences)[val_idx], tokenizer, max_len=max_seq_length)

    history = model.fit(x = train_input, y = labels.loc[train_idx, 'score'],
                        validation_data= (validation_input, labels.loc[val_idx, 'score']),
                        epochs = 5,
                        verbose = 1,
                        batch_size = 8)

My goal is to have a model at the end that is trained on this dataset and can output the embeddings (first layer of the hidden states (output[2][0])) whenever I give it a sentence so that I can get the mean of all the fine-tuned token embeddings of a sentence.

Answer 1

You can retrieve the embeddings with the get_input_embeddings function:

model = build_custome_model():

model.layers[3].get_input_embeddings()(input_ids)