Create Multi-Output Model with KerasNLP (CLS + MLM) - tensorflow

Question

I am trying to use Keras_NLP to pretrain a model through both Masked Language Modelling and using a proxy for the CLS token to classify certain labels. With just Masked Language Modelling, it works fine but when I try to make it multi-output and also classify labels using the CLS token it seems to be unable to work
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
!pip install -q keras-nlp
import keras_nlp
import numpy as np
#Creating Random Data
x = np.random.randint(3, 100, size = (1000, 140))
#Adding Int 2 at front of each vector to act as CLS
x_cls = np.hstack([np.full((1000,1),2), x])
y = np.random.randint(0, 3, size = (1000,1)) #Random Class for each x value
#Need data for actual problem in dataset format, so using it here
dataset = tf.data.Dataset.from_tensor_slices((x_cls, y)).batch(32)
#Keras Masker
masker = keras_nlp.layers.MLMMaskGenerator(99, 0.15, mask_token_id = 0, unselectable_token_ids= [0,1,2], mask_selection_length = 15)
def process(input, y):
outputs = masker(input)
features = {
"tokens": outputs["tokens"],
"mask_positions": outputs["mask_positions"],
}
labels = outputs["mask_ids"]
weights = outputs["mask_weights"]
return features, labels, weights , y
#Apply Mask to Dataset
dataset_mask = dataset.map(lambda x, y: process(x,y))
#Create Simple Encoder, for testing purposes
input = keras.Input(shape = (141,))
embedding = layers.Embedding(99, 20)(input)
x = layers.Dense(20, activation = "relu")(embedding)
encoder = keras.Model(inputs = input, outputs = x)
#MLM prediction + CLS prediction
inputs = {
"tokens": keras.Input(shape=(141,), dtype=tf.int32),
"mask_positions": keras.Input(shape=(15,), dtype=tf.int32),
}
encoded_tokens = encoder(inputs["tokens"])
outputs = keras_nlp.layers.MLMHead(vocabulary_size = 141, activation="softmax")(encoded_tokens, inputs["mask_positions"])
x = layers.Lambda(lambda x: x[:,0,:])(encoded_tokens)
#Use lambda layer to extract embedding from 1st dim, corresponds to CLS
output_2 = layers.Dense(3, activation = "softmax")(x)
pretraining_model = keras.Model(inputs, outputs = [outputs, output_2])
pretraining_model.compile(loss=["sparse_categorical_crossentropy" "sparse_categorical_crossentropy"], optimizer = "Adam",weighted_metrics=["sparse_categorical_accuracy"], jit_compile=True)
pretraining_model.fit(dataset_mask)
The error I get is
ValueError: Data is expected to be in format `x`, `(x,)`, `(x, y)`, or `(x, y, sample_weight)`, found: ({'tokens': <tf.Tensor 'data_1:0' shape=(None, 141) dtype=int64>, 'mask_positions': <tf.Tensor 'data:0' shape=(None, 15) dtype=int64>}, <tf.Tensor 'data_2:0' shape=(None, 15) dtype=int64>, <tf.Tensor 'data_3:0' shape=(None, 15) dtype=float32>, <tf.Tensor 'data_4:0' shape=(None, 1) dtype=int64>)
Any help would be appreciated!