En este artículo se explora la construcción de un sistema de recomendación de películas mediante un pipeline de Generación Aumentada por Recuperación (RAG). El objetivo es aprender a aprovechar el poder de las capacidades de búsqueda vectorial de MongoDB, transformar las descripciones de los datos en huellas digitales que se puedan buscar y crear un sistema que comprenda los matices de sus preferencias y su comunicación. En otras palabras, nuestro objetivo es construir un sistema de recomendación que no solo sea inteligente, sino también eficiente.
Al finalizar este artículo, habrás creado un sistema de recomendación de películas funcional. Este sistema podrá tomar la consulta de un usuario, como “Quiero ver una buena película de ciencia ficción que explore la inteligencia artificial” o “¿Qué película de animación sería buena para los adultos? ¿Por qué tu sugerencia es adecuada?” y devolver sugerencias de películas relevantes y el razonamiento de su elección.
Una secuencia de datos RAG se refiere al flujo secuencial de datos a través de una serie de pasos de procesamiento que combina las ventajas de los grandes modelos de lenguaje (LLM) con la recuperación de datos estructurados. Funciona recuperando primero la información relevante de una base de conocimiento y luego utilizando esta información para aumentar la entrada de un gran modelo de lenguaje, que genera el resultado final. El objetivo principal de una secuencia de datos de este tipo es generar respuestas más precisas, contextualmente apropiadas y personalizadas a las consultas específicas del usuario a partir de grandes bases de datos.
MongoDB es una base de datos NoSQL de código abierto que almacena datos en documentos flexibles, similares a JSON, lo que permite una fácil escalabilidad y manejo de diversos tipos y estructuras de datos. MongoDB desempeña un papel importante en este proyecto. Su modelo de documento se alinea bien con nuestros datos de películas, mientras que sus capacidades de búsqueda vectorial permiten búsquedas de similitud en nuestras incrustaciones (es decir, las representaciones numéricas del contenido de las películas). También podemos aprovechar las funciones de indexación y optimización de consultas para mantener una recuperación rápida de los datos incluso a medida que el conjunto de datos se expande.
Así es como se verá nuestro pipeline:
- Configurar el entorno y cargar datos de películas desde Hugging Face
- Modelar los datos utilizando Pydantic
- Generar incrustaciones para la información de las películas
- Ingerir los datos en una base de datos MongoDB
- Crear un índice de búsqueda vectorial en MongoDB Atlas
- Realizar operaciones de búsqueda de vectores para encontrar películas relevantes
- Gestionar consultas de usuarios con un modelo LLM
- Utilice el canal de RAG para obtener una recomendación de película
Paso 1: Configuración del entorno y carga del conjunto de datos
Primero, debemos importar las bibliotecas necesarias y configurar nuestro entorno. Esto también implica configurar nuestras claves API y la cadena de conexión que la aplicación utiliza para conectarse a una base de datos MongoDB:
import warnings
warnings.filterwarnings('ignore')import os
from dotenv import load_dotenv, find_dotenv
from datasets import load_dataset
import pandas as pd
from typing import List, Optional
from pydantic import BaseModel
from datetime import datetime
from pymongo.mongo_client import MongoClient
import openai
import time
_ = load_dotenv(find_dotenv())
MONGO_URI = os.environ.get("MONGO_URI")
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
openai.api_key = OPENAI_API_KEY
A continuación, cargamos nuestro conjunto de datos de películas:
dataset = load_dataset("Pablinho/movies-dataset", streaming=True, split="train")
dataset = dataset.take(200) # 200 movies for the sake of simplicity
dataset_df = pd.DataFrame(dataset)
El conjunto de datos contiene más de 9000 entradas. Sin embargo, para este ejercicio, estamos limitando nuestro conjunto de datos a 200 películas utilizando dataset.take(200)En un escenario del mundo real, probablemente utilizarías un conjunto de datos mucho más grande.
Paso 2: Modelado de datos con Pydantic
El modelado de datos es fundamental para garantizar la coherencia y la seguridad de tipos en nuestra aplicación. Por ello, utilizamos Pydantic para este fin:
class Movie(BaseModel):
Release_Date: Optional[str]
Title: str
Overview: str
Popularity: float
Vote_Count: int
Vote_Average: float
Original_Language: str
Genre: List[str]
Poster_Url: str
text_embeddings: List[float]
El uso de Pydantic ofrece varios beneficios, como la validación automática de datos, la verificación de tipos y la serialización/deserialización sencilla. Observe que también creamos un text_embeddings campo que almacenará nuestras incrustaciones generadas como una lista de flotantes
Paso 3: Generación de incrustaciones
Ahora, podemos usar la API de OpenAI y escribir una función para generar incrustaciones, de la siguiente manera:
def get_embedding(text):
if not text or not isinstance(text, str):
return None
try:
embedding = openai.embeddings.create(
input=text,
model="text-embedding-3-small", dimensions=1536).data[0].embedding
return embedding
except Exception as e:
print(f"Error in get_embedding: {e}")
return None
En las líneas de código anteriores, primero verificamos si la entrada es válida (cadena no vacía). Luego, usamos el método embeddings.create de OpenAI para generar la incrustación utilizando el modelo “text-embedding-3-small”, que genera incrustaciones de 1536 dimensiones.
Ahora podemos procesar cada registro y generar incrustaciones con la función anterior. También agregamos algunas líneas para procesar los 'Genre' campo, convirtiéndolo de una cadena (si existe) a una lista de géneros.
def process_and_embed_record(record):
for key, value in record.items():
if pd.isnull(value):
record[key] = Noneif record['Genre']:
record['Genre'] = record['Genre'].split(', ')
else:
record['Genre'] = []
text_to_embed = f"{record['Title']} {record['Overview']}"
embedding = get_embedding(text_to_embed)
record['text_embeddings'] = embedding
return record
records = [process_and_embed_record(record) for record in dataset_df.to_dict(orient='records')]
Estas incorporaciones nos permitirán realizar búsquedas semánticas más adelante y encontrar películas que sean conceptualmente similares a una consulta determinada. Tenga en cuenta que este proceso puede llevar algo de tiempo, especialmente para conjuntos de datos más grandes, ya que estamos realizando una llamada a la API para cada película.
Paso 4: Ingestión de datos en MongoDB
Establecemos una conexión a nuestra base de datos MongoDB:
def get_mongo_client(mongo_uri):
client = MongoClient(mongo_uri, appname="pmr.movie.python")
print("Connection to MongoDB successful")
return clientmongo_client = get_mongo_client(MONGO_URI)
database_name = "movies_dataset"
collection_name = "movies"
db = mongo_client.get_database(database_name)
collection = db.get_collection(collection_name)
collection.delete_many({})
Insertamos nuestros datos procesados e integrados en MongoDB, lo que nos permite almacenar y consultar de manera eficiente nuestros datos de películas, incluidas las incrustaciones de alta dimensión:
movies = [Movie(**record).dict() for record in records]
collection.insert_many(movies)
Paso 5: Creación de un índice de búsqueda vectorial en MongoDB Atlas
Antes de poder realizar operaciones de búsqueda de vectores, debemos crear un índice de búsqueda de vectores. Este paso se puede realizar directamente en la plataforma MongoDB Atlas:
- Inicie sesión en su cuenta de MongoDB Atlas
- Navegar a su clúster
- Vaya a la pestaña “Búsqueda y búsqueda de vectores”
- Haga clic en “Crear índice de búsqueda”
- Seleccione “Editor JSON” en la sección “Búsqueda de vectores de Atlas” y utilice la siguiente configuración:
{
"fields": [
{
"numDimensions": 1536,
"path": "text_embeddings",
"similarity": "cosine",
"type": "vector"
}
]
}
La idea es crear un índice de búsqueda vectorial llamado "vector_index_text" sobre el "text_embeddings" Utilizamos la similitud de coseno porque nos ayuda a encontrar películas con temas o contenidos similares comparando la dirección de sus vectores de inserción, ignorando las diferencias en longitud o cantidad de detalles, lo que es realmente bueno para hacer coincidir la consulta de un usuario con las descripciones de las películas.
Paso 6: Implementación de la búsqueda vectorial
Ahora, implementamos la función de búsqueda de vectores. La siguiente función está pensada para realizar una búsqueda de vectores en nuestra colección MongoDB. Primero genera una incrustación para la consulta del usuario. Luego, construye una canalización de agregación de MongoDB utilizando el operador $vectorSearch. La búsqueda busca los 20 vecinos más cercanos entre 150 candidatos.
def vector_search(user_query, db, collection, vector_index="vector_index_text", max_retries=3):
query_embedding = get_embedding(user_query)
if query_embedding is None:
return "Invalid query or embedding generation failed."vector_search_stage = {
"$vectorSearch": {
"index": vector_index,
"queryVector": query_embedding,
"path": "text_embeddings",
"numCandidates": 150,
"limit": 20
}
}
pipeline = [vector_search_stage]
for attempt in range(max_retries):
try:
results = list(collection.aggregate(pipeline))
if results:
explain_query_execution = db.command(
'explain', {
'aggregate': collection.name,
'pipeline': pipeline,
'cursor': {}
},
verbosity='executionStats')
vector_search_explain = explain_query_execution['stages'][0]['$vectorSearch']
millis_elapsed = vector_search_explain['explain']['collectStats']['millisElapsed']
print(f"Total time for the execution to complete on the database server: {millis_elapsed} milliseconds")
return results
else:
print(f"No results found on attempt {attempt + 1}. Retrying...")
time.sleep(2)
except Exception as e:
print(f"Error on attempt {attempt + 1}: {str(e)}")
time.sleep(2)
return "Failed to retrieve results after multiple attempts."
Implementamos un mecanismo de reintento (hasta 3 intentos) para manejar posibles problemas transitorios. La función ejecuta la explain comando también, que proporciona información detallada sobre la ejecución de la consulta.
Paso 7: Manejo de consultas de usuarios con un LLM
Finalmente, podemos gestionar las consultas de los usuarios. Primero, definimos un SearchResultItem clase para estructurar nuestros resultados de búsqueda. Luego, el handle_user_query La función une todo: realiza una búsqueda vectorial basada en la consulta del usuario, formatea los resultados de la búsqueda en un DataFrame de pandas y luego utiliza el modelo GPT de OpenAI (es decir, gpt-3.5-turbo) para generar una respuesta basada en los resultados de la búsqueda y la consulta del usuario, y muestra los resultados y la respuesta generada:
class SearchResultItem(BaseModel):
Title: str
Overview: str
Genre: List[str]
Vote_Average: float
Popularity: floatdef handle_user_query(query, db, collection):
get_knowledge = vector_search(query, db, collection)
if isinstance(get_knowledge, str):
return get_knowledge, "No source information available."
search_results_models = [SearchResultItem(**result) for result in get_knowledge]
search_results_df = pd.DataFrame([item.dict() for item in search_results_models])
completion = openai.chat.completions.create(
model="gpt-3.5-turbo",
messages=[
{"role": "system", "content": "You are a movie recommendation system."},
{"role": "user", "content": f"Answer this user query: {query} with the following context:\n{search_results_df}"}
]
)
system_response = completion.choices[0].message.content
print(f"- User Question:\n{query}\n")
print(f"- System Response:\n{system_response}\n")
return system_response
Esta función realmente demuestra el valor central de este RAG: generamos una respuesta contextualmente apropiada recuperando información relevante de nuestra base de datos.
8. Uso de la tubería RAG
Para utilizar esta canalización RAG, ahora puede realizar consultas como esta:
query = """
I'm in the mood for a highly-rated action movie. Can you recommend something popular?
Include a reason for your recommendation.
"""
handle_user_query(query, db, collection)
El sistema daría una respuesta similar a esta:
I recommend "Spider-Man: No Way Home" as a popular and highly-rated action
movie for you to watch. With a vote average of 8.3 and a popularity score
of 5083.954, this film has garnered a lot of attention and positive
reviews from audiences. "Spider-Man: No Way Home" is a thrilling action-packed movie that brings
together multiple iterations of Spider-Man in an epic crossover event. It
offers a blend of intense action sequences, emotional depth, and nostalgic
moments that fans of the superhero genre will surely enjoy. So, if you're
in the mood for an exciting action movie with a compelling storyline and
fantastic visual effects, "Spider-Man: No Way Home" is an excellent choice
for your movie night.
La creación de una canalización RAG implica varios pasos, desde la carga y el modelado de datos hasta la generación de incrustaciones y la búsqueda de vectores. Este ejemplo muestra cómo una canalización RAG puede proporcionar respuestas informativas y contextuales mediante la combinación de los datos de películas específicos de nuestra base de datos con las capacidades de comprensión y generación de lenguaje natural del modelo de lenguaje. Además de esto, utilizamos MongoDB porque es ideal para este tipo de flujo de trabajo debido a sus capacidades nativas de búsqueda de vectores, su modelo de documento flexible y su escalabilidad.
Puede ampliar este sistema agregando más datos, ajustando sus integraciones o implementando algoritmos de recomendación más complejos.
Para obtener el código completo y recursos adicionales, consulte Repositorio de GitHubEl conjunto de datos utilizado en este proyecto proviene de Kaggle y el autor original le ha otorgado la licencia CC0 1.0 Universal (CC0 1.0) Public Domain Dedication. Puede encontrar el conjunto de datos y más información aquí.