Búsqueda de similitud vectorial en artículos científicos de arXiv con Milvus 2.1
Introducción
Una de las mejores formas de aprender cualquier tema reciente de Ciencia de Datos es leyendo artículos de investigación de código abierto en arxiv.org. Sin embargo, la enorme cantidad de artículos de investigación puede ser abrumadora incluso para el investigador más experimentado a la hora de revisarlos. Herramientas como connected papers pueden ayudar, pero miden la similitud basándose en las citas y bibliografía compartidas entre artículos, no en el significado semántico del texto de estos documentos.
En esta publicación, me propuse construir un motor de búsqueda de similitud semántica, que toma como entrada un único artículo de “consulta” y utiliza NLP de última generación para encontrar los top-K artículos más similares del corpus de arxiv de aproximadamente 640K artículos de ciencias de la computación. ¡La búsqueda se ejecuta con una latencia de <50ms en un solo portátil! Específicamente, en esta publicación cubriré
- Configurar el entorno y descargar los datos de arXiv desde Kaggle
- Cargar los datos en Python usando Dask
- Implementar una aplicación de búsqueda de similitud semántica de artículos científicos usando la base de datos vectorial Milvus
Las técnicas utilizadas en esta publicación pueden usarse como plantilla para construir cualquier motor de búsqueda de similitud semántica de NLP, no solo para artículos científicos. La única diferencia sería el modelo preentrenado utilizado.
Para esta publicación, usaremos el arXiv Dataset from Kaggle, que los autores publicaron bajo la licencia CC0: Public Domain.
Describí las consideraciones de búsqueda de similitud vectorial a escala de producción en mi publicación anterior. Todas esas consideraciones también se aplican a este proyecto. La base de datos vectorial Milvus está tan bien diseñada que muchos de los pasos son exactamente los mismos y se replican aquí solo para que esté completo.
Configurar el entorno y descargar los datos de arxiv desde Kaggle.
Cornel University ha subido todo el corpus de arXiv a un Kaggle dataset y lo ha licenciado bajo la licencia CC0: Public Domain. Podemos descargar directamente el conjunto de datos usando la API de Kaggle. Si aún no lo has hecho, configura la API de Kaggle en tu sistema siguiendo estas instrucciones.
Usaremos un entorno conda para esta publicación llamado semantic_similarity. Si no has instalado conda en tu sistema, puedes hacerlo instalando el mini forge de código abierto desde su repositorio de GitHub. Los pasos siguientes crean los directorios necesarios y el entorno conda, instalan las bibliotecas de Python requeridas y descargan el arxiv dataset from Kaggle.
# Create the necessary directories
mkdir -p semantic_similarity/notebooks semantic_similarity/data semantic_similarity/milvus
# CD into the data directory
cd semantic_similarity/data
# Create and activate a conda environment
conda create -n semantic_similarity python=3.9
conda activate semantic_similarity
## Create Virtual Environment using venv if not using conda
# python -m venv semantic_similarity
# source semantic_similarity/bin/activate
# Pip install the necessary libraries
pip install jupyterlab kaggle matplotlib scikit-learn tqdm ipywidgets
pip install "dask[complete]" sentence-transformers
pip install pandas pyarrow pymilvus protobuf==3.20.0
# Download data using the kaggle API
kaggle datasets download -d Cornell-University/arxiv
# Unzip the data into the local directory
unzip arxiv.zip
# Delete the Zip file
rm arxiv.zip
Cargar los datos en Python usando Dask
Los datos que descargamos de Kaggle son un archivo JSON de 3,3 GB que contiene alrededor de 2 millones de artículos. Para procesar eficientemente un conjunto de datos tan grande, no es buena idea cargar todo el conjunto de datos en memoria usando pandas. En su lugar, podemos usar Dask para dividir los datos en múltiples particiones y cargar solo unas pocas particiones en memoria en un momento dado.
Dask
Dask es una biblioteca de código abierto que nos permite aplicar computación paralela fácilmente con una API similar a pandas. Es sencillo configurarla en tu máquina local ejecutando,pip install dask[complete] como se muestra en la sección de configuración. Empecemos importando primero las bibliotecas necesarias.
import dask.bag as db
import json
from datetime import datetime
import time
data_path = '../data/arxiv-metadata-oai-snapshot.json'
Usaremos dos componentes de Dask para procesar eficientemente el gran archivo JSON de arxiv.
- Dask Bag: Nos permite cargar el archivo JSON en bloques de un tamaño fijo y ejecutar algunas funciones de preprocesamiento en cada fila de datos.
- Dask DataFrame: Podemos convertir un dask bag en un dask dataframe para obtener acceso a APIs similares a las de pandas
Paso 1: Cargar el archivo JSON en un Dask bag
Carguemos el archivo JSON en un dask bag donde cada bloque tenga un tamaño de 10 MB. Puedes ajustar el argumento blocksize para controlar qué tan grande quieres que sea cada bloque. Luego aplicamos la función json.loads a cada fila del dask bag usando la función .map() para analizar la cadena JSON y convertirla en un diccionario de Python.
# Read the file in blocks of 10MB and parse the JSON.
papers_db = db.read_text(data_path, blocksize="10MB").map(json.loads)
# Print the first row
papers_db.take(1)
Imagen del autor
Paso 2: Escribir funciones auxiliares de preprocesamiento
A partir de la impresión, vemos que cada fila contiene varios metadatos relacionados con un artículo. Escribamos tres funciones auxiliares para ayudarnos a preprocesar el conjunto de datos.
v1_date(): Esta función sirve para extraer la fecha en que los autores subieron la primera versión del artículo a arXiv. Convertiremos la fecha a tiempo UNIX y la almacenaremos como un nuevo campo en esa fila.text_col(): Esta función sirve para combinar los campos “title” y “abstract” usando un token “[SEP]” para que podamos introducir estos textos en el modelo de embeddings SPECTRE. Hablaremos más sobre SPECTRE en la siguiente sección.filters(): Esta función conserva solo las filas que cumplen algunos criterios, como la longitud máxima del texto en varias columnas y artículos en la categoría de Ciencias de la Computación.
def v1_date(row):
"""
For each row in the dask bag,
find the date of the first version of the paper
and add it to the row as a new column
Args:
row: a row of the dask bag
Returns:
A row of the dask bag with added "unix_time" column
"""
versions = row["versions"]
date = None
for version in versions:
if version["version"] == "v1":
date = datetime.strptime(version["created"], "%a, %d %b %Y %H:%M:%S %Z")
date = int(time.mktime(date.timetuple()))
row["unix_time"] = date
return row
def text_col(row):
"""
It takes a row of a dataframe, adds a new column called 'text'
that is the concatenation of the 'title' and 'abstract' columns
Args:
row: the row of the dataframe
Returns:
A row with the text column added.
"""
row["text"] = row["title"] + "[SEP]" + row["abstract"]
return row
def filters(row):
"""
Para cada fila en la bolsa de dask, solo conserva la fila si cumple los criterios de filtrado
Args:
row: la fila del dataframe
Returns:
Máscara booleana
"""
return ((len(row["id"])<16) and
(len(row["categories"])<200) and
(len(row["title"])<4096) and
(len(row["abstract"])<65535) and
("cs." in row["categories"]) # Conservar solo artículos de CS
)
Paso 3: Ejecutar las funciones auxiliares de preprocesamiento en la Dask bag
Podemos usar fácilmente las funciones .map() y .filter() para ejecutar las funciones auxiliares en cada fila de la Dask bag, como se muestra a continuación. Dado que Dask admite el encadenamiento de métodos, aprovechamos esta oportunidad para conservar solo unas pocas columnas esenciales en nuestra Dask bag y descartar el resto.
# Especificar columnas que se conservarán en la tabla final
cols_to_keep = ["id", "categories", "title", "abstract", "unix_time", "text"]
# Aplicar el preprocesamiento
papers_db = (
papers_db.map(lambda row: v1_date(row))
.map(lambda row: text_col(row))
.map(
lambda row: {
key: value
for key, value in row.items()
if key in cols_to_keep
}
)
.filter(filters)
)
# Imprimir la primera fila
papers_db.take(1)
Imagen del autor
Paso 4: Convertir la Dask Bag en un Dask DataFrame
El paso final de la carga de datos es convertir la Dask Bag en un Dask Dataframe para usar APIs similares a las de pandas en cada bloque o partición de los datos.
# Convertir la Dask Bag en un Dask Dataframe
schema = {
"id": str,
"title": str,
"categories": str,
"abstract": str,
"unix_time": int,
"text": str,
}
papers_df = papers_db.to_dataframe(meta=schema)
# Mostrar las primeras 5 filas
papers_df.head()
Imagen del autor
Implementar una aplicación de búsqueda de similitud semántica de artículos científicos usando la base de datos vectorial Milvus
Milvus es una de las bases de datos vectoriales de código abierto más populares, creada para búsquedas de similitud altamente escalables y extremadamente rápidas. Usaremos Milvus Standalone para esta publicación, ya que solo ejecutaremos Milvus en nuestra máquina local.
Paso 1: Instalar la base de datos vectorial Milvus en tu equipo local
Instalar la base de datos vectorial Milvus es muy sencillo usando Docker, así que primero debemos instalar Docker y Docker Compose. Luego, todo lo que necesitamos hacer es descargar un docker-compose.yml e iniciar los contenedores de docker, como se muestra en el fragmento de código a continuación. El sitio web de milvus.io ofrece muchas otras opciones para instalar tanto Milvus standalone como Milvus Cluster; échale un vistazo si necesitas instalarlo en un clúster de Kubernetes o instalarlo sin conexión.
# Cambiar al directorio de milvus
cd semantic_similarity/milvus
# Descargar la versión Standalone de Milvus docker compose
wget https://github.com/milvus-io/milvus/releases/download/v2.1.0/milvus-standalone-docker-compose.yml -O ./docker-compose.yml
# Ejecutar el contenedor docker del servidor Milvus en tu equipo local
sudo docker-compose up -d
Paso 2: Crear una colección de Milvus
Ahora que tenemos el servidor de la base de datos vectorial Milvus ejecutándose en nuestra máquina local, podemos interactuar con él usando la biblioteca pymilvus. Primero, importemos los módulos necesarios y conectémonos al servidor Milvus que se ejecuta en localhost. Siéntete libre de cambiar los parámetros alias y collection_name. El modelo que usamos para convertir nuestro texto en embeddings determina el valor del parámetro emb_dim. En el caso de SPECTRE, los embeddings son de 768d.
# Make sure a Milvus server is already running
from pymilvus import connections, utility
from pymilvus import Collection, CollectionSchema, FieldSchema, DataType
# Connect to Milvus server
connections.connect(alias="default", host="localhost", port="19530")
# Collection name
collection_name = "arxiv"
# Embedding size
emb_dim = 768
# # Check for existing collection and drop if exists
# if utility.has_collection(collection_name):
# print(utility.list_collections())
# utility.drop_collection(collection_name)
Opcionalmente, puedes comprobar si la colección especificada por collection_name ya está presente en tu servidor Milvus. Para este ejemplo, si la colección ya está disponible, la elimino. Pero en un servidor de producción, no harías esto y, en su lugar, omitirías el código de creación de la colección que aparece a continuación.
Una colección de Milvus es análoga a una tabla en una base de datos tradicional. Para crear una colección para almacenar datos, primero necesitamos especificar el schema de la colección. En este ejemplo, estamos aprovechando la capacidad de Milvus 2.1 para almacenar índices y campos de cadenas para almacenar todos los metadatos necesarios relacionados con cada artículo. La clave primaria idx y otros campos categories, title, abstract tienen el tipo de dato VARCHAR con longitudes máximas razonables, mientras que embedding es un campo FLOAT_VECTOR que contiene los embeddings de dimensión emb_dim. Milvus admite una amplia variedad de tipos de datos, como se muestra en nuestra página de documentación.
# Create a schema for the collection
idx = FieldSchema(name="id", dtype=DataType.VARCHAR, is_primary=True, max_length=16)
categories = FieldSchema(name="categories", dtype=DataType.VARCHAR, max_length=200)
title = FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=4096)
abstract = FieldSchema(name="abstract", dtype=DataType.VARCHAR, max_length=65535)
unix_time = FieldSchema(name="unix_time", dtype=DataType.INT64)
embedding = FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=emb_dim)
# Fields in the collection
fields = [idx, categories, title, abstract, unix_time, embedding]
schema = CollectionSchema(
fields=fields, description="Semantic Similarity of Scientific Papers"
)
# Create a collection with the schema
collection = Collection(
name=collection_name, schema=schema, using="default", shards_num=10
)
Una vez que se ha creado una colección, ahora estamos listos para cargar nuestros textos y vectores en ella.
Paso 3: Iterar sobre las particiones de nuestro dataframe de Dask, incrustar los textos usando SPECTER y cargarlos en la base de datos vectorial de Milvus
Primero, necesitamos convertir los textos en el dataframe de Dask en un vector de embedding para ejecutar una búsqueda de similitud semántica. Mi publicación a continuación comparte cómo podemos convertir textos en embeddings. En particular, usaremos un modelo SBERT Bi-Encoder llamado SPECTRE para convertir artículos científicos en embeddings.
SPECTER [artículo] [Github]: Scientific Paper Embeddings using Citation-informed TransformERs es un modelo para convertir artículos científicos en embeddings.
- Los textos del Título y el Resumen de cada artículo se concatenan con el token [SEP] y se convierten en embeddings usando el token [CLS] de un modelo Transformer preentrenado (SciBERT).
- Usa las citas como una señal proxy de relación entre documentos. Si un artículo cita a otro, podemos inferir que ambos están relacionados.
- Objetivo de entrenamiento con pérdida tripleta: entrenamos el modelo Transformer para que los artículos con citas compartidas estén más cerca en el espacio de embeddings.
- En otras palabras, un artículo Positivo es un artículo citado en el artículo de Consulta, mientras que un artículo Negativo es un artículo no citado por el artículo de Consulta. Los negativos muestreados aleatoriamente son negativos “fáciles”.
- Para mejorar el rendimiento, creamos negativos “difíciles” usando artículos que NO son citados por el artículo de Consulta pero SÍ son citados por el artículo Positivo.
- Solo necesitamos el Título y el Resumen durante la inferencia. No se requieren citas, así que SPECTER puede producir embeddings incluso para artículos nuevos que aún no tienen ninguna cita.
- SPECTER ofrece un rendimiento excelente (mejor que SciBERT) en clasificación de temas, predicción de citas y recomendación de artículos científicos.
Imagen del autor usando capturas de pantalla del artículo SPECTER de código abierto
Usar el modelo SPECTRE preentrenado es simple con la biblioteca Sentence Transformer. Podemos descargar el modelo preentrenado con solo una línea de código, como se muestra a continuación. También escribimos una función auxiliar simple para convertir una columna completa de textos de la partición del dataframe de Dask en embeddings.
from sentence_transformers import SentenceTransformer
from tqdm import tqdm
# Scientific Papers SBERT Model
model = SentenceTransformer('allenai-specter')
def emb_gen(partition):
return model.encode(partition['text']).tolist()
Necesitamos iterar sobre las particiones del dataframe de Dask para cargar los datos en nuestra colección de Milvus. Durante cada iteración, cargamos en memoria solo las filas de esa partición y añadimos los datos de las columnas de metadatos a una variable data. Podemos usar la API .map_partitions() de dask para aplicar la generación de embeddings a cada fila de la partición y anexar los resultados de nuevo a la misma variable data. Finalmente, podemos cargar los datos a Milvus con collection.insert.
# Initialize
collection = Collection(collection_name)
for partition in tqdm(range(papers_df.npartitions)):
# Get the dask dataframe for the partition
subset_df = papers_df.get_partition(partition)
# Check if dataframe is empty
if len(subset_df.index) != 0:
# Metadata
data = [
subset_df[col].values.compute().tolist()
for col in ["id", "categories", "title", "abstract", "unix_time"]
]
# Embeddings
data += [
subset_df
.map_partitions(emb_gen)
.compute()[0]
]
# Insert data
collection.insert(data)
¡Ten en cuenta que el orden de las columnas añadidas a la variable data debe seguir el mismo orden que la variable fields que definimos durante la creación del esquema!
Paso 4: Crear un índice de Vecinos Más Cercanos Aproximados (ANN) sobre los datos cargados
Después de insertar todos los embeddings en la base de datos vectorial de Milvus, necesitamos crear un índice ANN para acelerar la búsqueda. En este ejemplo, estoy usando el tipo de índice HNSW, uno de los índices ANN más rápidos y precisos. Consulta la documentación de Milvus para obtener más información sobre el índice HNSW y sus parámetros.
# Add an ANN index to the collection
index_params = {
"metric_type": "L2",
"index_type": "HNSW",
"params": {"efConstruction": 128, "M": 8},
}
collection.create_index(field_name="embedding", index_params=index_params)
Paso 5: ¡Ejecuta tus consultas de Búsqueda por Similitud Vectorial!
Finalmente, los datos de nuestra colección de Milvus están listos para consultarse. Primero, debemos cargar la colección en memoria para ejecutar consultas sobre ella.
# Load the collection into memory
collection = Collection(collection_name)
collection.load()
A continuación, he creado una función auxiliar simple que recibe un query_text, lo convierte en el embedding de SPECTRE, ejecuta una búsqueda ANN en la colección de Milvus e imprime los resultados. Podemos controlar la calidad y la velocidad de la búsqueda usando los search_params descritos en la página de documentación de HNSW.
def query_and_display(query_text, collection, num_results=10):
# Embed the Query Text
query_emb = [model.encode(query_text)]
# Search Params
search_params = {"metric_type": "L2", "params": {"ef": 128}}
# Search
query_start = datetime.now()
results = collection.search(
data=query_emb,
anns_field="embedding",
param=search_params,
limit=num_results,
expr=None,
output_fields=["title", "abstract"],
)
query_end = datetime.now()
# Print Results
print(f"Query Speed: {(query_end - query_start).total_seconds():.2f} s")
print("Results:")
for res in results[0]:
title = res.entity.get("title").replace("\n ", "")
print(f"➡️ ID: {res.id}. L2 Distance: {res.distance:.2f}")
print(f"Title: {title}")
print(f"Abstract: {res.entity.get('abstract')}")
Ahora podemos usar la función auxiliar con solo una línea de código para ejecutar una búsqueda semántica de artículos de arXiv sobre los ~640K artículos de Computer Science almacenados en nuestra colección de Milvus. Por ejemplo, estoy buscando algunos artículos similares al artículo de SimCSE que analicé en detalle en mi publicación anterior. ¡Los 10 resultados principales son bastante relevantes para mi consulta de búsqueda, ya que están relacionados principalmente con el aprendizaje contrastivo de embeddings de oraciones! Es aún más impresionante que toda la búsqueda haya tardado solo 30 ms ejecutándose en mi laptop, lo cual está muy dentro de los requisitos de uso típicos para la mayoría de las aplicaciones.
# Query for papers that are similar to the SimCSE paper
title = "SimCSE: Simple Contrastive Learning of Sentence Embeddings"
abstract = """This paper presents SimCSE, a simple contrastive learning framework that greatly advances state-of-the-art sentence embeddings. We first describe an unsupervised approach, which takes an input sentence and predicts itself in a contrastive objective, with only standard dropout used as noise. This simple method works surprisingly well, performing on par with previous supervised counterparts. We find that dropout acts as minimal data augmentation, and removing it leads to a representation collapse. Then, we propose a supervised approach, which incorporates annotated pairs from natural language inference datasets into our contrastive learning framework by using "entailment" pairs as positives and "contradiction" pairs as hard negatives. We evaluate SimCSE on standard semantic textual similarity (STS) tasks, and our unsupervised and supervised models using BERT base achieve an average of 76.3% and 81.6% Spearman's correlation respectively, a 4.2% and 2.2% improvement compared to the previous best results. We also show -- both theoretically and empirically -- that the contrastive learning objective regularizes pre-trained embeddings' anisotropic space to be more uniform, and it better aligns positive pairs when supervised signals are available."""
query_text = f"{title}[SEP]{abstract}"
query_and_display(query_text, collection, num_results=10)
Imagen del autor
Si no necesitamos ejecutar más consultas, podemos liberar la colección para liberar memoria de nuestra máquina. Eliminar una colección de la memoria no causa pérdida de datos, ya que todavía está almacenada en nuestro disco y se puede cargar de nuevo cuando sea necesario.
# Liberar la colección de la memoria cuando ya no sea necesaria
collection.release()
Si quieres detener el servidor de Milvus y eliminar todos los datos del disco, puedes seguir las instrucciones para detener Milvus. ¡Cuidado! Esta operación es irreversible y eliminará todos los datos de tu clúster de Milvus.
Conclusión
En esta publicación, implementamos un servicio de búsqueda semántica ultrascalable de artículos científicos utilizando embeddings de SPECTRE y la base de datos vectorial Milvus en unos pocos pasos sencillos. Este enfoque es escalable en producción a cientos de millones o incluso miles de millones de vectores. Probamos la búsqueda utilizando una consulta de artículo de muestra que devolvió los 10 mejores resultados en solo 30 ms. ¡La reputación de Milvus como una base de datos de búsqueda de similitud vectorial altamente escalable y ultrarrápida está bien merecida!
Para más inspiración sobre las aplicaciones de Milvus, visita las demos de la base de datos vectorial Milvus y el Bootcamp.
Sigue leyendo
Introducing Zilliz Cloud Global Cluster: Region-Level Resilience for Mission-Critical AI
Zilliz Cloud Global Cluster delivers multi-region resilience, automatic failover, and fast global AI search with built-in security and compliance.
Introducing Zilliz MCP Server: Natural Language Access to Your Vector Database
Developers can easily manage and query vector databases with natural language via Zilliz MCP Server in AI-native environments.
Build for the Boom: Why AI Agent Startups Should Build Scalable Infrastructure Early
Explore strategies for developing AI agents that can handle rapid growth. Don't let inadequate systems undermine your success during critical breakthrough moments.