Keras Layer to map raw indices to labels and only outputing the top n predictions (tensorflow 2)

MappingLayer.py

class MappingLayer(tf.keras.layers.Layer):
    """
    Converts probability outputs to TopN predictions applying a label mapping from indices to ids
    """

    def __init__(self, mapping, topN):
        super(MappingLayer, self).__init__()
        initializer = tf.lookup.KeyValueTensorInitializer(
            keys=tf.range(len(mapping)),
            values=tf.constant(mapping, tf.int32),
            key_dtype=tf.int32,
            value_dtype=tf.int32)
        self.table = tf.lookup.StaticHashTable(initializer, default_value=-1)
        self.topN = topN

    def call(self, input, **kwargs):
        probas, indices = tf.math.top_k(
            input, k=self.topN, sorted=True, name='mapped_out'
        )
        mapped = self.table.lookup(indices)
        return [probas, mapped]