Recherche de similarité vectorielle dans les articles scientifiques d’ArXiv avec Milvus 2.1
Introduction
L’une des meilleures façons d’apprendre n’importe quel sujet récent de Data Science consiste à lire des articles de recherche open source sur arxiv.org. Cependant, le très grand nombre d’articles de recherche peut être écrasant, même pour le chercheur le plus chevronné qui souhaite les trier. Des outils comme connected papers peuvent aider, mais ils mesurent la similarité en fonction des citations et de la bibliographie partagées entre les articles, et non du sens sémantique du texte dans ces documents.
Dans cet article, j’ai entrepris de créer un moteur de recherche par similarité sémantique, qui prend en entrée un seul article « requête » et utilise le NLP de pointe pour trouver les top-K articles les plus similaires dans le corpus arxiv d’environ 640 000 articles d’informatique ! La recherche s’exécute avec une latence <50 ms sur un seul ordinateur portable ! Plus précisément, dans cet article, je couvrirai
- Configurer l’environnement et télécharger les données arXiv depuis Kaggle
- Charger les données dans Python à l’aide de Dask
- Mettre en œuvre une application de recherche par similarité sémantique d’articles scientifiques à l’aide de la base de données vectorielle Milvus
Les techniques utilisées dans cet article peuvent servir de modèle pour créer n’importe quel moteur de recherche par similarité sémantique NLP, pas seulement pour des articles scientifiques. La seule différence serait le modèle pré-entraîné utilisé.
Pour cet article, nous utiliserons le jeu de données arXiv de Kaggle, que les auteurs ont publié sous la licence CC0 : Domaine public.
J’ai décrit les considérations de recherche par similarité vectorielle à l’échelle de la production dans mon article précédent. Toutes ces considérations s’appliquent également à ce projet. La base de données vectorielle Milvus est si bien conçue que de nombreuses étapes sont exactement les mêmes et sont reprises ici uniquement par souci d’exhaustivité.
Configurer l’environnement et télécharger les données arxiv depuis Kaggle.
Cornel University a téléversé l’intégralité du corpus arXiv dans un jeu de données Kaggle et l’a placé sous la licence CC0 : Domaine public. Nous pouvons télécharger directement le jeu de données à l’aide de l’API Kaggle. Si vous ne l’avez pas déjà fait, veuillez configurer l’API Kaggle sur votre système en suivant ces instructions.
Nous utiliserons un environnement conda pour cet article appelé semantic_similarity. Si vous n’avez pas installé conda sur votre système, vous pouvez le faire en installant le mini forge open source depuis son dépôt GitHub. Les étapes ci-dessous créent les répertoires nécessaires et l’environnement conda, installent les bibliothèques Python requises et téléchargent le jeu de données arxiv depuis 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
Charger les données dans Python à l’aide de Dask
Les données que nous avons téléchargées depuis Kaggle sont un fichier JSON de 3,3 Go contenant environ 2 millions d’articles ! Pour traiter efficacement un ensemble de données aussi volumineux, ce n’est pas une bonne idée de charger l’ensemble du jeu de données en mémoire avec pandas. À la place, nous pouvons utiliser Dask pour diviser les données en plusieurs partitions et ne charger que quelques partitions en mémoire à un moment donné.
Dask
Dask est une bibliothèque open-source qui nous permet d’appliquer facilement le calcul parallèle avec une API similaire à pandas. Elle est simple à configurer sur votre machine locale en exécutant,pip install dask[complete] comme indiqué dans la section de configuration. Commençons par importer les bibliothèques nécessaires.
import dask.bag as db
import json
from datetime import datetime
import time
data_path = '../data/arxiv-metadata-oai-snapshot.json'
Nous utiliserons deux composants de Dask pour traiter efficacement le grand fichier JSON d’arxiv.
- Dask Bag : il nous permet de charger le fichier JSON par blocs de taille fixe et d’exécuter certaines fonctions de prétraitement sur chaque ligne de données.
- Dask DataFrame : nous pouvons convertir un dask bag en dask dataframe pour accéder à des API similaires à celles de pandas
Étape 1 : Charger le fichier JSON dans un Dask bag
Chargeons le fichier JSON dans un dask bag où chaque bloc a une taille de 10 Mo. Vous pouvez ajuster l’argument blocksize pour contrôler la taille souhaitée de chaque bloc. Nous appliquons ensuite la fonction json.loads à chaque ligne du dask bag en utilisant la fonction .map() pour analyser la chaîne JSON en un dictionnaire 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)
Image par l’auteur
Étape 2 : Écrire des fonctions auxiliaires de prétraitement
D’après l’affichage, nous voyons que chaque ligne contient plusieurs métadonnées liées à un article. Écrivons trois fonctions auxiliaires pour nous aider à prétraiter le jeu de données.
v1_date(): cette fonction sert à extraire la date à laquelle les auteurs ont téléversé la première version de l’article sur arXiv. Nous convertirons la date en temps UNIX et la stockerons comme un nouveau champ dans cette ligne.text_col(): cette fonction sert à combiner les champs “title” et “abstract” en utilisant un token “[SEP]” afin que nous puissions fournir ces textes au modèle d’embedding SPECTRE. Nous parlerons davantage de SPECTRE dans la section suivante.filters(): cette fonction ne conserve que les lignes qui répondent à certains critères, tels que la longueur maximale du texte dans diverses colonnes et les articles de la catégorie Computer Science.
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
)
Étape 3 : Exécuter les fonctions auxiliaires de prétraitement sur le Dask bag
Nous pouvons facilement utiliser les fonctions .map() et .filter() pour exécuter les fonctions auxiliaires sur chaque ligne du Dask bag, comme indiqué ci-dessous. Étant donné que Dask prend en charge le chaînage de méthodes, nous profitons de cette occasion pour ne conserver que quelques colonnes essentielles dans notre Dask bag et supprimer le reste.
# 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)
Image par l’auteur
Étape 4 : Convertir le Dask Bag en Dask DataFrame
La dernière étape du chargement des données consiste à convertir le Dask Bag en Dask Dataframe afin d’utiliser des API similaires à pandas sur chaque bloc ou partition des données.
# 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()
Image par l’auteur
Implémenter une application de recherche de similarité sémantique d’articles scientifiques à l’aide de la base de données vectorielle Milvus
Milvus est l’une des bases de données vectorielles open-source les plus populaires, conçue pour une recherche de similarité hautement évolutive et extrêmement rapide. Nous utiliserons Milvus Standalone pour cet article, car nous exécutons Milvus uniquement sur notre machine locale.
Étape 1 : Installer la base de données vectorielle Milvus sur votre machine locale
L’installation de la base de données vectorielle Milvus est un jeu d’enfant avec Docker ; nous devons donc d’abord installer Docker et Docker Compose. Ensuite, il nous suffit de télécharger un docker-compose.yml et de démarrer les conteneurs docker, comme indiqué dans l’extrait de code ci-dessous ! Le site web milvus.io propose de nombreuses autres options pour installer à la fois Milvus standalone et Milvus Cluster ; n’hésitez pas à le consulter si vous devez l’installer sur un cluster Kubernetes ou l’installer hors ligne.
# 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
Étape 2 : Créer une collection Milvus
Maintenant que le serveur de base de données vectorielle Milvus fonctionne sur notre machine locale, nous pouvons interagir avec lui à l’aide de la bibliothèque pymilvus. Tout d’abord, importons les modules nécessaires et connectons-nous au serveur Milvus exécuté sur localhost. N’hésitez pas à modifier les paramètres alias et collection_name. Le modèle que nous utilisons pour convertir notre texte en embeddings détermine la valeur du paramètre emb_dim. Dans le cas de SPECTRE, les embeddings sont en 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)
Facultativement, vous pouvez vérifier si la collection spécifiée par collection_name est déjà présente sur votre serveur Milvus. Pour cet exemple, si la collection est déjà disponible, je la supprime. Mais sur un serveur de production, vous ne feriez pas cela et vous ignoreriez plutôt le code de création de collection ci-dessous.
Une collection Milvus est analogue à une table dans une base de données traditionnelle. Pour créer une collection afin de stocker des données, nous devons d’abord spécifier le schema de la collection. Dans cet exemple, nous exploitons la capacité de Milvus 2.1 à stocker des index et des champs de chaînes de caractères pour stocker toutes les métadonnées nécessaires liées à chaque article. La clé primaire idx et les autres champs categories, title, abstract ont le type de données VARCHAR avec des longueurs maximales raisonnables, tandis que embedding est un champ FLOAT_VECTOR contenant les embeddings de dimension emb_dim. Milvus prend en charge une grande variété de types de données, comme indiqué dans notre page de documentation.
# 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
)
Une fois qu’une collection a été créée, nous sommes maintenant prêts à y téléverser nos textes et nos vecteurs.
Étape 3 : Parcourir les partitions de notre dataframe Dask, intégrer les textes à l’aide de SPECTER et les téléverser dans la base de données vectorielle Milvus
Tout d’abord, nous devons convertir les textes du dataframe Dask en un vecteur d’embedding afin d’effectuer une recherche de similarité sémantique. Mon article ci-dessous explique comment nous pouvons convertir des textes en embeddings. En particulier, nous utiliserons un modèle SBERT Bi-Encoder appelé SPECTRE pour convertir des articles scientifiques en embeddings.
SPECTER [article] [Github] : Scientific Paper Embeddings using Citation-informed TransformERs est un modèle permettant de convertir des articles scientifiques en embeddings.
- Les textes du titre et du résumé de chaque article sont concaténés avec le jeton [SEP] et convertis en embeddings à l’aide du jeton [CLS] d’un modèle Transformer pré-entraîné (SciBERT).
- Utilisez les citations comme signal proxy de la relation entre documents. Si un article en cite un autre, nous pouvons en déduire qu’ils sont tous deux liés.
- Objectif d’entraînement avec perte triplet : nous entraînons le modèle Transformer afin que les articles partageant des citations soient plus proches dans l’espace d’embedding.
- En d’autres termes, un article positif est un article cité dans l’article requête, tandis qu’un article négatif est un article non cité par l’article requête. Les négatifs échantillonnés aléatoirement sont des négatifs « faciles ».
- Pour améliorer les performances, nous créons des négatifs « difficiles » à l’aide d’articles qui ne sont PAS cités par l’article requête mais qui SONT cités par l’article positif.
- Nous n’avons besoin que du titre et du résumé pendant l’inférence. Aucune citation n’est requise, SPECTER peut donc produire des embeddings même pour de nouveaux articles qui n’ont pas encore de citations !
- SPECTER offre d’excellentes performances (meilleures que SciBERT) en classification thématique, prédiction de citations et recommandation d’articles scientifiques.
Image par l’auteur utilisant des captures d’écran issues de l’article SPECTER open source
L’utilisation du modèle SPECTRE pré-entraîné est simple avec la bibliothèque Sentence Transformer. Nous pouvons télécharger le modèle pré-entraîné avec une seule ligne de code, comme indiqué ci-dessous. Nous écrivons également une fonction d’assistance simple pour convertir toute une colonne de textes de la partition du dataframe 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()
Nous devons itérer sur les partitions du dataframe Dask pour téléverser les données dans notre collection Milvus. À chaque itération, nous chargeons uniquement les lignes de cette partition en mémoire et ajoutons les données des colonnes de métadonnées à une variable data. Nous pouvons utiliser l’API dask .map_partitions() pour appliquer la génération d’embeddings à chaque ligne de la partition et ajouter les résultats à la même variable data. Enfin, nous pouvons téléverser les données vers Milvus avec 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)
Veuillez noter que l’ordre des colonnes ajoutées à la variable data doit suivre le même ordre que la variable fields que nous avons définie lors de la création du schéma !
Étape 4 : Créer un index Approximate Nearest Neighbors (ANN) sur les données téléversées
Après avoir inséré tous les embeddings dans la base de données vectorielle Milvus, nous devons créer un index ANN pour accélérer la recherche. Dans cet exemple, j’utilise le type d’index HNSW, l’un des index ANN les plus rapides et les plus précis. Consultez la documentation Milvus pour plus d’informations sur l’index HNSW et ses paramètres.
# 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)
Étape 5 : Exécutez vos requêtes de recherche de similarité vectorielle !
Enfin, les données de notre collection Milvus sont prêtes à être interrogées. Tout d’abord, nous devons charger la collection en mémoire pour exécuter des requêtes dessus.
# Load the collection into memory
collection = Collection(collection_name)
collection.load()
Ensuite, j’ai créé une fonction d’aide simple qui prend un query_text, le convertit en embedding SPECTRE, exécute une recherche ANN dans la collection Milvus, puis affiche les résultats. Nous pouvons contrôler la qualité et la vitesse de la recherche à l’aide des search_params décrits sur la page de documentation 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')}")
Nous pouvons maintenant utiliser la fonction d’aide avec une seule ligne de code pour lancer une recherche sémantique d’articles arXiv sur l’ensemble des quelque 640 000 articles d’informatique stockés dans notre collection Milvus. Par exemple, je recherche des articles similaires à l’article SimCSE dont j’ai parlé en détail dans mon précédent article. Les 10 premiers résultats sont tout à fait pertinents par rapport à ma requête de recherche, car ils sont pour la plupart liés à l’apprentissage contrastif d’embeddings de phrases ! Il est encore plus impressionnant que l’ensemble de la recherche n’ait pris que 30 ms en s’exécutant simplement sur mon ordinateur portable, ce qui est largement conforme aux exigences d’utilisation typiques de la plupart des applications !
# 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)
Image de l’auteur
Si nous n’avons plus besoin d’exécuter de requêtes, nous pouvons libérer la collection afin de libérer la mémoire de notre machine. Supprimer une collection de la mémoire n’entraîne pas de perte de données, car elle est toujours stockée sur notre disque et peut être chargée de nouveau lorsque nécessaire.
# Release the collection from memory when it's not needed anymore
collection.release()
Si vous souhaitez arrêter le serveur Milvus et supprimer toutes les données du disque, vous pouvez suivre les instructions pour arrêter Milvus. Attention ! Cette opération est irréversible et supprimera toutes les données de votre cluster Milvus.
Conclusion
Dans cet article, nous avons mis en œuvre en quelques étapes simples un service de recherche sémantique ultra-évolutif d’articles scientifiques utilisant les embeddings SPECTRE et la base de données vectorielle Milvus. Cette approche est évolutive en production jusqu’à des centaines de millions, voire des milliards de vecteurs. Nous avons testé la recherche à l’aide d’un exemple de requête d’article qui a renvoyé les 10 meilleurs résultats en seulement 30 ms ! La réputation de Milvus en tant que base de données de recherche de similarité vectorielle hautement évolutive et ultra-rapide est amplement méritée !
Pour plus d’inspiration sur les applications de Milvus, veuillez consulter les démos de la base de données vectorielle Milvus et le Bootcamp.
Continuer à lire
Top 10 Context Engineering Techniques You Should Know for Production RAG
A practical guide to context engineering for production LLM systems, covering RAG, context processing, memory, agents, and multimodal context.
Zilliz Cloud Launches in AWS Australia, Expanding Global Reach to Australia and Neighboring Markets
We're thrilled to announce that Zilliz Cloud is now available in the AWS Sydney, Australia region (ap-southeast-2).
Vector Databases vs. Time Series Databases
Use a vector database for similarity search and semantic relationships; use a time series database for tracking value changes over time.