Build and train a recommender system in 10 minutes using Keras and JAX

Today, we are excited to announce the launch of Keras Recommenders, a new library that puts state-of-the-art recommendation techniques at your fingertips.

Power digital experiences with recommendation systems

Recommendation systems power many of the interactions you have with technology today. Open up any app on your phone and you’ll likely find yourself interacting with a recommendation model right away, from the homefeed on your go-to social media platform to video suggestions on YouTube to even the ads that pop up in your favorite game. As the world of AI continues to evolve, delivering personalized experiences is more important than ever. Large language models can’t do everything, and recommender systems are responsible for creating many top-tier digital experiences today.

To help developers create performant and accurate recommender systems, Keras Recommenders (KerasRS) contains a set of APIs with building blocks designed for tasks such as ranking and retrieval. For example, at Google, we use KerasRS to help power the feed in Google Play.

Install KerasRS with JAX, TensorFlow, or PyTorch

To get started, pip install the keras-rs package. Then set the backend to JAX (or TensorFlow or PyTorch). Now you are on your way to crafting your own state-of-the-art recommender system.

import os
os.environ["KERAS_BACKEND"] = "jax"

import keras
import keras_rs

class SequentialRetrievalModel(keras.Model):
    def __init__(self):
        self.query_model = keras.Sequential([
            keras.layers.Embedding(query_count, embed_dim),
            keras.layers.GRU(embed_dim),
        ])
        self.candidate_model = keras.layers.Embedding(candidate_count, embed_dim)
        self.retrieval = keras_rs.layers.BruteForceRetrieval(k=10)
        self.loss_fn = keras.losses.CategoricalCrossentropy(from_logits=True)

    def call(self, inputs):
        query_embeddings = self.query_model(inputs)
        predictions = self.retrieval(query_embeddings)
        return {"query_embeddings": query_embeddings, "predictions": predictions}

Python

In this example, we show a popular retrieval architecture in which we identify a set of candidate recommendations. KerasRS provides everything you need to implement this architecture, with specialized layers, losses, and metrics designed specifically for recommender tasks. You can also follow along in this colab notebook.

And of course, all these building blocks work with the standard Keras APIs of model.compile to build your model and model.fit to easily configure your training loop.

model.compile(
    loss=keras_rs.losses.PairwiseHingeLoss(),
    metrics=[keras_rs.metrics.NDCG(k=8, name="ndcg")],
    optimizer=keras.optimizers.Adagrad(learning_rate=3e-4),
)
model.fit(train_ds, validation_data=val_ds, epochs=5)

Python

In the coming months, we plan to release the keras_rs.layers.DistributedEmbedding class for leveraging SparseCore chips on TPU for doing large embedding lookups distributed across machines. Additionally, we will add popular model implementations to our library continuously, making it even easier to build state-of-the-art recommender systems.

Explore the KerasRS documentation and examples

We also want to highlight all the documentation we have for Keras Recommenders on our recently redesigned keras.io website. On keras.io/keras_rs, you can find starter examples involving the classic Deep and Cross Network (DCN) and two-tower embedding model that show the step-by-step processes for writing and training your first recommender. There are also more advanced tutorials, such as SASRec, showing an end-to-end example of training a transformer model.

Get started

Visit our website today for more examples, documentation, and guides to build your very own recommendation system. You can also browse the code and contribute at https://github.com/keras-team/keras-rs (feel free to give it a star ⭐ too while you’re there!).

We look forward to seeing all the excellent recommendation systems that get built with Keras Recommenders.

Acknowledgements

^{Shout-out to Fabien Hertschuh and Abheesht Sharma for building Keras Recommenders. We also want to thank the Keras and ML Frameworks teams as well as all our collaborators and leadership for helping us pull this off.}