Hi folks,
Spent way too much time on this, but as a hint, check the following:

Are you using padding with -100? Then be sure you mask it in the loss and metric!
model.fit does not do this for you automatically, whereas TFTrainer with TFTokenClassificationLoss does this for you!
Be sure to check the activation of the last layer, by default it is linear, which means you need to use (Sparse)CategoricalEntropy(from_logits=True); But then your accuracy scores will be messed up by default!
Essentially, align your final activation to be softmax, use from_logits=False, write a custom loss & metric wrapper for masking out padding & special tokens.


def sparse_categorical_accuracy_masked(y_true, y_pred):
    mask_value = -100
    active_loss = tf.reshape(y_true, (-1,)) != mask_value
    reduced_logits = tf.boolean_mask(tf.reshape(y_pred, (-1, shape_list(y_pred)[2])), active_loss)
    y_true = tf.boolean_mask(tf.reshape(y_true, (-1,)), active_loss)
    reduced_logits = tf.cast(tf.argmax(reduced_logits, axis=-1), tf.keras.backend.floatx())
    equality = tf.equal(y_true, reduced_logits)
    return tf.reduce_mean(tf.cast(equality, tf.keras.backend.floatx()))



def sparse_crossentropy_masked(y_true, y_pred):
    mask_value = -100
    y_true_masked = tf.boolean_mask(y_true, tf.not_equal(y_true, mask_value))
    y_pred_masked = tf.boolean_mask(y_pred, tf.not_equal(y_true, mask_value))
    return tf.reduce_mean(tf.keras.losses.sparse_categorical_crossentropy(y_true_masked, 
y_pred_masked))