TensorFlow.js is giving me random probabilities for predictions

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.