Two-Tower desde cero: tutorial con TensorFlow Recommenders y MovieLens

pip install tensorflow==2.15 tensorflow-recommenders==0.7.3 numpy pandas

import urllib.request
import zipfile
import os

# Descargar MovieLens 100K
url = "https://files.grouplens.org/datasets/movielens/ml-100k.zip"
urllib.request.urlretrieve(url, "ml-100k.zip")

with zipfile.ZipFile("ml-100k.zip", "r") as zip_ref:
    zip_ref.extractall(".")

print("Descarga completa.")

import pandas as pd

# Cargar ratings
ratings = pd.read_csv(
    "ml-100k/u.data",
    sep="\t",
    names=["user_id", "movie_id", "rating", "timestamp"]
)

# Cargar información de películas
movies = pd.read_csv(
    "ml-100k/u.item",
    sep="|",
    encoding="latin-1",
    usecols=[0, 1],
    names=["movie_id", "title"]
)

# Unir
df = ratings.merge(movies, on="movie_id")
df["user_id"] = df["user_id"].astype(str)
df["movie_id"] = df["movie_id"].astype(str)

print(df.head())

print(f"Total de ratings:       {len(df):,}")
print(f"Usuarios únicos:        {df['user_id'].nunique():,}")
print(f"Películas en catálogo:  {df['movie_id'].nunique():,}")
print()
print("Distribución de ratings:")
print(df['rating'].value_counts().sort_index())

import tensorflow as tf
import numpy as np

# Crear tf.data.Dataset desde el DataFrame
dataset = tf.data.Dataset.from_tensor_slices({
    "user_id": df["user_id"].values,
    "movie_title": df["title"].values,
})

# Vocabularios
user_ids = df["user_id"].unique().tolist()
movie_titles = df["title"].unique().tolist()

print(f"Vocabulario de usuarios: {len(user_ids)}")
print(f"Vocabulario de películas: {len(movie_titles)}")

# División 80/20 (80k train, 20k test)
tf.random.set_seed(42)
shuffled = dataset.shuffle(100_000, seed=42, reshuffle_each_iteration=False)

train = shuffled.take(80_000)
test  = shuffled.skip(80_000).take(20_000)

# Batching
train = train.batch(256).cache()
test  = test.batch(256).cache()

print(f"Batches de entrenamiento: {len(list(train))}")
print(f"Batches de evaluación:    {len(list(test))}")

import tensorflow_recommenders as tfrs

class UserModel(tf.keras.Model):
    def __init__(self, user_ids):
        super().__init__()
        self.user_lookup = tf.keras.layers.StringLookup(
            vocabulary=user_ids, mask_token=None
        )
        self.user_embedding = tf.keras.layers.Embedding(
            len(user_ids) + 1, 32
        )
        self.dense1 = tf.keras.layers.Dense(64, activation="relu")
        self.dense2 = tf.keras.layers.Dense(32)

    def call(self, user_id):
        x = self.user_lookup(user_id)
        x = self.user_embedding(x)
        x = self.dense1(x)
        return self.dense2(x)


class MovieModel(tf.keras.Model):
    def __init__(self, movie_titles):
        super().__init__()
        self.movie_lookup = tf.keras.layers.StringLookup(
            vocabulary=movie_titles, mask_token=None
        )
        self.movie_embedding = tf.keras.layers.Embedding(
            len(movie_titles) + 1, 32
        )
        self.dense1 = tf.keras.layers.Dense(64, activation="relu")
        self.dense2 = tf.keras.layers.Dense(32)

    def call(self, movie_title):
        x = self.movie_lookup(movie_title)
        x = self.movie_embedding(x)
        x = self.dense1(x)
        return self.dense2(x)

# Dataset de películas para el índice de retrieval
movies_dataset = tf.data.Dataset.from_tensor_slices(movie_titles)

class MovielensModel(tfrs.Model):
    def __init__(self):
        super().__init__()
        self.user_model  = UserModel(user_ids)
        self.movie_model = MovieModel(movie_titles)

        self.task = tfrs.tasks.Retrieval(
            metrics=tfrs.metrics.FactorizedTopK(
                candidates=movies_dataset.batch(128).map(self.movie_model)
            )
        )

    def compute_loss(self, features, training=False):
        user_embeddings  = self.user_model(features["user_id"])
        movie_embeddings = self.movie_model(features["movie_title"])
        return self.task(user_embeddings, movie_embeddings)

model = MovielensModel()

model.compile(optimizer=tf.keras.optimizers.Adagrad(learning_rate=0.1))

early_stopping = tf.keras.callbacks.EarlyStopping(
    monitor="val_factorized_top_k/top_100_categorical_accuracy",
    patience=3,
    mode="max",
    restore_best_weights=True,
    verbose=1,
)

history = model.fit(
    train,
    validation_data=test,
    epochs=30,
    callbacks=[early_stopping],
    verbose=1,
)

# Construir el índice BruteForce
index = tfrs.layers.factorized_top_k.BruteForce(model.user_model)

index.index_from_dataset(
    movies_dataset.batch(128).map(
        lambda title: (title, model.movie_model(title))
    )
)

def recomendar(user_id: str, k: int = 5):
    _, titles = index(tf.constant([user_id]))
    print(f"\nRecomendaciones para usuario {user_id}:")
    for title in titles[0, :k].numpy():
        print(f"  - {title.decode('utf-8')}")

# Películas vistas por cada usuario (para contexto)
def peliculas_vistas(user_id: str, n: int = 3):
    vistas = df[df["user_id"] == user_id].nlargest(n, "rating")["title"].tolist()
    print(f"\nUsuario {user_id} calificó bien: {', '.join(vistas)}")

# Usuarios de ejemplo
for uid in ["42", "100", "7"]:
    peliculas_vistas(uid)
    recomendar(uid)

# Extraer embeddings de todas las películas
all_titles = tf.constant(movie_titles)
all_embeddings = model.movie_model(all_titles)

# Normalizar para similitud coseno
norms = tf.norm(all_embeddings, axis=1, keepdims=True)
normalized = all_embeddings / norms

def peliculas_similares(title: str, k: int = 5):
    # Buscar la película en el vocabulario
    idx = movie_titles.index(title)
    query = normalized[idx:idx+1]  # (1, 32)

    # Similitud coseno con todas las películas
    similarities = tf.squeeze(query @ tf.transpose(normalized))  # (N,)

    # Top-K excluyendo la propia película
    top_indices = tf.argsort(similarities, direction="DESCENDING").numpy()
    top_indices = [i for i in top_indices if i != idx][:k]

    print(f"\nPelículas similares a '{title}':")
    for i in top_indices:
        sim = similarities[i].numpy()
        print(f"  {movie_titles[i]:<45} (similitud: {sim:.2f})")

peliculas_similares("Star Wars (1977)")
peliculas_similares("Godfather, The (1972)")

pip install faiss-cpu

import faiss
import numpy as np

# Normalizar embeddings para similitud coseno
norms = np.linalg.norm(all_movie_embs, axis=1, keepdims=True)
normalized_embs = all_movie_embs / (norms + 1e-8)
user_norm = user_emb / (np.linalg.norm(user_emb) + 1e-8)

# --- FAISS IndexFlatIP (búsqueda exacta) ---
faiss_flat = faiss.IndexFlatIP(32)  # 32 = dimensión del embedding
faiss_flat.add(normalized_embs)
scores, indices = faiss_flat.search(user_norm.reshape(1, -1), 10)

# --- FAISS IndexIVFFlat (búsqueda aproximada) ---
nlist = 10  # número de clusters
quantizer = faiss.IndexFlatIP(32)
faiss_ivf = faiss.IndexIVFFlat(quantizer, 32, nlist, faiss.METRIC_INNER_PRODUCT)
faiss_ivf.train(normalized_embs)
faiss_ivf.add(normalized_embs)
faiss_ivf.nprobe = 3  # buscar en 3 de 10 clusters
scores, indices = faiss_ivf.search(user_norm.reshape(1, -1), 10)

# --- FAISS HNSW (grafo de vecinos) ---
faiss_hnsw = faiss.IndexHNSWFlat(32, 32, faiss.METRIC_INNER_PRODUCT)
faiss_hnsw.add(normalized_embs)
scores, indices = faiss_hnsw.search(user_norm.reshape(1, -1), 10)