Tab Transformer Predictions

I am working on implementing the tab transformer model into my dataset, and have it working so far and giving a fair ROC_AUC.I am using the code as seen below for my dataset. The only variation I have made is splitting the train dataset into train and validation, rather than using the test dataset as validation as they do below. https://keras.io/examples/structured_data/tabtransformer/ , https://colab.research.google.com/github/keras-team/keras-io/blob/master/examples/structured_data/ipynb/tabtransformer.ipynb

I am now wanting to create a confusion matrix and ROC_AUC plot for my train and test dataset, but am struggling to work out how to do so. Within the run experiment function, I attempted to split the train_dataset into X_train and y_train however cannot do that because the previous function returns a cache dataset which is needed.

What I would like to do is run code like this:

y_predicted_train = model.predict(X_train)
y_pred_train= []
for pred in y_predicted_train:
    if pred > 0.5:
        y_pred_train.append(1)
    else:
        y_pred_train.append(0)
y_predicted_test = model.predict(X_test)
y_pred_test= []
for pred in y_predicted_test:
    if pred > 0.5:
        y_pred_test.append(1)
    else:
        y_pred_test.append(0)   

And as soon as I have my y_pred_test and y_pred_train I can compute the metrics i need. But I am unsure how to get to that point. Any help is greatly appreciated