Sistema de recomendación de películas con Gemini y Qdrant
Copyright 2026 Google LLC.
# @title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
Nota: Este notebook requiere límites de tasa de nivel de pago para ejecutarse correctamente. (Consulta precios para más detalles).
Descripción general
La API de Gemini proporciona acceso a una familia de modelos de IA generativa para crear contenido y resolver problemas. Estos modelos están diseñados y entrenados para manejar texto e imágenes como entrada.
Qdrant es un motor de búsqueda de similitud vectorial de código abierto diseñado para una búsqueda semántica eficiente y escalable. Ofrece una API simple pero potente para almacenar y buscar vectores de alta dimensión, admite el filtrado con metadatos (payloads) y se integra fácilmente en sistemas de producción. Qdrant se puede autoalojar o acceder a través de su servicio en la nube administrado, lo que lo hace rápido de configurar e ideal para una amplia gama de aplicaciones de IA que dependen de la comprensión y recuperación semántica.
En este notebook, aprenderás a realizar una búsqueda de similitud en datos de un sitio web con la ayuda de la API de Gemini y Qdrant.
from google.colab import userdata
from google import genai
GEMINI_API_KEY=userdata.get('GEMINI_API_KEY')
client = genai.Client(api_key=GEMINI_API_KEY)
Construyendo el índice vectorial de películas
Esta sección cubre la preparación del conjunto de datos de películas, la generación de embeddings usando Gemini y su indexación en Qdrant para la búsqueda de similitud.
1. Carga el conjunto de datos de Kaggle
Comienza cargando el conjunto de datos de Kaggle usando la biblioteca kagglehub. El conjunto de datos utilizado en este notebook es el TMDB Movie Dataset 2024, que contiene aproximadamente más de 1 millón de entradas de películas.
/tmp/ipykernel_26130/2431045845.py:6: DeprecationWarning: load_dataset is deprecated and will be removed in a future version.
df = kagglehub.load_dataset(
2. Inspecciona la estructura del conjunto de datos
Dado que el conjunto de datos es grande, inspeccionarlo te ayuda a identificar campos útiles y a filtrar datos irrelevantes desde el principio.
print("\nDataset Columns:")
print(df.columns)
print("\nMissing Values per Column:")
print(df.isnull().sum())
print(f"\nNumber of rows: {len(df)}")
print(f"Number of unique IDs: {df['id'].nunique()}")
Este paso filtra el conjunto de datos para mantener solo los metadatos útiles para la búsqueda semántica: id, title, overview, genres, keywords, tagline y release_date. Estos campos proporcionan suficiente contexto para generar embeddings significativos.
Se eliminan las entradas (filas) a las que les falta un title o que carecen tanto de overview como de genres, ya que no tienen suficientes datos descriptivos para recomendaciones precisas.
import pandas as pd
import numpy as np
import ast
print(f"Original rows: {len(df)}")
columns_to_keep = ['id', 'title', 'overview', 'genres', 'keywords', 'tagline', 'release_date']
df_relevant = df[columns_to_keep].copy()
print(f"Rows before dropping missing title: {len(df_relevant)}")
df_relevant.dropna(subset=['title'], inplace=True)
df_relevant = df_relevant[~(df_relevant['genres'].isna() & df_relevant['overview'].isna())]
print(f"Rows after dropping missing title and dropping missing (genres and overview): {len(df_relevant)}")
# Fill missing text columns with empty strings
text_cols_to_fill = ['overview', 'genres', 'keywords', 'tagline']
for col in text_cols_to_fill:
df_relevant[col] = df_relevant[col].fillna('')
# Extract release year from the release_date string
def get_year(date_str):
if pd.isna(date_str) or not isinstance(date_str, str) or len(date_str) < 4:
return None
try:
return int(date_str[:4])
except (ValueError, TypeError):
return None
df_relevant['release_year'] = df_relevant['release_date'].apply(get_year)
print("\nSample data after cleaning (keeping missing overviews):")
print(df_relevant[['id', 'title', 'overview', 'genres', 'keywords', 'tagline', 'release_year']].head())
Original rows: 1254611
Rows before dropping missing title: 1254611
Rows after dropping missing title and dropping missing (genres and overview): 1109811
Sample data after cleaning (keeping missing overviews):
id title overview \
0 27205 Inception Cobb, a skilled thief who commits corporate es...
1 157336 Interstellar The adventures of a group of explorers who mak...
2 155 The Dark Knight Batman raises the stakes in his war on crime. ...
3 19995 Avatar In the 22nd century, a paraplegic Marine is di...
4 24428 The Avengers When an unexpected enemy emerges and threatens...
genres \
0 Action, Science Fiction, Adventure
1 Adventure, Drama, Science Fiction
2 Drama, Action, Crime, Thriller
3 Action, Adventure, Fantasy, Science Fiction
4 Science Fiction, Action, Adventure
keywords \
0 rescue, mission, dream, airplane, paris, franc...
1 rescue, future, spacecraft, race against time,...
2 joker, sadism, chaos, secret identity, crime f...
3 future, society, culture clash, space travel, ...
4 new york city, superhero, shield, based on com...
tagline release_year
0 Your mind is the scene of the crime. 2010.0
1 Mankind was born on Earth. It was never meant ... 2014.0
2 Welcome to a world without rules. 2008.0
3 Enter the world of Pandora. 2009.0
4 Some assembly required. 2012.0
4. Prepara el texto para embedding
Este paso prepara los metadatos de la película para embedding combinando campos relevantes en una única cadena de texto estructurada. Esta representación incluye el título, la sinopsis, los géneros, las palabras clave, el eslogan y el año de lanzamiento (si está disponible). La salida se almacena en una nueva columna llamada text_for_embedding.
Los embeddings son representaciones vectoriales numéricas de texto que capturan el significado semántico y las relaciones. Estos vectores se pueden usar para tareas como la búsqueda de similitud y la agrupación.
Aprende más sobre los embeddings de texto y explora el notebook de embeddings de Gemini.
def create_embedding_text(row):
"""Combines available movie metadata into a single string for embedding."""
# Title is always present, so it can be included directly
title_str = f"Title: {row['title']}"
overview_str = f"Overview: {row['overview']}" if row['overview'] else ""
year_str = f"Release Year: {int(row['release_year'])}" if pd.notna(row['release_year']) else ""
genre_str = f"Genres: {row['genres']}" if row['genres'] else ""
keywords_str = f"Keywords: {row['keywords']}" if row['keywords'] else ""
tagline_str = f"Tagline: {row['tagline']}" if row['tagline'] else ""
parts = [
title_str,
overview_str,
year_str,
genre_str,
keywords_str,
tagline_str
]
return "\n".join(part for part in parts if part)
df_relevant['text_for_embedding'] = df_relevant.apply(create_embedding_text, axis=1)
# Use this to inspect how movie data has been transformed for embedding
print(df_relevant[['id', 'title', 'text_for_embedding']].head())
id title text_for_embedding
0 27205 Inception Title: Inception\nOverview: Cobb, a skilled th...
1 157336 Interstellar Title: Interstellar\nOverview: The adventures ...
2 155 The Dark Knight Title: The Dark Knight\nOverview: Batman raise...
3 19995 Avatar Title: Avatar\nOverview: In the 22nd century, ...
4 24428 The Avengers Title: The Avengers\nOverview: When an unexpec...
5. Muestra un subconjunto para desarrollo
Para que el notebook sea fácil de ejecutar y garantizar un desarrollo eficiente, querrás iterar rápidamente y minimizar el uso de recursos. En lugar de usar el conjunto de datos completo, este paso muestrea 5,000 películas de los datos limpios, a menos que el conjunto de datos ya sea más pequeño, en cuyo caso se usan todas las entradas.
SAMPLE_SIZE = 5000
if len(df_relevant) > SAMPLE_SIZE:
print(f"\nTaking a random sample of {SAMPLE_SIZE} movies for development.")
df_sample = df_relevant.sample(n=SAMPLE_SIZE, random_state=42)
else:
print(f"\nCleaned dataset size ({len(df_relevant)}) is smaller than or equal to SAMPLE_SIZE. Using the full cleaned dataset.")
df_sample = df_relevant
print(f"Working with {len(df_sample)} movies for the next steps.")
print(df_sample[['id', 'title', 'release_year']].head())
columns_for_payload = ['title', 'overview', 'genres', 'keywords', 'tagline', 'release_year']
columns_final = ['id', 'text_for_embedding'] + columns_for_payload
df_sample = df_sample[columns_final]
print("\nFinal sample DataFrame structure for embedding/indexing:")
print(df_sample.info())
Taking a random sample of 5000 movies for development.
Working with 5000 movies for the next steps.
id title release_year
95090 714945 Gather 2020.0
198838 117602 California Girls 1985.0
827619 648297 After Jake 2013.0
801274 637062 Duas vezes Senzala 2017.0
846938 342211 Rise NaN
Final sample DataFrame structure for embedding/indexing:
<class 'pandas.core.frame.DataFrame'>
Index: 5000 entries, 95090 to 650934
Data columns (total 8 columns):
# Column Non-Null Count Dtype
--- ------ -------------- -----
0 id 5000 non-null int64
1 text_for_embedding 5000 non-null object
2 title 5000 non-null object
3 overview 5000 non-null object
4 genres 5000 non-null object
5 keywords 5000 non-null object
6 tagline 5000 non-null object
7 release_year 4260 non-null float64
dtypes: float64(1), int64(1), object(6)
memory usage: 351.6+ KB
None
6. Inicializa Qdrant para la indexación vectorial
Con los datos preparados, el siguiente paso es configurar Qdrant, un motor de búsqueda de similitud vectorial optimizado para almacenar y consultar vectores de alta dimensión. Admite una indexación, filtrado y búsqueda de similitud rápidos en millones de vectores.
Qdrant puede ejecutarse:
Localmente como un servicio independiente
En la nube para implementaciones de producción
O completamente en memoria para un uso rápido y temporal durante el desarrollo
En este notebook, Qdrant se inicializa usando el modo en memoria pasando ":memory:" al cliente. Esto almacena datos solo en la RAM, lo que significa que no persistirá después de que termine la sesión. Esto es adecuado para la experimentación, pero no para guardar resultados a largo plazo.
También configuras lo siguiente:
COLLECTION_NAME: El nombre de la colección de Qdrant para almacenar vectores de películas
VECTOR_SIZE: Establecido en 3072 para que coincida con la dimensionalidad de los embeddings de texto generados por Gemini
DISTANCE_METRIC: Establecido en distancia coseno, que es ideal para medir la similitud semántica entre vectores de embedding
from qdrant_client import QdrantClient, models
import time
COLLECTION_NAME = "tmdb_movies_sample"
VECTOR_SIZE = 3072
DISTANCE_METRIC = models.Distance.COSINE
# Initialize Qdrant client using in-memory storage
qdrant_client = QdrantClient(":memory:")
7. Define la función de embedding por lotes
Este paso define la función get_embeddings_batch, que genera embeddings de texto para lotes de datos de películas utilizando el modelo de embedding de Gemini (gemini-embedding-001), incluyendo reintentos automáticos para mayor robustez.
import time
from google.api_core import exceptions, retry
MODEL_FOR_EMBEDDING = "gemini-embedding-001" # @param ["gemini-embedding-001", "gemini-embedding-2-preview"] {"allow-input":true, isTemplate: true}
BATCH_SIZE = 25
QDRANT_BATCH_SIZE = 3072
@retry.Retry(timeout=3000)
def get_embeddings_batch(texts: list[str], task_type="RETRIEVAL_DOCUMENT") -> list[list[float]] | None:
"""
Generates embeddings for a batch of texts using Gemini API with retry.
Args:
texts: A list of strings to embed.
task_type: The task type for the embedding model.
Returns:
A list of embedding vectors (list of floats), or None if a non-retryable error occurs.
"""
if not texts:
return []
try:
response = client.models.embed_content(
model=MODEL_FOR_EMBEDDING,
contents=texts,
config={
"task_type":task_type,
}
)
return response.embeddings
except exceptions.RetryError as e:
print(f"Embedding batch failed after retries: {e}")
return None
except Exception as e:
print(f"An unexpected error occurred during embedding: {e}")
return None
# Example of what an embedding looks like
sample_embedding = get_embeddings_batch(["Example movie about space and survival"])[0]
print("Example embedding vector:", sample_embedding.values[:10])
Una colección en Qdrant es como una tabla en una base de datos, almacena vectores junto con metadatos opcionales (payload). Cada colección tiene su propia configuración, incluyendo el tamaño del vector y la métrica de similitud.
# In case someone tries running the whole notebook again they would want to create the collection again
try:
qdrant_client.delete_collection(collection_name=COLLECTION_NAME)
print(f"Existing collection '{COLLECTION_NAME}' deleted.")
except Exception as e:
print(f"Error deleting collection (it might not exist): {e}")
try:
qdrant_client.create_collection(
collection_name=COLLECTION_NAME,
vectors_config=models.VectorParams(
size=VECTOR_SIZE,
distance=DISTANCE_METRIC
)
)
print(f"Collection '{COLLECTION_NAME}' created successfully.")
except Exception as e:
print(f"Error creating collection: {e}")
Existing collection 'tmdb_movies_sample' deleted.
Collection 'tmdb_movies_sample' created successfully.
9. Crea payloads para el almacenamiento de metadatos
En Qdrant, además de almacenar embeddings vectoriales, puedes adjuntar información adicional llamada payload a cada vector. Estos metadatos ayudan a filtrar o recuperar resultados relevantes basados en atributos como el título, los géneros o el año de lanzamiento.
La función create_payload prepara el payload extrayendo columnas específicas de cada registro de película, manejando los valores faltantes y asegurando que los tipos de datos sean compatibles con Qdrant.
payload_columns = [
'title', 'overview', 'genres', 'keywords', 'tagline', 'release_year'
]
def create_payload(row, payload_columns):
payload = {}
for col in payload_columns:
value = row[col]
if pd.isna(value):
payload[col] = None
elif isinstance(value, (np.int64, np.int32)):
payload[col] = int(value)
elif isinstance(value, (np.float64, np.float32)):
payload[col] = float(value)
else:
payload[col] = value
return payload
10. Embedding por lotes e indexación en Qdrant
Este paso procesa el conjunto de datos de películas muestreadas en lotes para generar embeddings vectoriales utilizando la API de Gemini y cargar (upsert) estos embeddings junto con sus payloads de metadatos en la colección de Qdrant.
Puntos clave de este proceso:
El conjunto de datos se divide en lotes de tamaño BATCH_SIZE para la generación de embeddings para mantenerse dentro de los límites de la API.
Los datos de texto de cada lote se envían a la API de embedding de Gemini con reintentos manejados en la función de embedding.
Para cada lote incrustado con éxito, el código prepara puntos (cada uno contiene un ID, un embedding vectorial y un payload de metadatos) para ser cargados en Qdrant.
Los puntos se almacenan en búfer y se cargan en fragmentos de tamaño QDRANT_BATCH_SIZE para optimizar el rendimiento.
El proceso incluye manejo de errores y lógica de reintento para evitar que las fallas detengan toda la operación.
Al final, se cargan los puntos restantes en el búfer.
Se imprimen estadísticas resumidas de los elementos procesados, fallidos y cargados con éxito.
from tqdm import tqdm
print(f"Starting batch embedding and indexing process for {len(df_sample)} movies...")
print(f"Using Gemini Batch Size: {BATCH_SIZE}, Qdrant Upsert Batch Size: {QDRANT_BATCH_SIZE}")
points_to_upsert_buffer = []
total_processed = 0
total_failed_embedding = 0
total_upserted = 0
num_batches = (len(df_sample) + BATCH_SIZE - 1) // BATCH_SIZE
for i in tqdm(range(0, len(df_sample), BATCH_SIZE), total=num_batches, desc="Processing Batches"):
batch_df = df_sample.iloc[i : i + BATCH_SIZE]
batch_texts = batch_df['text_for_embedding'].tolist()
batch_ids = batch_df['id'].tolist()
if not batch_texts:
continue
# Generate embeddings for the current batch of movie texts
batch_embeddings = get_embeddings_batch(batch_texts, task_type="RETRIEVAL_DOCUMENT")
# Check if embeddings were successfully generated and correspond to batch size
if batch_embeddings and len(batch_embeddings) == len(batch_texts):
for j in range(len(batch_ids)):
item_id = batch_ids[j]
item_embedding = batch_embeddings[j]
row_data = batch_df.iloc[j]
# Prepare metadata payload for this movie
payload = create_payload(row_data, payload_columns)
# Create a Qdrant PointStruct with id, embedding vector, and payload
point = models.PointStruct(
id=int(item_id),
vector=item_embedding.values,
payload=payload
)
points_to_upsert_buffer.append(point)
total_processed += len(batch_ids)
else:
print(f"Failed to get embeddings for batch starting at index {i}. Skipping {len(batch_ids)} items.")
total_failed_embedding += len(batch_ids)
continue
# Upload buffered points to Qdrant if buffer reached batch size or end of data
if len(points_to_upsert_buffer) >= QDRANT_BATCH_SIZE or (i + BATCH_SIZE >= len(df_sample)):
if points_to_upsert_buffer:
try:
qdrant_client.upsert(
collection_name=COLLECTION_NAME,
points=points_to_upsert_buffer,
wait=False
)
total_upserted += len(points_to_upsert_buffer)
points_to_upsert_buffer = []
except Exception as e:
print(f"Error upserting chunk to Qdrant: {e}")
points_to_upsert_buffer = []
time.sleep(5)
# Pause before retrying to avoid hammering the service after an error
# Upload any remaining points left in buffer after loop completion
if points_to_upsert_buffer:
print(f"Upserting final remaining chunk of {len(points_to_upsert_buffer)} points.")
try:
qdrant_client.upsert(
collection_name=COLLECTION_NAME,
points=points_to_upsert_buffer,
wait=True
)
total_upserted += len(points_to_upsert_buffer)
points_to_upsert_buffer = []
except Exception as e:
print(f"Error upserting final chunk: {e}")
print("Batch embedding and indexing finished.")
print(f"Total items processed (attempted embedding): {total_processed}")
print(f"Total points successfully prepared for upsert: {total_upserted}")
Starting batch embedding and indexing process for 5000 movies...
Using Gemini Batch Size: 25, Qdrant Upsert Batch Size: 3072
Batch embedding and indexing finished.
Total items processed (attempted embedding): 5000
Total points successfully prepared for upsert: 5000
# Waiting for collection to settle
time.sleep(5)
try:
count = qdrant_client.count(collection_name=COLLECTION_NAME, exact=True)
print(f"\nVerification: Collection '{COLLECTION_NAME}' now contains {count.count} points.") # it should print 5000
except Exception as e:
print(f"Error verifying collection count: {e}")
Verification: Collection 'tmdb_movies_sample' now contains 5000 points.
11. Busca y recomienda películas similares usando embeddings vectoriales
Con todos los vectores de películas indexados en Qdrant, ahora puedes realizar búsquedas semánticas. Esto te permite tomar cualquier consulta de usuario (como una frase, descripción de película o concepto), convertirla en un embedding usando el mismo modelo de Gemini y recuperar los vectores de películas más similares de la colección usando la similitud del coseno.
Esta función recommend_movies demuestra cómo:
Generar un embedding a partir de tu consulta de entrada usando la API de Gemini.
Realizar una búsqueda de similitud usando el método search() de Qdrant.
Recuperar las k entradas de películas más similares, incluyendo sus metadatos y puntuaciones de similitud.
Este es el paso final donde la base de datos vectorial funciona como un motor de recomendación.
def recommend_movies(query_text, top_k=5):
"""
Finds movies similar to the query_text using the Qdrant index.
Args:
query_text (str): The user's query (e.g., movie title, description, theme).
top_k (int): The maximum number of recommendations to return.
Returns:
list: A list of dictionaries, where each dictionary contains the
payload (movie details) and similarity score of a recommended movie.
Returns an empty list if query embedding fails or no results found.
"""
print(f"Searching for recommendations based on: '{query_text}'")
# Generate embedding for the user query using Gemini
query_embedding = get_embeddings_batch(query_text, task_type="RETRIEVAL_QUERY")[0].values
print(f"Query embedding: {query_embedding}")
if query_embedding is None:
print("Error: Could not generate embedding for the query.")
return []
try:
# Perform a semantic search on Qdrant using the query vector
search_result = qdrant_client.search(
collection_name=COLLECTION_NAME,
query_vector=query_embedding,
limit=top_k,
with_payload=True
)
recommendations = []
if search_result:
print(f"Found {len(search_result)} potential recommendations:")
for hit in search_result:
recommendation = {
"id": hit.id,
"score": hit.score,
"payload": hit.payload
}
recommendations.append(recommendation)
else:
print("No recommendations found matching the query.")
return recommendations
except Exception as e:
print(f"Error during Qdrant search: {e}")
return []
Prueba tu recomendador de películas
Ahora puedes probar tu recomendador de películas describiendo un tema, género o concepto en lenguaje natural. El sistema devolverá las películas semánticamente más similares de tu conjunto de datos basándose en la búsqueda de similitud vectorial utilizando embeddings generados por Gemini.
query = """
I want to watch something with my girlfriends that’s
both funny and teaches something.
"""
recommendations = recommend_movies(query, top_k=5)
if recommendations:
print("\n--- Recommendations ---")
for rec in recommendations:
print(f" - Score: {rec['score']:.4f}")
print(f" Title: {rec['payload'].get('title', 'N/A')}")
print(f" Genre: {rec['payload'].get('genres', 'N/A')}")
print(f" Year: {rec['payload'].get('release_year', 'N/A')}")
print("-" * 10)
/tmp/ipykernel_26130/3969524323.py:26: DeprecationWarning: `search` method is deprecated and will be removed in the future. Use `query_points` instead.
search_result = qdrant_client.search(
Próximos pasos
Este notebook demostró cómo construir un sistema de recomendación de películas combinando las capacidades de embedding de la API de Gemini con la búsqueda vectorial de Qdrant.
Referencias útiles de la API
Para una comprensión más detallada y para explorar funciones avanzadas, consulta la documentación oficial:
Documentación de Embeddings de la API de Gemini: Aprende a generar embeddings de texto, comprender los parámetros del modelo y usarlos eficazmente para tareas de búsqueda semántica y similitud.
Documentación del cliente Python de Qdrant: Comprende cómo administrar colecciones, insertar y buscar vectores, configurar la indexación e interactuar con la base de datos vectorial.
Ejemplos relacionados
Para explorar más casos de uso y obtener inspiración adicional, consulta estos notebooks relacionados en este directorio:
Descripción general del SDK de Google GenAI: Te guía a través de la instalación y configuración del SDK, la creación de prompts de texto y multimodales, el conteo de tokens, los filtros de seguridad, el chat de múltiples turnos, la llamada a funciones, la carga de archivos, el almacenamiento en caché de contexto y más.
Embeddings de texto con la API de Gemini: Se centra en la generación y el trabajo con embeddings de texto utilizando la API de Gemini, ideal para construir sistemas de búsqueda y recomendación basados en vectores.
Lección del curso «Gemini API Cookbook (examples)» de Google, publicado con licencia Apache 2.0. Traducción y adaptación al español de IA con Clase. IA con Clase no está afiliado a Google. Ver el original · Licencia
Esta lección es gratuita. El resto del curso se abre con la Membresía de IA con Clase, que incluye todos los cursos del catálogo. Ver precios
