Как в tf.keras добавить слой с softmax выпрямителем

Проектирую нейросеть в Keras. Не до конца еще разобрался с тем, как тут добавлять слои.

Вот у меня такой код

bert_model = TFBertForSequenceClassification.from_pretrained('bert-base-multilingual-cased',num_labels=num_classes)

import tensorflow as tf
import keras

log_dir='/content/drive/My Drive/BSG/sms_spam/tb_bert'
model_save_path='/content/drive/My Drive/BSG/sms_spam/bert_model.h5'

callbacks = [tf.keras.callbacks.ModelCheckpoint(filepath=model_save_path,
                                                save_weights_only=True,
                                                monitor='val_loss',
                                                mode='min',
                                                save_best_only=True),
             keras.callbacks.TensorBoard(log_dir=log_dir)]

print('\nBert Model',bert_model.summary())

loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
metric = tf.keras.metrics.SparseCategoricalAccuracy('accuracy')
optimizer = tf.keras.optimizers.Adam(learning_rate=2e-5,epsilon=1e-08)

bert_model.compile(loss=loss,optimizer=optimizer,metrics=[metric])

history=bert_model.fit([train_inp,train_mask],
                       train_label,
                       batch_size=128,
                       epochs=epochs,
                       validation_data=([val_inp,val_mask],val_label),
                       callbacks=callbacks)

В итоге я получаю на выходе np.array и через argmax(axis=1) получаю предсказание класса. А как мне добавить в нейросеть слой (или куда-то еще), чтобы мне выводился не np.array, а предсказание вероятности через softmax? В данном случае моя задача подразумевает, что мне нужно получить в большей степени вероятность, чем класс, поэтому вот думаю как решить.

Спасибо.