Add an MLMHead layer

This layer can go with the MLMMaskGenerator layer to make it easy to write a masked language model transformer.

The layer should map two inputs:

• embedded tokens of shape (batch, sequence_length, model_dim)
• mask positions of shape (batch, prediction_length)

To one output:

• predictions of shape (batch, prediction_length, vocabulary_size)

Where the outputs contains either logits or softmax probabilities per word in the vocabulary.

The layer should optionally allow reusing the embedding weights to reverse word embeddings from model_dim -> vocabulary_size.

[mattdangerw]

@aflah02 thanks, but this one I already have 😃, made it while working on a guide for keras-nlp. We will open up more issues this week!

[aflah02]

Sure I’d be interested to check that out!