Building a Recommenders System for Massive Dataset using Keras-RS and SparseCore

Imagine you have a massive dataset of millions of user-item interactions, and you need to build a recommendation system that can scale seamlessly without bottlenecking your hardware. How can we achieve this? The answer lies in leveraging specialized hardware accelerators and modern frameworks.

In this tutorial, we’ll dive into building a highly scalable book recommendation system using a Books/Ratings dataset, leveraging the power of Keras 3, JAX, and the specialized keras-rs library to unlock the potential of TPU SparseCore.

By the end of this tutorial, you’ll understand:

• How to initialize native TPU distribution using JAX in Keras 3.
• How to preprocess and densify sparse datasets so they can be effectively learned by the model.
• How to build a Two-Tower retrieval model utilizing keras_rs Distributed Embeddings for SparseCore acceleration.
• How to construct custom data loaders to handle JAX/SparseCore data routing.
• How to generate and evaluate top-k recommendations from your model.

Before we start, we need to create the TPU Cloud VM to run our code. We will try to utilize a Cloud TPU for our tutorial. There are several step-by-steps that need to be done for utilizing the Cloud TPU.

1. Create a project in Google Cloud Console.
2. Createn a new VM in our Google Cloud project by utilizing this code.

gcloud alpha compute tpus tpu-vm create tpu-rs \
  --project=<your project id> \
  --zone=europe-west4-a \
  --accelerator-type=v6e-1 \
  --version=v2-alpha-tpuv6e \
  --provisioning-model=SPOT

gcloud alpha compute tpus tpu-vm create tpu-tunix

• **gcloud alpha**: This invokes the "alpha" release track of the CLI, meaning you are accessing early-release, experimental, or preview features that aren't yet in general availability.
• **compute tpus tpu-vm create**: This instructs Google Cloud to create a new TPU architecture where the VM is physically attached to the TPU host (unlike older Node architectures where the VM and TPU communicated over the network). This bare-metal access is what allows JAX to compile and execute with zero network bottleneck.
• **tpu-rs**: This is simply the custom name assigned to the instance.
• **--project=<google project ID>** Directs the billing and resource allocation to this specific Google Cloud Project ID.
• **--zone=europe-west4-a** Specifies the data center that used and this the place where the hardware will be provisioned.
• **--accelerator-type=v6e-1 This defines the exact hardware slice. `v6e** designates the 6th generation TPU (Trillium). The-1` indicates a single-chip topology (one TPU core). It is the ideal entry point for building and debugging a training script before executing a complex sharding strategy across a larger v6e-4 or v6e-8 pod.
• **--version=v2-alpha-tpuv6e** This selects the foundational software image loaded onto the VM. It ensures the operating system has the correct drivers and XLA compiler versions specifically built to support the alpha v6e hardware.

3. After creating the TPU VM, we can login to the TPU VM by using the CLI.

gcloud compute tpus tpu-vm ssh tpu-tunix --zone=europe-west4-a

4. We copy all the dataset to the Cloud Storage Bucket so it can be used by the training process that run in the TPU VM.
5. Environment & Hardware Setup

Before writing any model logic, we need to instruct Keras to use JAX as its backend. JAX is exceptionally fast on TPUs thanks to XLA (Accelerated Linear Algebra) compilation. We then initialize the TPU distribution natively through Keras 3’s DataParallel strategy.

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

import jax
import jax.numpy as jnp
import keras
import keras_rs
import numpy as np
import pandas as pd

def setup_hardware():
    try:
        distribution = keras.distribution.DataParallel(devices=jax.devices("tpu"))
        keras.distribution.set_distribution(distribution)
        print(f"JAX TPU Strategy Initialized with {jax.device_count()} devices.")
    except Exception as e:
        print(f"Running without TPU/Distribution Strategy: {e}")

This strategy automatically splits our data and computations across all available TPU cores, which is essential for training on massive datasets without hitting memory bottlenecks.

2. Loading and Preprocessing the Dataset

In this article, we use the book recommendation dataset from Kaggle. Real-world recommendation data is notoriously sparse (most users only rate a few items, and most books have very few reviews). Our first step is to densify the data. We filter out users and books with fewer than 50 interactions. This ensures the model has enough signal to learn meaningful patterns rather than overfitting to noise.

Once densified, we map the raw string IDs (like ISBNs) into sequential integers, as neural networks require numeric inputs. Finally, we perform a user-based Train/Test split, ensuring we hold out a specific fraction of interactions for every single user to properly evaluate the system later.

def load_split_and_preprocess_data(bucket_path='gs://your-bucket-name/path/to', file_name='Ratings.csv', test_fraction=0.2):
    print("Loading and densifying data...")
    file_path = f"{bucket_path}/{file_name}"
    ratings = pd.read_csv(file_path)

    user_counts = ratings['User-ID'].value_counts()
    book_counts = ratings['ISBN'].value_counts()

    active_users = user_counts[user_counts >= 50].index
    popular_books = book_counts[book_counts >= 50].index

    filtered_ratings = ratings[
        (ratings['User-ID'].isin(active_users)) &
        (ratings['ISBN'].isin(popular_books)) &
        (ratings['Book-Rating'] >= 1)
    ].copy()

    user_ids = filtered_ratings['User-ID'].astype('category')
    book_ids = filtered_ratings['ISBN'].astype('category')

    filtered_ratings['user_int_id'] = user_ids.cat.codes
    filtered_ratings['book_int_id'] = book_ids.cat.codes

    num_users_true = len(user_ids.cat.categories)
    num_books_true = len(book_ids.cat.categories)

    num_users = num_users_true + 1
    num_books = num_books_true + 1

    test_df = filtered_ratings.groupby('user_int_id').sample(frac=test_fraction, random_state=42)
    train_df = filtered_ratings.drop(test_df.index)

    print(f"Train set size: {len(train_df)} ({len(train_df)/len(filtered_ratings)*100:.1f}%)")
    print(f"Test set size:  {len(test_df)} ({len(test_df)/len(filtered_ratings)*100:.1f}%)")

    train_user_array = train_df['user_int_id'].to_numpy(dtype=np.int32)
    train_book_array = train_df['book_int_id'].to_numpy(dtype=np.int32)

    return train_user_array, train_book_array, test_df, num_users, num_books, user_ids.cat.categories, book_ids.cat.categories

3. Building the Two-Tower Architecture with SparseCore

Our recommender system will be a retrieval model, aiming to learn embeddings for users and books such that the embeddings of books a user is likely to read are “close” to the user’s embedding in a latent space.

Standard embedding layers struggle when vocabulary sizes hit the millions because they consume too much high-bandwidth memory. Here, we use keras_rs.layers.TableConfig and DistributedEmbedding. This powerful feature instructs the TPU to route the embedding lookups to its SparseCore—specialized hardware designed specifically for sparse, massive table lookups, freeing up the main matrix multiplication units

class BookTwoTower(keras.Model):
    def __init__(self, num_users, num_books, embedding_dim=32, batch_size=512, **kwargs):
        super().__init__(**kwargs)

        # Define SparseCore Tables
        table_config = keras_rs.layers.TableConfig(
            name='sparse_embedding_table',
            vocabulary_size=max(num_users, num_books) + 1000,
            embedding_dim=embedding_dim,
            initializer="uniform", # Prevent zero-initialization NaNs
            optimizer=keras.optimizers.Adagrad(0.01)
        )

        # Distributed Embedding Layer (Translates lookups to JAX/SparseCore)
        self.sparse_embedding = keras_rs.layers.DistributedEmbedding({
            'user_id': keras_rs.layers.FeatureConfig(name='user_id_feature', table=table_config, input_shape=(batch_size, 1), output_shape=(batch_size, embedding_dim)),
            'book_id': keras_rs.layers.FeatureConfig(name='book_id_feature', table=table_config, input_shape=(batch_size, 1), output_shape=(batch_size, embedding_dim))
        })

        # Keras 3 Dense Towers
        self.user_tower = keras.Sequential([
            keras.layers.Dense(64, activation="swish"),
            keras.layers.Dense(embedding_dim),
            keras.layers.UnitNormalization(axis=-1) # Safe normalization
        ])

        self.book_tower = keras.Sequential([
            keras.layers.Dense(64, activation="swish"),
            keras.layers.Dense(embedding_dim),
            keras.layers.UnitNormalization(axis=-1) # Safe normalization
        ])

    def call(self, inputs):
        """Forward pass to generate embeddings."""
        embeddings = self.sparse_embedding(inputs)
        user_vector = self.user_tower(embeddings['user_id'])
        book_vector = self.book_tower(embeddings['book_id'])

        return user_vector, book_vector

    def compute_loss(self, x=None, y=None, y_pred=None, sample_weight=None):
        """Pure JAX loss computation."""
        user_vector, book_vector = y_pred

        # Compute In-Batch Softmax Loss (dot product of all pairs)
        logits = keras.ops.matmul(user_vector, keras.ops.transpose(book_vector))

        # Scale logits by temperature because vectors are L2-normalized
        logits = logits / 0.05

        # Positive pairs are on the diagonal
        batch_size = keras.ops.shape(logits)[0]
        labels = keras.ops.arange(batch_size)

        loss = keras.losses.sparse_categorical_crossentropy(
            labels, logits, from_logits=True
        )
        return keras.ops.mean(loss)

Instead of calculating the loss against millions of negative books, we use an In-Batch Softmax Loss. We compute the dot product of all users and books currently in the batch, scale it by a temperature variable (/ 0.05) to sharpen the gradients, and use the diagonal (the matching user-book pairs) as our positive ground truth.

Because we rely on custom Keras-RS distributed embeddings, standard dataloaders won’t work out of the box. We need to explicitly tell the model to route the features to the SparseCore before the forward pass happens. We do this by calling self.embedding_layer.preprocess(features, training=True) inside our custom PyDataset.

class JAXSparseCoreDataset(keras.utils.PyDataset):
    def __init__(self, user_array, book_array, batch_size, embedding_layer, **kwargs):
        super().__init__(**kwargs)
        self.user_array = user_array
        self.book_array = book_array
        self.batch_size = batch_size
        self.embedding_layer = embedding_layer

    def __getitem__(self, idx):
        start_idx = idx * self.batch_size
        end_idx = (idx + 1) * self.batch_size

        features = {
            'user_id': np.expand_dims(self.user_array[start_idx:end_idx], axis=-1),
            'book_id': np.expand_dims(self.book_array[start_idx:end_idx], axis=-1)
        }

        return self.embedding_layer.preprocess(features, training=True)

    def on_epoch_end(self):
        indices = np.arange(len(self.user_array))
        np.random.shuffle(indices)
        self.user_array = self.user_array[indices]
        self.book_array = self.book_array[indices]

5. Training and Generating Recommendations

Finally, we compile our model and pass our customized dataset. We define several parameter including the dataset path. Path for save model and some hyperparameter for Two Tower Training.

BUCKET_PATH = "gs://dataset_rs/dataset"
SAVE_PATH = "gs://dataset_rs"
FILE_NAME = "Ratings.csv"

EPOCHS = 75
LEARNING_RATE = 0.001
BATCH_SIZE = 64
EMBEDDING_DIM = 16

Next step is we try to setup the TPU by calling setup_hardware() function. The next steps is try to prepare the dataset and split into training and testing. The next steps is to build the model for the Recommendation System.

# 1. Setup Hardware
setup_hardware()

# 2. Prepare Data and Split
(train_user_data, train_book_data, test_df, n_users, n_books, user_mapping, book_mapping) = load_split_and_preprocess_data(BUCKET_PATH, FILE_NAME)

# 3. Build Model
print("\nBuilding model...")
model = BookTwoTower(n_users, n_books, embedding_dim=EMBEDDING_DIM, batch_size=BATCH_SIZE)

To prevent Keras from getting confused during initialization with our dynamic SparseCore layers, we eagerly pass a “dummy batch” through the model first. This forces Keras to build all the weights cleanly before training begins!

# Eagerly initializing model to bypass JAX tracing issues
dummy_features = {
    'user_id': np.zeros((BATCH_SIZE, 1), dtype=np.int32),
    'book_id': np.zeros((BATCH_SIZE, 1), dtype=np.int32)
}
preprocessed_dummy = model.sparse_embedding.preprocess(dummy_features, training=True)
_ = model(preprocessed_dummy) 

model.compile(optimizer=keras.optimizers.Adam(learning_rate=LEARNING_RATE))

Once the model is trained, finding recommendations is simply a matter of calculating the dot product between a specific user’s vector and the vector of all books in our catalog, sorting the scores, and masking out the books the user has already read so we only recommend fresh content.

train_dataset = JAXSparseCoreDataset(
        train_user_data,
        train_book_data,
        batch_size=BATCH_SIZE,
        embedding_layer=model.sparse_embedding
    )
print("\nStarting Training...")
model.fit(train_dataset, epochs=EPOCHS)

Instead of using a standard tf.data.Dataset or a basic Keras loader, the code uses the custom JAXSparseCoreDataset which is a Custom Keras 3 Data Loader from keras.utils.PyDataset . And we continue with training process. After the training process is done, we can utilize the model to create the recommendation system. We create a custom data loader for Evaluation or Inference the model.

class EvalSparseCoreDataset(keras.utils.PyDataset):
    """Dataset for extracting embeddings safely via Keras predict."""
    def __init__(self, user_array, book_array, batch_size, embedding_layer, **kwargs):
        super().__init__(**kwargs)
        remainder = len(user_array) % batch_size
        if remainder != 0:
            pad_len = batch_size - remainder
            self.user_array = np.pad(user_array, (0, pad_len), mode='constant')
            self.book_array = np.pad(book_array, (0, pad_len), mode='constant')
            self.actual_len = len(user_array)
        else:
            self.user_array = user_array
            self.book_array = book_array
            self.actual_len = len(user_array)

        self.batch_size = batch_size
        self.embedding_layer = embedding_layer

    def __len__(self):
        return len(self.user_array) // self.batch_size

    def __getitem__(self, idx):
        start_idx = idx * self.batch_size
        end_idx = (idx + 1) * self.batch_size

        features = {
            'user_id': np.expand_dims(self.user_array[start_idx:end_idx], axis=-1),
            'book_id': np.expand_dims(self.book_array[start_idx:end_idx], axis=-1)
        }

        return self.embedding_layer.preprocess(features, training=False)

TPUs process data using XLA (Accelerated Linear Algebra). XLA compiles a highly optimized, static computational graph based on the exact input shapes it receives. Suppose your batch size is 512, but your evaluation dataset has 1,030 users.

• Batch 1: 512 users (XLA compiles the graph).
• Batch 2: 512 users (XLA reuses the graph — blazing fast).
• Batch 3: 6 leftover users. (XLA panics. The shape changed from 512 to 6, so it drops everything and painstakingly recompiles a brand new graph just for those 6 users).

This recompilation can cause massive latency spikes, make your inference freeze, or crash the TPU entirely out of memory. The code calculates the remainder. If there are leftover items, it uses np.pad to append dummy zeros (mode='constant') to the end of the array until the final batch perfectly equals 512.

Now, XLA only ever sees one shape, completely bypassing the recompilation trap. Because we just injected a bunch of fake “Zero” users and books at the end of our dataset to appease the TPU, the model will generate embeddings for them. We need to remember exactly where the real data stopped (self.actual_len).

If you look back at the evaluation function in the main script, you’ll see this: user_catalog_embeddings = user_preds[:user_eval_ds.actual_len] This cleanly slices off the padded dummy embeddings so they don't corrupt your actual recommendation metrics.

Just like in the training dataloader, we must explicitly route the features through keras_rs to hit the SparseCore. However, we set training=False.

• For standard Keras layers, this turns off things like Dropout.
• For SparseCore, this tells the specialized hardware that we are only doing lookups (reads) and should absolutely not calculate gradients or attempt to update the massive embedding tables (writes).

In the context of your recommendation system, the neural network only outputs cold, hard numbers (like “User 5 will like Book ID 1042”). This block of code is responsible for loading the actual book details (Titles and Authors) so that at the end of your script, you can print human-readable recommendations like “The Hobbit by J.R.R. Tolkien” instead of just “ISBN 0345339681”.

books_df = pd.read_csv(f"{BUCKET_PATH}/Books.csv", sep=';', encoding='latin-1', on_bad_lines='skip', low_memory=False)
if len(books_df.columns) == 1:
    books_df = pd.read_csv(f"{BUCKET_PATH}/Books.csv", sep=',', on_bad_lines='skip', low_memory=False)

Here is the line-by-line breakdown of what actually process in this code:

• **sep=';'**: It first tries to read the file assuming the columns are separated by semicolons.
• **encoding='latin-1'*: Book titles often have special characters or accents (e.g., Les Misérables*). Standard UTF-8 encoding sometimes crashes on these; latin-1 is a safer fallback for older text datasets.
• **on_bad_lines='skip'**: If a row is corrupted (e.g., it has 5 columns instead of 4), Pandas will just skip it instead of crashing your entire program.
• **low_memory=False**: Tells Pandas to read the whole file into memory before guessing the data types of each column, preventing annoying "mixed datatype" warnings.
• The if statement: This is a fallback. If the dataset was actually comma-separated, the semicolon read will result in a dataframe with just 1 giant column. If that happens, the script immediately re-reads the file using commas (sep=',').

books_df.columns = books_df.columns.str.replace('"', '').str.strip()

if 'ISBN' in books_df.columns and 'Book-Title' in books_df.columns:
    books_df['ISBN'] = books_df['ISBN'].astype(str).str.strip()
    books_df['Book-Author'] = books_df['Book-Author'].fillna('Unknown Author').astype(str)
    books_df['Book-Title'] = books_df['Book-Title'].fillna('Unknown Title').astype(str)

    meta_strings = books_df['Book-Title'] + " by " + books_df['Book-Author']
    isbn_to_meta = pd.Series(meta_strings.values, index=books_df['ISBN']).to_dict()
else:
    isbn_to_meta = {}

Sometimes CSV headers look like "ISBN", "Book-Title". This code strips away those extra quotation marks and any accidental trailing spaces so you can reliably call books_df['ISBN'] later.

The next process is for extracting the data. Here is the breakdown of our code.

• IF statement verifies the columns actually exist to prevent a KeyError crash.
• **astype(str).str.strip()**: Forces the ISBNs to be treated as text. This is crucial because if Pandas thinks an ISBN is a math number, it might delete leading zeros (turning 0123 into 123), which breaks everything.
• **fillna('...')**: Real-world data is messy. If a book has no listed author or title (a NaN value), it safely fills that blank space with the text "Unknown Author" so string operations later don't fail.
• **meta_strings: This concatenates the title and author columns into a single, pretty string. Example: “Harry Potter” + “ by “ + “J.K. Rowling” = “Harry Potter by J.K. Rowling”**
• **to_dict(): This is the magic step. It converts the two columns into a fast Python dictionary where the Key is the ISBN and the Value* is the pretty string. Example Output:* {"0439139597": "Harry Potter by J.K. Rowling", "0553588486": "A Game of Thrones by George R.R. Martin"}
• **else: isbn_to_meta = {}**: If the CSV loading failed completely, it creates an empty dictionary. This ensures your downstream evaluation code doesn't break; it will just print "Unknown Book" for everything instead of crashing the app.

user_mapping_list = list(user_mapping)
target_users = user_mapping_list[:5]

users_to_predict_int = []
user_warnings = {}

next_oov_id = len(user_mapping_list)

for u in target_users:
    if u in user_mapping_list:
        users_to_predict_int.append(user_mapping_list.index(u))
    else:
        users_to_predict_int.append(next_oov_id)
        user_warnings[u] = "WARNING: Data is not sufficient (user filtered out). The result may be false/random."
        next_oov_id += 1

print(f"\nExtracting all book catalog embeddings...")
book_ds = EvalSparseCoreDataset(np.zeros_like(all_unique_book_ids, dtype=np.int32), all_unique_book_ids, BATCH_SIZE, model.sparse_embedding)
_, book_emb = model.predict(book_ds, batch_size=BATCH_SIZE, verbose=0)
book_catalog = book_emb[:book_ds.actual_len]

print(f"Extracting embeddings for {len(users_to_predict_int)} users...")
user_ds = EvalSparseCoreDataset(np.array(users_to_predict_int, dtype=np.int32), np.zeros(len(users_to_predict_int), dtype=np.int32), BATCH_SIZE, model.sparse_embedding)
user_embs, _ = model.predict(user_ds, verbose=0)
user_embs = user_embs[:len(users_to_predict_int)]

book_id_to_index = {book_id: idx for idx, book_id in enumerate(all_unique_book_ids)}

In production systems, users sometimes will ask for recommendations for a user ID that doesn’t exist in the training data (e.g., a brand new user, or a user who was filtered out during our densification step because they had fewer than 50 reviews). This is known as the Cold Start Problem.

Instead of letting the application crash with an IndexError, this code elegantly handles it. If you recall from the model architecture earlier, we defined the SparseCore table vocabulary as max(num_users, num_books) + 1000. This loop checks if a user is known. If they aren't, it assigns them one of those 1,000 extra "blank" integer IDs. The model won't crash; it will simply return a randomly initialized embedding, and the code flags a warning that the resulting recommendations are essentially random guesses until that user reads more books.

Because our model is a BookTwoTower that expects both a user_id and a book_id in its forward pass, how do we extract just the book embeddings?

The answer is the Dummy Array Trick. Notice the np.zeros_like(all_unique_book_ids). We pass an array of all zeros for the user IDs. Because the User Tower and the Book Tower process their inputs completely independently (they only interact at the very end during the loss calculation), feeding "User 0" to the user tower has absolutely zero impact on the vector coming out of the book tower.

We perform this once to build the book_catalog (a matrix containing the embeddings of every book in existence), and then we do the exact same thing in reverse (using dummy books) to extract the user embeddings.

At this point, book_catalog is just a massive 2D NumPy array of numbers. This final line creates a "treasure map." It creates a dictionary mapping the actual book integer ID to the row index in that matrix. Now, we will create the code for extracting the result from the model as recommendation for some user inputs.

for i, user_int_id in enumerate(users_to_predict_int):
        user_str = target_users[i]
        u_emb = user_embs[i]

        seen_books_int = train_book_data[train_user_data == user_int_id]
        seen_books_str = [book_mapping[b] for b in seen_books_int if b < len(book_mapping)]

        seen_books_meta = []
        for isbn in seen_books_str:
            meta = isbn_to_meta.get(str(isbn), 'Unknown Book')
            if str(meta).lower() == 'nan' or 'nan by' in str(meta): meta = 'Unknown Book'
            seen_books_meta.append(f"{isbn} ({meta})")

        print("\n" + "="*50)
        if user_str in user_warnings:
            print(user_warnings[user_str])

        print(f"User {user_str} read {len(seen_books_str)} books. Sample:")
        for meta in seen_books_meta[:5]:
            print(f"  - {meta}")

        scores = np.dot(book_catalog, u_emb)
        seen_indices = [book_id_to_index[b] for b in seen_books_int if b in book_id_to_index]
        scores[seen_indices] = -np.inf

        top_10_indices = np.argsort(scores)[-10:][::-1]
        top_10_books = [all_unique_book_ids[idx] for idx in top_10_indices]
        top_10_isbn = [book_mapping[b] for b in top_10_books]
        top_10_scores = [scores[idx] for idx in top_10_indices]

        print(f"\nTop 10 Recommended books for user {user_str}:")
        for rank in range(10):
            isbn = top_10_isbn[rank]
            meta = isbn_to_meta.get(str(isbn), 'Unknown Book')
            if str(meta).lower() == 'nan' or 'nan by' in str(meta): meta = 'Unknown Book'
            print(f"{rank+1}. {meta}")
            print(f"   ISBN: {isbn} | Score: {top_10_scores[rank]:.4f}")

Before making new recommendations, the system needs to know what the user has already consumed. This block searches the original training arrays to find all the books (as integers) the user has interacted with, and maps them back to their original string ISBNs. The script then fetches the metadata (Title and Author) so we can print a nice summary of the user’s reading history.

The line of code scores=np.dot(book_catalog, u_emb)is the mathematical foundation model of the entire retrieval system.

• **book_catalog** is a massive 2D matrix containing the embeddings of every single book.
• **u_emb** is the 1D embedding vector for our specific user.
• By performing a dot product (np.dot), we are simultaneously calculating the similarity score between this user and every single book in existence in a fraction of a millisecond. Because we applied UnitNormalization to our model towers earlier, this dot product effectively calculates the Cosine Similarity. The higher the score, the closer the user and the book are in the latent embedding space.

We already calculated the similarity scores for millions of books. However, it’s highly likely that the books with the absolute highest scores are the ones the user has already read (because the model learned to pull their vectors together during training).

To ensure we only recommend novel items, we grab the index locations of the user’s previously read books and forcefully set their similarity scores to negative infinity (-np.inf). This guarantees they will drop to the absolute bottom of the ranked list.

Now that the scores are calculated and the read books are masked out, we need the top recommendations.

• np.argsort(scores) sorts the array and returns the indices of the items, from lowest score to highest score.
• [-10:] slices off just the last 10 indices (which represent the highest scores).
• [::-1] reverses that list so it goes from highest (Rank 1) to lowest (Rank 10).

Finally, the code loops through these top 10 indices, maps them back to their human-readable ISBNs and Metadata, and prints out the final personalized recommendation list for the user. This is the example result from the demonstration result form the Recommendation Model.

User 243 read 11 books. Sample:
  - 0060977493 (The God of Small Things by Arundhati Roy)
  - 0316601950 (The Pilot's Wife : A Novel by Anita Shreve)
  - 0316776963 (Me Talk Pretty One Day by David Sedaris)
  - 0316899984 (River, Cross My Heart by Breena Clarke)
  - 0375400117 (Memoirs of a Geisha by Arthur Golden)

Top 10 Recommended books for user 243:
 1.⁠ ⁠Gap Creek: The Story Of A Marriage by Robert Morgan
   ISBN: 0743203631 | Score: 0.3854
 2.⁠ ⁠From the Corner of His Eye by Dean R. Koontz
   ISBN: 0553801341 | Score: 0.3785
 3.⁠ ⁠Tishomingo Blues by Elmore Leonard
   ISBN: 0060083948 | Score: 0.3781
 4.⁠ ⁠Finding Moon by Tony Hillerman
   ISBN: 0061092614 | Score: 0.3744
 5.⁠ ⁠Shell Seekers by Rosamunde Pilcher
   ISBN: 0440202043 | Score: 0.3728
 6.⁠ ⁠Five Quarters of the Orange by Joanne Harris
   ISBN: 0060198133 | Score: 0.3705
 7.⁠ ⁠Strangers by Dean R. Koontz
   ISBN: 0425181111 | Score: 0.3704
 8.⁠ ⁠The Heart of a Woman by Maya Angelou
   ISBN: 0553380095 | Score: 0.3688
 9.⁠ ⁠Christmas Box (Christmas Box Trilogy) by Richard Paul Evans
   ISBN: 0684814994 | Score: 0.3675
10.⁠ ⁠Last Chance Saloon by Marian Keyes
   ISBN: 0380820293 | Score: 0.3673

In your Two-Tower model, User 243 is represented by a single integer ID. During training, the “User Tower” learned to adjust User 243’s embedding vector based on these 11 books. Because the user read these specific titles, the model pulled User 243’s vector closer to the vectors of these books in the multi-dimensional latent space.

We can see a clear pattern in User 243’s taste. They lean heavily toward popular 1990s/early 2000s literary fiction, memoirs, and book club favorites (e.g., Oprah’s Book Club picks like The Pilot’s Wife and River, Cross My Heart). There is a strong theme of dramatic, character-driven narratives, often with female protagonists.

This script took User 243’s vector and performed a massive mathematical operation (a dot product) against the vectors of every single book in your entire catalog. It then successfully applied your “masking trick” — ensuring none of the 11 books User 243 has already read appeared in this final list.

By examining the output for User 243, we can see the Two-Tower model working exactly as intended. Without being fed any text, genres, or publication dates, the model analyzed the user’s history of 90s book club favorites and successfully mapped them to similar historical dramas and era-appropriate bestsellers. The model successfully masked out previously read items, relying purely on the Cosine Similarity of the learned embeddings to rank the top 10 unseen books.

And there you have it! You’ve successfully built a highly optimized, production-ready recommendation system that scales effortlessly using Keras-RS and TPU SparseCores.

References:

1. https://docs.cloud.google.com/tpu/docs/intro-to-tpu
2. https://keras.io/keras_rs/examples/distributed_embedding_jax/
3. https://docs.jax.dev/en/latest/

*Thanks to Jonathan Kenrick for helping me on finishing this articles.

TPUSprint #Keras #KerasRS #SparseCore #TPU