TensorFlow.js is giving me random probabilities for predictions - tensorflow

Question

I've been stuck with this problem for a while and I can't find a solution here on StackOverflow for this case.
I'm trying to build a prediction model for a chatbot. The predictor is a message in a bag of words format, in a tensor of shape [1,108] while the target variable is another tensor with dummies from a nominal classificatory variable, with shape [1,15].
When I run a prediction it returns an array with random probabilities for each of the dummies.
I used brain.js using the same functions and it gave me a good prediction, however I am having trouble doing the same thing here on TensorFlow.
Here is my code:
async function trainModel() {
XandY = await createTrainingData()
let X = XandY[0];
var y = XandY[1];
const inputShape = [X.length,X[0].length] // [# of observations, # of features]
const outputShape = [y.length,y[0].length]
X = tf.tensor2d( X, inputShape )
y = tf.tensor2d( y, outputShape)
traningDataObject = {
data: X,
target: y,
}
const model = tf.sequential();
model.add(tf.layers.dense(
{ units: 128, activation: 'relu', inputShape: [108] }));
model.add(tf.layers.dense(
{ units: 64, activation: 'relu' }));
model.add(tf.layers.dense(
{ units: 32, activation: 'relu' }));
model.add(tf.layers.dense(
{ units: 15, activation: 'softmax' }));
model.compile({
optimizer: tf.train.adam(0.01),
loss: 'categoricalCrossentropy',
metrics: ['accuracy']
});
model.fit(X, y,
{epochs: 150, validationData: [X,y]});
return model
}
async function getPrediction(message,wordset) {
model = await trainModel();
bow = bagOfWords(message, wordset);
const input = tf.tensor([bagOfWords(message, wordset)]);
var prediction = model.predict(input);
var prediction_values = prediction.dataSync();
var prediction_array = Array.from(prediction_values);
console.log(prediction_array)
let greatestProba = 0;
prediction_array.forEach((element) => {
if (greatestProba < element) {
greatestProba = element;
}
});
if (greatestProba > 0.02) {
return intents[prediction_array.indexOf(greatestProba)];
} else {
return 'undefined'
}
}
async function main(){
const wordset = await getWordset();
const message = "Meu pagamento não caiu";
console.log(await getPrediction(message,wordset));
}
main()
How can I solve this? Where is the problem?
Using the same network in python gives me good predictions, but here not.

Answer 1

Alright, so the problem seemed to be the number of Epochs I used...
It is odd because I'm pretty sure I used a very reasonable number of epochs...
I removed the ValidationData as well...
If you do it like this:
const res = await model.fit(X, y, { epochs: 50 });
console.log(res.history.loss[0]);
it should work.