from transformers import BertTokenizer, BertForMaskedLM
import torch

mask = "[MASK]",
checkpoint = "bert-base-uncased"

tokenizer = BertTokenizer.from_pretrained(f"{checkpoint}")
model = BertForMaskedLM.from_pretrained(f"{checkpoint}")

inputs = tokenizer(f"The capital of France is {mask}.", return_tensors="pt")

with torch.no_grad():
    logits = model(**inputs).logits

# retrieve index of {mask}
mask_token_index = (inputs.input_ids == tokenizer.mask_token_id)[0].nonzero(as_tuple=True)[0]
predicted_token_id = logits[0, mask_token_index].argmax(axis=-1)
expected_output = tokenizer.decode(predicted_token_id)


print(mask_token_index)  # tensor([8]): row dimension from `nonzero()`
print(predicted_token_id)  # tensor([3000])
print(expected_output)  # paris

from transformers import BertTokenizer, TFBertForMaskedLM
import tensorflow as tf

tokenizer = BertTokenizer.from_pretrained(f"{checkpoint}")
model = TFBertForMaskedLM.from_pretrained(f"{checkpoint}")

inputs = tokenizer(f"The capital of France is {mask}.", return_tensors="tf")
logits = model(**inputs).logits

# retrieve index of {mask}
mask_token_index = tf.where(inputs.input_ids == tokenizer.mask_token_id)[0][1]
predicted_token_id = tf.math.argmax(logits[0, mask_token_index], axis=-1)
expected_output = tokenizer.decode(predicted_token_id)

print(mask_token_index)  # tf.Tensor(8, shape=(), dtype=int64): no row dimension
print(predicted_token_id)  # tf.Tensor(3000, shape=(), dtype=int64)
print(tokenizer.decode(predicted_token_id))  # p a r i s (not good)

# retrieve index of {mask}
mask_token_index = tf.where((inputs.input_ids == tokenizer.mask_token_id)[0])
selected_logits = tf.gather_nd(logits[0], indices=mask_token_index)
predicted_token_id = tf.math.argmax(selected_logits, axis=-1)
expected_output = tokenizer.decode(predicted_token_id)

print(mask_token_index)  # tf.Tensor([[8]], shape=(1, 1), dtype=int64): with row dimension
print(predicted_token_id)  # tf.Tensor([3000], shape=(1,), dtype=int64)
print(tokenizer.decode(predicted_token_id))  # paris