Uso de una base de datos vectorial para buscar discursos de la Casa Blanca
Este artículo se publicó originalmente en The New Stack.
Se acerca la temporada de campaña para la campaña presidencial de EE. UU. Es un buen momento para echar la vista atrás a algunos de los discursos pronunciados por la administración Biden durante sus primeros dos años en el cargo. ¿No sería genial buscar en algunas transcripciones de discursos para saber más sobre el mensaje de la Casa Blanca en torno a determinados temas hasta ahora?
Digamos que queremos buscar el contenido de un discurso. ¿Cómo lo haríamos? Podríamos usar búsqueda semántica. La búsqueda semántica es uno de los temas más candentes en inteligencia artificial (IA) en este momento. Se ha vuelto más importante a medida que hemos visto el aumento de popularidad de aplicaciones de procesamiento del lenguaje natural (NLP) como ChatGPT. En lugar de hacer ping repetidamente a GPT, lo cual es costoso tanto económica como ecológicamente, podemos usar una base de datos vectorial para almacenar en caché los resultados (como con GPTCache).
En este tutorial, pondremos en marcha una base de datos vectorial localmente para poder buscar en los discursos de Biden de 2021 a 2022 por contenido. El conjunto de datos que usamos es el conjunto de datos “The White House (Speeches and Remarks) 12/10/2022”, que encontramos en Kaggle y pusimos a disposición para su descarga mediante Google Drive para este ejemplo. Hay un cuaderno de recorrido de este tutorial disponible en GitHub.
Antes de sumergirnos en el código, asegúrate de descargar los requisitos previos. Necesitamos cuatro bibliotecas: PyMilvus, Milvus, Sentence-Transformers y gdown. Puedes obtener las bibliotecas necesarias desde PyPi ejecutando: pip3 install pymilvus==2.2.5 sentence-transformers gdown milvus.
Preparación del conjunto de datos de discursos de la Casa Blanca
Como en casi cualquier proyecto de IA/ML basado en conjuntos de datos del mundo real, primero necesitamos preparar los datos. Usamos gdown para descargar el conjunto de datos y zipfile para extraerlo en una carpeta local. Después de ejecutar el código siguiente, esperamos ver un archivo titulado “The white house speeches.csv” en una carpeta titulada “white_house_2021_2022”.
import gdown
url = 'https://drive.google.com/uc?id=10_sVL0UmEog7mczLedK5s1pnlDOz3Ukf'
output = './white_house_2021_2022.zip'
gdown.download(url, output)
import zipfile
with zipfile.ZipFile("./white_house_2021_2022.zip","r") as zip_ref:
zip_ref.extractall("./white_house_2021_2022")
Usamos pandas para cargar e inspeccionar los datos CSV.
import pandas as pd
df = pd.read_csv("./white_house_2021_2022/The white house speeches.csv")
df.head()
Cuando echamos un vistazo al head de los datos, ¿qué notas? Lo primero que noto es que los datos tienen cuatro columnas: una columna de título, fecha y hora, ubicación y discurso. Lo segundo es que hay valores nulos. Los valores nulos no siempre son un problema, pero sí lo son para nuestros datos.
Limpieza del conjunto de datos
Los discursos sin ningún contenido sustancial (valores nulos en la columna “Speech”) son completamente inútiles para nosotros. Eliminemos nuestros valores nulos y volvamos a examinar los datos.
df = df.dropna()
df
Ahora vemos que en realidad hay un segundo problema que no era evidente de inmediato al mirar solo el head de los datos. Si miras la última entrada, verás que esta entrada es solo una hora. “12:18 P.M. EST” difícilmente es un discurso. No tiene sentido guardar esta entrada. No podemos obtener ningún valor al guardar una incrustación vectorial.
Deshagámonos de todos los discursos que tienen menos de cierta longitud. Para este ejemplo, he elegido 50, pero puedes elegir lo que tenga sentido para ti. Elegí 50 explorando muchos números diferentes. Si buscas transcripciones de discursos de entre 20 y 50 caracteres, verás que muchas son ubicaciones u horas con algunas frases aleatorias incluidas.
cleaned_df = df.loc[(df["Speech"].str.len() > 50)]cleaned_df
Una vez resueltos los discursos cortos y sin sustancia, volvemos a mirar nuestros datos para ver un problema más. Muchos de los discursos contienen valores \r\n: saltos de línea y retornos. Estos caracteres se usan para el formato, pero no contienen ningún valor semántico. El siguiente paso en nuestro proceso de limpieza de datos es deshacernos de ellos.
cleaned_df["Speech"] = cleaned_df["Speech"].str.replace("\r\n", "")
cleaned_df
Eso se ve mucho mejor. El paso final es convertir la columna “Date_time” a un formato mejor para almacenarla en nuestra base de datos vectorial y compararla con otras fechas y horas. Usamos la biblioteca datetime para convertir simplemente este formato de fecha y hora en un formato universal YYYY-MM-DD.
import datetime
# Convert the 'date' column to datetime objects
cleaned_df["Date_time"] = pd.to_datetime(cleaned_df["Date_time"], format="%B %d, %Y")
cleaned_df
Configuración de una base de datos vectorial para búsqueda semántica
Nuestros datos ya están limpios y listos para trabajar con ellos. El siguiente paso es poner en marcha una base de datos vectorial para buscar realmente los discursos por su contenido. Para este ejemplo, usamos Milvus Lite, una versión ligera de Milvus que puedes poner en funcionamiento sin Docker, Kubernetes ni tener que lidiar con ningún tipo de archivo YAML.
Lo primero que hacemos es definir algunas de nuestras constantes. Necesitamos un nombre de colección (para la base de datos vectorial), el número de dimensiones de nuestro vector embebido, un tamaño de lote y un número que defina cuántos resultados queremos obtener cuando buscamos. Este ejemplo usa el transformador de oraciones MiniLM L6 v2, que produce vectores de embeddings de 384 dimensiones.
COLLECTION_NAME = "white_house_2021_2022"
DIMENSION = 384
BATCH_SIZE = 128
TOPK = 3
Usamos el default_server de Milvus. Luego, usamos el SDK de PyMilvus para conectarnos a nuestro servidor local de Milvus. Si hay una colección en nuestra base de datos vectorial que tiene el mismo nombre que el nombre de colección que definimos antes, eliminamos esa colección para asegurarnos de comenzar con una pizarra en blanco.
from milvus import default_server
from pymilvus import connections, utility
default_server.start()
connections.connect(host="127.0.0.1", port=default_server.listen_port)
if utility.has_collection(COLLECTION_NAME):
utility.drop_collection(COLLECTION_NAME)
Como en la mayoría de las otras bases de datos, necesitamos un esquema para cargar datos en la base de datos vectorial Milvus. Primero, definimos los campos de datos que queremos que tenga cada objeto. Menos mal que miramos los datos antes. Usamos cinco campos de datos: las cuatro columnas que teníamos antes y una columna de ID. Excepto que esta vez usamos el embedding vectorial del discurso en lugar del texto real.
from pymilvus import FieldSchema, CollectionSchema, DataType, Collection
# object should be inserted in the format of (title, date, location, speech embedding)
fields = [
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=500),
FieldSchema(name="date", dtype=DataType.VARCHAR, max_length=100),
FieldSchema(name="location", dtype=DataType.VARCHAR, max_length=200),
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=DIMENSION)
]
schema = CollectionSchema(fields=fields)
collection = Collection(name=COLLECTION_NAME, schema=schema)
Lo último que necesitamos definir antes de estar listos para cargar datos en la base de datos vectorial es el índice. Hay muchos índices y patrones vectoriales, pero para este ejemplo usamos el índice IVF_FLAT con 128 clústeres. Las aplicaciones más grandes suelen usar más de 128 clústeres, pero de todos modos solo tenemos un poco más de 600 entradas. Para nuestra distancia, medimos usando la norma L2. Una vez que definimos los parámetros de nuestro índice, creamos el índice en nuestra colección y lo cargamos para su uso.
index_params = {
"index_type": "IVF_FLAT",
"metric_type": "L2",
"params": {"nlist": 128},
}
collection.create_index(field_name="embedding", index_params=index_params)
collection.load()
Obtener embeddings vectoriales de discursos
Gran parte de lo que hemos repasado hasta ahora se aplica al trabajar con casi cualquier base de datos. Limpiamos algunos datos, pusimos en marcha una instancia de base de datos y definimos un esquema para nuestra base de datos. Además de definir un índice, otra cosa que debemos hacer para las bases de datos vectoriales en particular es obtener los embeddings.
Primero, obtenemos el modelo de transformador de oraciones MiniLM L6 v2 como se mencionó anteriormente. Luego creamos una función que realiza una transformación en los datos y los inserta en la colección. Esta función toma un lote de datos, obtiene los embeddings para las transcripciones de los discursos, crea un objeto para insertar y lo inserta en la colección.
Para dar contexto, esta función realiza una actualización por lotes. En este ejemplo, estamos insertando por lotes 128 entradas a la vez. La única transformación de datos que hacemos en nuestra inserción es convertir el texto del discurso en un embedding.
from sentence_transformers import SentenceTransformer
transformer = SentenceTransformer('all-MiniLM-L6-v2')
# expects a list of (title, date, location, speech)
def embed_insert(data: list):
embeddings = transformer.encode(data[3])
ins = [
data[0],
data[1],
data[2],
[x for x in embeddings]
]
collection.insert(ins)
Poblar tu base de datos vectorial
Con una función que crea embeddings por lotes y realiza inserciones completa, estamos listos para poblar la base de datos. Para este ejemplo, recorremos cada fila de nuestro dataframe y la añadimos a una lista de listas que usamos para agrupar nuestros datos en lotes. Una vez que alcanzamos nuestro tamaño de lote, llamamos a la función embed_insert y restablecemos nuestro lote.
Si queda algún dato sobrante en el lote de datos después de terminar el bucle, generamos el embedding e insertamos los datos restantes. Finalmente, para terminar de poblar nuestra base de datos vectorial, llamamos a flush para asegurarnos de que la base de datos esté actualizada e indexada.
data_batch = [[], [], [], []]
for index, row in cleaned_df.iterrows():
data_batch[0].append(row["Title"])
data_batch[1].append(str(row["Date_time"]))
data_batch[2].append(row["Location"])
data_batch[3].append(row["Speech"])
if len(data_batch[0]) % BATCH_SIZE == 0:
embed_insert(data_batch)
data_batch = [[], [], [], []]
# Embed and insert the remainder
if len(data_batch[0]) != 0:
embed_insert(data_batch)
# Call a flush to index any unsealed segments.
collection.flush()
Búsqueda semántica de discursos de la Casa Blanca basada en descripciones
Digamos que me interesa encontrar un discurso en el que el presidente habló sobre el impacto de la energía renovable en el National Renewable Energy Lab (NREL) y un discurso en el que hablan la vicepresidenta y el primer ministro de Canadá. Puedo encontrar los títulos de los discursos más similares pronunciados por los miembros de la Casa Blanca en 2021-2022 usando la base de datos vectorial que acabamos de crear.
Podemos buscar en nuestra base de datos vectorial los discursos más similares a nuestras descripciones. Luego, todo lo que tenemos que hacer es convertir la descripción en un embedding vectorial usando el mismo modelo que usamos para obtener los embeddings de los discursos y después buscar en la base de datos vectorial.
Una vez que convertimos las descripciones en una incrustación vectorial, usamos la función search en nuestra colección. Pasamos las incrustaciones como los datos de búsqueda, pasamos el campo que estamos buscando, agregamos algunos parámetros sobre cómo buscar, un límite para el número de resultados y el campo que queremos devolver. En este ejemplo, los parámetros de búsqueda que necesitamos pasar son el tipo de métrica, que tiene que ser del mismo tipo que usamos al crear el índice (norma L2), y el número de clústeres que queremos buscar (estableciendo nprobe en 10).
import time
search_terms = ["The President speaks about the impact of renewable energy at the National Renewable Energy Lab.", "The Vice President and the Prime Minister of Canada both speak."]
# Search the database based on input text
def embed_search(data):
embeds = transformer.encode(data)
return [x for x in embeds]
search_data = embed_search(search_terms)
start = time.time()
res = collection.search(
data=search_data, # Embeded search value
anns_field="embedding", # Search across embeddings
param={"metric_type": "L2",
"params": {"nprobe": 10}},
limit = TOPK, # Limit to top_k results per search
output_fields=["title"] # Include title field in result
)
end = time.time()
for hits_i, hits in enumerate(res):
print("Title:", search_terms[hits_i])
print("Search Time:", end-start)
print("Results:")
for hit in hits:
print( hit.entity.get("title"), "----", hit.distance)
print()
Cuando buscamos las oraciones en este ejemplo, esperamos ver una salida como la imagen de abajo. Esta fue una búsqueda exitosa porque los títulos son los que esperamos ver. La primera descripción devuelve el título de un discurso pronunciado por el presidente Biden en NREL, y la segunda descripción devuelve un título que refleja un discurso pronunciado por la vicepresidenta Harris y el primer ministro Trudeau.
Resumen
En este tutorial, aprendimos cómo usar una base de datos vectorial para buscar semánticamente entre discursos pronunciados por la administración Biden antes de las elecciones de mitad de mandato de 2022. La búsqueda semántica nos permite tomar un fragmento de texto y buscar texto semánticamente similar, no solo texto sintácticamente similar. Esto nos permite buscar una descripción general de un discurso en lugar de buscar un discurso basándonos en oraciones o citas específicas. Para la mayoría de nosotros, eso hace que encontrar discursos que nos interesarían sea mucho más fácil.
Sigue leyendo
8 Latest RAG Advancements Every Developer Should Know
Explore eight advanced RAG variants that can solve real problems you might be facing: slow retrieval, poor context understanding, multimodal data handling, and resource optimization.
What Exactly Are AI Agents? Why OpenAI and LangChain Are Fighting Over Their Definition?
AI agents are software programs powered by AI that can perceive their environment, make decisions, and take actions to achieve a goal—often autonomously.
Selecting the Right ETL Tools for Unstructured Data to Prepare for AI
Learn the right ETL tools for unstructured data to power AI. Explore key challenges, tool comparisons, and integrations with Milvus for vector search.