Ricerca per similarità vettoriale di articoli scientifici ArXiv con Milvus 2.1
Introduzione
Uno dei modi migliori per imparare qualsiasi argomento recente di Data Science è leggere articoli di ricerca open-sourced su arxiv.org. Tuttavia, l’enorme numero di articoli di ricerca può essere travolgente anche per il ricercatore più esperto da esaminare. Strumenti come connected papers possono aiutare, ma misurano la similarità in base alle citazioni e alla bibliografia condivise tra gli articoli, non al significato semantico del testo in questi documenti.
In questo post, mi sono proposto di costruire un motore di ricerca per similarità semantica, che prende un singolo articolo “query” come input e utilizza NLP allo stato dell’arte per trovare i top-K articoli più simili dal corpus arxiv di circa 640K articoli di informatica! La ricerca viene eseguita con una latenza <50ms su un singolo laptop! Nello specifico, in questo post tratterò
- Configurare l’ambiente e scaricare i dati arXiv da Kaggle
- Caricare i dati in Python utilizzando Dask
- Implementare un’applicazione di ricerca per similarità semantica di articoli scientifici utilizzando il database vettoriale Milvus
Le tecniche utilizzate in questo post possono essere usate come modello per costruire qualsiasi motore di ricerca NLP per similarità semantica, non solo per articoli scientifici. L’unica differenza sarebbe il modello pre-addestrato utilizzato.
Per questo post, useremo l’arXiv Dataset from Kaggle, che gli autori hanno rilasciato con licenza CC0: Public Domain.
Ho delineato le considerazioni sulla ricerca di similarità vettoriale su scala di produzione nel mio post precedente. Tutte queste considerazioni si applicano anche a questo progetto. Il database vettoriale Milvus è progettato così bene che molti dei passaggi sono esattamente gli stessi e vengono replicati qui solo per completezza.
Configurare l’ambiente e scaricare i dati arxiv da Kaggle.
La Cornel University ha caricato l’intero corpus arXiv in un dataset Kaggle e lo ha concesso in licenza con la licenza CC0: Public Domain. Possiamo scaricare direttamente il dataset utilizzando l’API Kaggle. Se non lo hai già fatto, configura l’API Kaggle sul tuo sistema seguendo queste istruzioni.
Useremo un ambiente conda per questo post chiamato semantic_similarity. Se non hai installato conda sul tuo sistema, puoi farlo installando il mini forge open-sourced dal suo repository GitHub. I passaggi seguenti creano le directory necessarie e l’ambiente conda, installano le librerie Python richieste e scaricano l’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
Caricare i dati in Python utilizzando Dask
I dati che abbiamo scaricato da Kaggle sono un file JSON da 3,3 GB contenente circa 2 milioni di paper! Per elaborare in modo efficiente un dataset così grande, non è una buona idea caricare l’intero dataset in memoria usando pandas. Invece, possiamo usare Dask per suddividere i dati in più partizioni e caricare in memoria solo poche partizioni alla volta.
Dask
Dask è una libreria open-source che ci permette di applicare facilmente il calcolo parallelo con un’API simile a pandas. È semplice da configurare sulla tua macchina locale eseguendo,pip install dask[complete] come mostrato nella sezione di configurazione. Iniziamo importando prima le librerie necessarie.
import dask.bag as db
import json
from datetime import datetime
import time
data_path = '../data/arxiv-metadata-oai-snapshot.json'
Useremo due componenti di Dask per elaborare in modo efficiente il grande file JSON di arxiv.
- Dask Bag: Ci consente di caricare il file JSON in blocchi di dimensione fissa ed eseguire alcune funzioni di pre-elaborazione su ogni riga di dati.
- Dask DataFrame: Possiamo convertire un dask bag in un dask dataframe per accedere ad API simili a pandas
Passaggio 1: Caricare il file JSON in un Dask bag
Carichiamo il file JSON in un dask bag in cui ogni blocco ha una dimensione di 10 MB. Puoi modificare l’argomento blocksize per controllare quanto grande vuoi che sia ogni blocco. Applichiamo quindi la funzione json.loads a ogni riga del dask bag usando la funzione .map() per analizzare la stringa JSON in un dizionario 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)
Immagine dell’autore
Passaggio 2: Scrivere funzioni helper di pre-elaborazione
Dalla stampa, vediamo che ogni riga contiene diversi metadati relativi a un paper. Scriviamo tre funzioni helper per aiutarci a pre-elaborare il dataset.
v1_date(): Questa funzione serve a estrarre la data in cui gli autori hanno caricato la prima versione del paper su arXiv. Convertiremo la data in tempo UNIX e la memorizzeremo come nuovo campo in quella riga.text_col(): Questa funzione serve a combinare i campi “title” e “abstract” usando un token “[SEP]” in modo da poter fornire questi testi al modello di embedding SPECTRE. Parleremo di più di SPECTRE nella prossima sezione.filters(): Questa funzione conserva solo le righe che soddisfano alcuni criteri, come la lunghezza massima del testo in varie colonne e i paper nella categoria Informatica.
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):
"""
For each row in the dask bag, only keep the row if it meets the filter criteria
Args:
row: the row of the dataframe
Returns:
Boolean mask
"""
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"]) # Keep only CS papers
)
Passaggio 3: Eseguire le funzioni helper di pre-elaborazione sul Dask bag
Possiamo usare facilmente le funzioni .map() e .filter() per eseguire le funzioni helper su ogni riga del Dask bag, come mostrato di seguito. Poiché Dask supporta il concatenamento dei metodi, sfruttiamo questa opportunità per mantenere solo alcune colonne essenziali nel nostro Dask bag ed eliminare il resto.
# Specify columns to keep in the final table
cols_to_keep = ["id", "categories", "title", "abstract", "unix_time", "text"]
# Apply the pre-processing
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)
)
# Print the first row
papers_db.take(1)
Immagine dell'autore
Passaggio 4: Convertire il Dask Bag in un Dask DataFrame
Il passaggio finale del caricamento dei dati è convertire il Dask Bag in un Dask Dataframe per utilizzare API simili a pandas su ciascun blocco o partizione dei dati.
# Convert the Dask Bag to a Dask Dataframe
schema = {
"id": str,
"title": str,
"categories": str,
"abstract": str,
"unix_time": int,
"text": str,
}
papers_df = papers_db.to_dataframe(meta=schema)
# Display first 5 rows
papers_df.head()
Immagine dell'autore
Implementare un’applicazione di ricerca di similarità semantica per articoli scientifici usando il database vettoriale Milvus
Milvus è uno dei database vettoriali open-source più popolari, creato per una ricerca di similarità altamente scalabile e incredibilmente veloce. Useremo Milvus Standalone per questo post, poiché eseguiremo Milvus solo sulla nostra macchina locale.
Passaggio 1: Installare il database vettoriale Milvus in locale
Installare il database vettoriale Milvus è semplicissimo usando Docker, quindi per prima cosa dobbiamo installare Docker e Docker Compose. Poi, tutto ciò che dobbiamo fare è scaricare un docker-compose.yml e avviare i container docker, come mostrato nello snippet di codice qui sotto! Il sito web milvus.io offre molte altre opzioni per installare sia Milvus standalone sia Milvus Cluster; consultatelo se avete bisogno di installarlo su un cluster Kubernetes o installarlo offline.
# CD into milvus directory
cd semantic_similarity/milvus
# Download the Standalone version of Milvus docker compose
wget https://github.com/milvus-io/milvus/releases/download/v2.1.0/milvus-standalone-docker-compose.yml -O ./docker-compose.yml
# Run the Milvus server docker container on your local
sudo docker-compose up -d
Passaggio 2: Creare una collection Milvus
Ora che abbiamo il server del database vettoriale Milvus in esecuzione sulla nostra macchina locale, possiamo interagire con esso usando la libreria pymilvus. Per prima cosa, importiamo i moduli necessari e connettiamoci al server Milvus in esecuzione su localhost. Sentitevi liberi di modificare i parametri alias e collection_name. Il modello che usiamo per convertire il nostro testo in embedding determina il valore del parametro emb_dim. Nel caso di SPECTRE, gli embedding sono 768d.
# Assicurati che un server Milvus sia già in esecuzione
from pymilvus import connections, utility
from pymilvus import Collection, CollectionSchema, FieldSchema, DataType
# Connettiti al server Milvus
connections.connect(alias="default", host="localhost", port="19530")
# Nome della collection
collection_name = "arxiv"
# Dimensione dell'embedding
emb_dim = 768
# # Verifica la presenza di una collection esistente ed eliminala se esiste
# if utility.has_collection(collection_name):
# print(utility.list_collections())
# utility.drop_collection(collection_name)
Facoltativamente, puoi verificare se la collection specificata da collection_name è già presente sul tuo server Milvus. Per questo esempio, se la collection è già disponibile, la elimino. Ma in un server di produzione, non lo faresti e invece salteresti il codice di creazione della collection qui sotto.
Una collection Milvus è analoga a una tabella in un database tradizionale. Per creare una collection per archiviare i dati, dobbiamo prima specificare lo schema della collection. In questo esempio, stiamo sfruttando la capacità di Milvus 2.1 di archiviare indici e campi stringa per archiviare tutti i metadati necessari relativi a ciascun paper. La chiave primaria idx e gli altri campi categories, title, abstract hanno tipo di dato VARCHAR con lunghezze massime ragionevoli, mentre embedding è un campo FLOAT_VECTORcontenente gli embedding di dimensione emb_dim. Milvus supporta un'ampia varietà di tipi di dato, come mostrato nella nostra pagina di documentazione.
# Crea uno schema per la 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)
# Campi nella collection
fields = [idx, categories, title, abstract, unix_time, embedding]
schema = CollectionSchema(
fields=fields, description="Semantic Similarity of Scientific Papers"
)
# Crea una collection con lo schema
collection = Collection(
name=collection_name, schema=schema, using="default", shards_num=10
)
Una volta creata una collection, ora siamo pronti a caricare i nostri testi e vettori al suo interno.
Passaggio 3: Iterare sulle partizioni del nostro dataframe Dask, incorporare i testi usando SPECTER e caricarli nel database vettoriale Milvus
Per prima cosa, dobbiamo convertire i testi nel dataframe Dask in un vettore di embedding per eseguire una ricerca di similarità semantica. Il mio post qui sotto mostra come possiamo convertire testi in embedding. In particolare, useremo un modello SBERT Bi-Encoder chiamato SPECTRE per convertire paper scientifici in embedding.
SPECTER [paper] [Github]: Scientific Paper Embeddings using Citation-informed TransformERs è un modello per convertire paper scientifici in embedding.
- Il Titolo e il testo dell’Abstract di ciascun articolo vengono concatenati con il token [SEP] e convertiti in embedding usando il token [CLS] di un modello Transformer pre-addestrato (SciBERT).
- Usa le citazioni come segnale proxy per la correlazione tra documenti. Se un articolo ne cita un altro, possiamo dedurre che siano entrambi correlati.
- Obiettivo di addestramento con triplet loss: addestriamo il modello Transformer in modo che gli articoli con citazioni condivise siano più vicini nello spazio degli embedding.
- In altre parole, un articolo Positivo è un articolo citato nell’articolo Query, mentre un articolo Negativo è un articolo non citato dall’articolo Query. I negativi campionati casualmente sono negativi “facili”.
- Per migliorare le prestazioni, creiamo negativi “difficili” usando articoli che NON sono citati dall’articolo Query ma SONO citati dall’articolo Positivo.
- Durante l’inferenza abbiamo bisogno solo del Titolo e dell’Abstract. Non sono richieste citazioni, quindi SPECTER può produrre embedding anche per nuovi articoli che non hanno ancora citazioni!
- SPECTER offre prestazioni eccellenti (migliori di SciBERT) nella classificazione per argomento, nella previsione delle citazioni e nella raccomandazione di Articoli Scientifici.
Immagine dell’Autore usando screenshot dall’articolo SPECTER open-source
Usare il modello SPECTRE pre-addestrato è semplice con la libreria Sentence Transformer. Possiamo scaricare il modello pre-addestrato con una sola riga di codice, come mostrato di seguito. Scriviamo anche una semplice funzione di supporto per convertire un’intera colonna di testi dalla partizione del dataframe Dask in embedding.
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()
Dobbiamo iterare sulle partizioni del dataframe Dask per caricare i dati nella nostra collection Milvus. Durante ogni iterazione, carichiamo in memoria solo le righe di quella partizione e aggiungiamo i dati dalle colonne di metadati a una variabile data. Possiamo usare l’API dask .map_partitions() per applicare la generazione degli embedding a ogni riga nella partizione e aggiungere i risultati alla stessa variabile data. Infine, possiamo caricare i dati su 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)
Tieni presente che l’ordine delle colonne aggiunte nella variabile data deve seguire lo stesso ordine della variabile fields che abbiamo definito durante la creazione dello schema!
Passaggio 4: Creare un indice Approximate Nearest Neighbors (ANN) sui dati caricati
Dopo aver inserito tutti gli embedding nel database vettoriale Milvus, dobbiamo creare un indice ANN per velocizzare la ricerca. In questo esempio, sto usando il tipo di indice HNSW, uno degli indici ANN più veloci e accurati. Consulta la documentazione di Milvus per maggiori informazioni sull’indice HNSW e sui suoi parametri.
# 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)
Passaggio 5: Esegui le tue query di Ricerca per Similarità Vettoriale!
Infine, i dati nella nostra collection Milvus sono pronti per essere interrogati. Per prima cosa, dobbiamo caricare la collection in memoria per eseguire query su di essa.
# Carica la collection in memoria
collection = Collection(collection_name)
collection.load()
Successivamente, ho creato una semplice funzione helper che prende in input un query_text, lo converte nell’embedding SPECTRE, esegue una ricerca ANN sulla collection Milvus e stampa i risultati. Possiamo controllare la qualità e la velocità della ricerca usando i search_params descritti nella pagina della documentazione HNSW.
def query_and_display(query_text, collection, num_results=10):
# Incorpora il testo della query
query_emb = [model.encode(query_text)]
# Parametri di ricerca
search_params = {"metric_type": "L2", "params": {"ef": 128}}
# Ricerca
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()
# Stampa i risultati
print(f"Velocità della query: {(query_end - query_start).total_seconds():.2f} s")
print("Risultati:")
for res in results[0]:
title = res.entity.get("title").replace("\n ", "")
print(f"➡️ ID: {res.id}. Distanza L2: {res.distance:.2f}")
print(f"Titolo: {title}")
print(f"Abstract: {res.entity.get('abstract')}")
Ora possiamo usare la funzione helper con una sola riga di codice per eseguire una ricerca semantica di paper arXiv sull’intero insieme di ~640K paper di Computer Science archiviati nella nostra collection Milvus. Per esempio, sto cercando alcuni paper simili al paper SimCSE di cui ho discusso in dettaglio nel mio post precedente. I primi 10 risultati sono molto rilevanti per la mia query di ricerca, poiché sono per lo più correlati all’apprendimento contrastivo degli embedding di frasi! È ancora più impressionante che l’intera ricerca abbia richiesto solo 30 ms eseguita sul mio laptop, un tempo ampiamente entro i requisiti d’uso tipici della maggior parte delle applicazioni!
# Query per paper simili al paper SimCSE
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)
Immagine dell’autore
Se non dobbiamo eseguire altre query, possiamo rilasciare la collection per liberare la memoria della nostra macchina. Rimuovere una collection dalla memoria non causa perdita di dati, poiché è ancora archiviata sul nostro disco e può essere caricata di nuovo quando necessario.
# Rilascia la collection dalla memoria quando non è più necessaria
collection.release()
Se vuoi arrestare il server Milvus ed eliminare tutti i dati dal disco, puoi seguire le istruzioni per arrestare Milvus. Attenzione! Questa operazione è irreversibile ed eliminerà tutti i dati nel tuo cluster Milvus.
Conclusione
In questo post, abbiamo implementato un servizio di ricerca semantica ultra-scalabile di articoli scientifici utilizzando gli embedding SPECTRE e il database vettoriale Milvus in pochi semplici passaggi. Questo approccio è scalabile in produzione fino a centinaia di milioni o persino miliardi di vettori. Abbiamo testato la ricerca usando una query di esempio su un articolo che ha restituito i primi 10 risultati in soli 30 ms! La reputazione di Milvus come database di ricerca per similarità vettoriale altamente scalabile e velocissimo è ben meritata!
Per ulteriori spunti sulle applicazioni di Milvus, vai a Milvus database vettoriale demo e Bootcamp.
Continua a leggere
Vector Lakebase: End the AI Data Silo
Learn how Vector Lakebase unifies vector search, data lakes, and AI data operations so teams can serve RAG and agents without copy-and-sync pipelines.
Why Deepseek is Waking up AI Giants Like OpenAI And Why You Should Care
Discover how DeepSeek R1's open-source AI model with superior reasoning capabilities and lower costs is disrupting the AI landscape and challenging tech giants like OpenAI.
DeepSeek-VL2: Mixture-of-Experts Vision-Language Models for Advanced Multimodal Understanding
Explore DeepSeek-VL2, the open-source MoE vision-language model. Discover its architecture, efficient training pipeline, and top-tier performance.