DEV Community

Ahirton Lopes
Ahirton Lopes

Posted on

You've Got Mail📨 (and Recommendations!): Delivering Recs with Keras, JAX & KerasRS

#programming #ai #keras #jax

Recommendations are everywhere — from your email inbox to the shopping carts you abandon (but never escape 😅). Whether it’s suggesting a movie, a product, or a next best action, recommender systems have become fundamental to today’s digital experience.

Until recently, building robust recommender pipelines meant stitching together lots of custom layers, custom losses, and custom evaluation metrics by hand. That’s why the introduction of KerasRS is such a game-changer.

What is KerasRS?

KerasRS (Keras Recommenders) is an open-source extension for Keras 3 that delivers building blocks specifically designed for recommender systems, including:

  • ✅ retrieval layers
  • ✅ ranking layers
  • ✅ specialized recommender losses
  • ✅ ranking metrics

Best of all, KerasRS is multi-backend, working seamlessly with TensorFlow, PyTorch, and JAX. That means you can combine the familiar Keras API with the high-performance JAX compiler and TPU acceleration for your recsys workflows.

☝️ Did you know? The Google Play feed uses KerasRS behind the scenes! (source)

Why JAX + KerasRS?

If you want fast training, JAX is your secret weapon:

  • JIT compilation for speed
  • XLA acceleration
  • Automatic vectorization
  • TPU support

Pairing JAX with KerasRS means you get production-grade recommender building blocks with superior performance. It’s like having your cake and eating it too. 🍰

Installing KerasRS 

pip install keras-rs
# or, for the nightly:
pip install --upgrade keras-rs-nightly
Enter fullscreen mode Exit fullscreen mode

Then set your backend:

import os
os.environ["KERAS_BACKEND"] = "jax"
Enter fullscreen mode Exit fullscreen mode

A Quick Retrieval Example

Let’s build a minimal retrieval recommender in just a few lines.

import keras
import keras_rs
import numpy as np
import tensorflow as tf

# Dummy user-item pairs
user_ids = np.array([0, 1, 2, 3])
item_ids = np.array([10, 20, 30, 40])

user_embedding = keras.layers.Embedding(input_dim=4, output_dim=32)
item_embedding = keras.layers.Embedding(input_dim=50, output_dim=32)

retrieval = keras_rs.layers.BruteForceRetrieval(k=5)

query_model = keras.Sequential([
    keras.Input(shape=(), dtype=tf.int32, name="user_id"),
    user_embedding,
    keras.layers.Flatten()
])

candidate_model = keras.Sequential([
    keras.Input(shape=(), dtype=tf.int32, name="item_id"),
    item_embedding,
    keras.layers.Flatten()
])

retrieval.index_from_dataset(
    tf.data.Dataset.from_tensor_slices(item_ids).map(candidate_model)
)

model = keras_rs.RetrievalModel(
    query_model=query_model,
    candidate_model=candidate_model,
    retrieval=retrieval
)

model.compile(
    optimizer=keras.optimizers.Adagrad(3e-4),
    loss=keras_rs.losses.PairwiseHingeLoss(),
    metrics=[keras_rs.metrics.NDCG(k=5)],
)

features = {
    "user_id": user_ids,
    "item_id": item_ids,
}

model.fit(features, item_ids, epochs=5, verbose=2)
Enter fullscreen mode Exit fullscreen mode

That’s a minimal retrieval system — you can expand this with categorical features, embeddings, or even sequence models.

Going Further: Transformers & Two-Tower

KerasRS supports more advanced recommender architectures too:

  • Deep & Cross Networks (DCN)
  • Two-Tower models
  • Sequence-based recommenders with transformers
  • SASRec-style sequence recommenders

If you want to see a transformer-based recommender in action, check out this Movielens demo.

What’s Next for KerasRS?

🚀 KerasRS is on a fast-moving roadmap, with upcoming features such as:

  • DistributedEmbedding for large-scale TPU sharded embedding tables
  • SparseCore support
  • Ultra-scalable retrieval across billions of items

It’s a great time to build recsys pipelines with these tools.

Try it Yourself

If you’re excited to build your own recommender with KerasRS, check out:

KerasRS makes scalable, production-grade recommendation models delightfully easy. Give it a try, and share your experiments!

📨 So next time you see that “You might like…” email, remember: there’s probably a KerasRS model working hard behind the scenes. 😉

If you liked this article, feel free to leave a comment or share! 🚀

Top comments (0)