tensorflow-recommenders

import tensorflow as tf import tensorflow_datasets as tfds import tensorflow_recommenders as tfrs # Load the MovieLens 100K dataset ratings = tfds.load('movielens/100k-ratings', split="train") movies = tfds.load('movielens/100k-movies', split="train") # Prepare data by selecting relevant features ratings = ratings.map(lambda x: {"movie_title": x["movie_title"], "user_id": x["user_id"]}) movies = movies.map(lambda x: x["movie_title"]) # Build vocabularies for user IDs and movie titles user_ids_vocabulary = tf.keras.layers.StringLookup(mask_token=None) user_ids_vocabulary.adapt(ratings.map(lambda x: x["user_id"])) movie_titles_vocabulary = tf.keras.layers.StringLookup(mask_token=None) movie_titles_vocabulary.adapt(movies) # Define user and movie models using Keras Sequential user_model = tf.keras.Sequential([ user_ids_vocabulary, tf.keras.layers.Embedding(user_ids_vocabulary.vocabulary_size(), 32) ]) movie_model = tf.keras.Sequential([ movie_titles_vocabulary, tf.keras.layers.Embedding(movie_titles_vocabulary.vocabulary_size(), 32) ]) # Define the retrieval task with FactorizedTopK metric task = tfrs.tasks.Retrieval( metrics=tfrs.metrics.FactorizedTopK( candidates=movies.batch(128).map(movie_model) ) ) # Create a TFRS model class MovieLensModel(tfrs.Model): def __init__(self, user_model, movie_model): super().__init__() self.movie_model: tf.keras.Model = movie_model self.user_model: tf.keras.Model = user_model self.task: tf.keras.layers.Layer = task def compute_loss(self, features: dict, training=False) -> tf.Tensor: user_embeddings = self.user_model(features["user_id"]) positive_movie_embeddings = self.movie_model(features["movie_title"]) return self.task(user_embeddings, positive_movie_embeddings) # Compile and train the model model = MovieLensModel(user_model, movie_model) model.compile(optimizer=tf.keras.optimizers.Adagrad(0.1)) model.fit(ratings.batch(4096), epochs=3) # Generate recommendations (example for a specific user) index = tfrs.layers.factorized_top_k.BruteForce(model.user_model) index.index_from_dataset( movies.batch(100).map(lambda title: (title, model.movie_model(title))) ) # Example: get recommendations for user with ID '42' _, titles = index(tf.constant(["42"])) print(f"Top 3 recommendations for user '42': {titles[0, :3].numpy().astype(str)}")