[Python] Why is Keras pretrained BERT MaskedLM producing inconsistent predictions?

I am trying to use keras-nlp with a pretrained masked BERT model to predict some tokens in a sequence. However the model produces inconsistent results. What could be wrong or am i misunderstanding something?

import keras
import keras_nlp
import numpy as np
import tensorflow as tf

preprocessor = keras_nlp.models.BertMaskedLMPreprocessor.from_preset(
"bert_base_en"
)
masked_lm = keras_nlp.models.BertMaskedLM.from_preset(
"bert_base_en",
load_weights=True,
preprocessor=None,
)

test_sequence = preprocessor(["The capital of France is Paris."])
outputs = masked_lm(test_sequence[0])
predicted_ids = tf.math.argmax(outputs, axis=-1)


In this example the preprocessor produces the following sequence [101, 1109, 2364, 1104, 1699, 103, 103, 119, 102,0,0,...] with mask positions [5,6,0,0,0,0....]. The predicted_ids are [28059, 18994, 21690, 21690,.....] which is [b'##saur', b'##hetic', b'##lani', b'##lani',....] instead of something remotely close to 'is', 'Paris'.

Continue reading...