Marking a TensorflowJs model to predict user movie preferences

DanielDecember 13, 2023

Let's build a TensorflowJs model that will act as a movie recommendation system.

It's a super simplified example, aiming just to demo how something like this can be done with a Javascript neuronal network.

Our model will take into account just two characteristics of a given movie:

• how long is the movie
• and the image quality of the movie

The training data is stored in a 2d tensor.

const movieFeatures = tf.tensor2d([
    // image quality , lenght
    [0.1, 0.1],
    [0.3, 0.3],
    [0.5, 0.6],
    [0.4, 0.8],
    [0.9, 0.1],
    [0.75, 0.4],
    [0.75, 0.9],
    [0.6, 0.9],
    [0.6, 0.75]
])

// 1 - liked the movie, 
// 0 - did not like the movie
const results = [0, 0, 1, 1, 0, 0, 1, 1, 1]

In order to normalize the data both characteristics of a movie are represented by values between 0 and 1.

If you take a zoom-out view of the data you can observe that our users enjoy movies that are both long and high-quality.

The aim is to build a model that can take a value for the image quality and for length of a movie and will predict if the user will like that movie or not:

const movie1 = [0.1, 0.6]
const movie2 = [0.7, 0.1]
model.predict(tf.tensor2d([movie1]))
model.predict(tf.tensor2d([movie2]))

Let's get to work 💪 🧑‍💻!!!

Building the neuronal network

The first step will be to build the neuronal network that will solve our problem:

We will use a TensorflowJs sequential model with :

• 2 input neurons; one for the movie length and one for the image quality score
• a hidden layer made of 3 neurons
• and an output of 2 neurons; each out neuron will return a value between 0 and 1 for how likely is for the user to like or not the given movie

The below code will implement this neuronal network:

let model

const buildModel = ()=> {
   model = tf.sequential();

   model.add(
      tf.layers.dense({
       inputShape: [2],
        units: 3,
        activation: "relu"
      })
   );

   model.add(
     tf.layers.dense({
       units: 2,
       activation: "softmax"
     })
   )

    model.compile({
       optimizer: tf.train.adam(0.1),
       loss: "binaryCrossentropy",
       metrics: ["accuracy"]
    })
}

Training the TensorflowJs model

Having the neuronal network ready it's time to train it. In order to do this we will use the data from the movieFeatures and results tensors defined above.

Before we actually pass the training data, we will need to do one hot encoding on the results tensors, so that we can make predictions with the format:

[
    chances_of_enjoying_the_movie,
    chances_of_not_enjoying_the_movie
]

The below code will do the one-hot encoding for the labels and train the model:

const encodeOneHot = (val, categoryCount) =>
  Array.from(tf.oneHot(val, categoryCount).dataSync())

const buildModel = ()=> {
  // build the neuronal network
  // ...
  // train the model
  const labels = tf.tensor(results.map(y => encodeOneHot(y, 2)))
  const options = {
    shuffle: true,
    epochs: 20
  }
  return model.fit(movieFeatures, lables, options)
}

Making predictions

Having the TensorflowJs model defined and trained we can use it to predict if a movie will be liked by a user or not.

Remember that the Tensorflow fit() method is async and returns a promise, so we will first need to wait for the promise to resolve and then we can call the new makePredictions() method:

buildModel().then(makePredictions)

And we can now test some movies:

const makePredictions = ()=> {
  const movie = [[0.8, 0.7]]
  model.predict(tf.tensor2d([movie])).print()
}

The model will output a tensor like the one below:

Tensor
     [[0.3198042, 0.6801958],]

This means that the probability of the user not enjoying the movie is smaller than the probability of the user liking the movie. You can use the TensorflowJs argMax() function to directly extract the result from this.

If you have enjoyed this tutorial you can also checkout this example of a TensorflowJs model that will predict the price of a house. It's a simpler version that only takes one feature into account.