mrr_metric = tfr.keras.metrics.MRRMetric(topn=10)
def create_classifier_bert_model():    
    inputs = layers.Input((config.MAX_LEN,), dtype=tf.int64)
    sequence_output = pretrained_bert_model(inputs)
    outputs = layers.Dense(vectorize_layer.vocabulary_size(), activation="softmax")(sequence_output)
    # outputs2 = layers.Dense(vectorize_layer.vocabulary_size(), activation="softmax")(sequence_output)
    classifer_model = keras.Model(inputs, outputs=[outputs], name="prediction")
    optimizer = keras.optimizers.Adam(learning_rate=config.LR)
    classifer_model.compile(
        optimizer=optimizer, loss="sparse_categorical_crossentropy", weighted_metrics=[tf.keras.metrics.SparseCategoricalAccuracy(), mrr_metric]
    )
    return classifer_model

classifer_model = create_classifier_bert_model()
print(config.LR)
classifer_model.fit(
    mlm_ds_ft_cs,
    epochs=5,
)