Векторный поиск похожих научных статей ArXiv с помощью Milvus 2.1
Введение
Один из лучших способов изучить любую новейшую тему в Data Science — читать исследовательские статьи с открытым доступом на arxiv.org. Однако огромное количество исследовательских статей может быть подавляющим даже для самого опытного исследователя, пытающегося в них разобраться. Такие инструменты, как connected papers, могут помочь, но они измеряют сходство на основе цитирований и библиографии, общих для статей, а не семантического смысла текста в этих документах.
В этом посте я решил создать поисковую систему семантического сходства, которая принимает одну “запросную” статью на вход и использует передовые NLP-методы, чтобы найти top-K наиболее похожих статей из корпуса arxiv, состоящего примерно из 640K статей по компьютерным наукам! Поиск выполняется с задержкой <50ms на одном ноутбуке! В частности, в этом посте я рассмотрю
- Настройку среды и загрузку данных arXiv из Kaggle
- Загрузку данных в Python с использованием Dask
- Реализацию приложения для поиска семантического сходства научных статей с использованием векторной базы данных Milvus
Методы, использованные в этом посте, можно применять как шаблон для создания любой поисковой системы семантического сходства на основе NLP, а не только для научных статей. Единственным отличием будет используемая предобученная модель.
Для этого поста мы будем использовать arXiv Dataset from Kaggle, который авторы выпустили под лицензией CC0: Public Domain.
Я изложил соображения по векторному поиску сходства в производственном масштабе в моем предыдущем посте. Все эти соображения применимы и к этому проекту. Векторная база данных Milvus настолько хорошо спроектирована, что многие шаги полностью совпадают и воспроизведены здесь только для полноты.
Настройка среды и загрузка данных arxiv из Kaggle.
Cornel University загрузил весь корпус arXiv в Kaggle dataset и лицензировал его под лицензией CC0: Public Domain. Мы можем напрямую загрузить набор данных с помощью Kaggle API. Если вы еще этого не сделали, настройте Kaggle API в своей системе, следуя этим инструкциям.
В этом посте мы будем использовать conda-среду под названием semantic_similarity. Если conda еще не установлена в вашей системе, вы можете сделать это, установив open-sourced mini forge из ее GitHub repository. Приведенные ниже шаги создают необходимые директории и conda-среду, устанавливают требуемые Python-библиотеки и загружают 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
Загрузка данных в Python с использованием Dask
Данные, которые мы скачали с Kaggle, представляют собой JSON-файл размером 3,3 ГБ, содержащий около 2 миллионов статей! Чтобы эффективно обрабатывать такой большой датасет, не стоит загружать весь датасет в память с помощью pandas. Вместо этого мы можем использовать Dask, чтобы разбить данные на несколько разделов и загружать в память лишь несколько разделов в любой момент времени.
Dask
Dask — это библиотека с открытым исходным кодом, которая позволяет нам легко применять параллельные вычисления с API, похожим на pandas. Ее просто настроить на локальной машине, выполнив pip install dask[complete], как показано в разделе настройки. Давайте начнем с импорта необходимых библиотек.
import dask.bag as db
import json
from datetime import datetime
import time
data_path = '../data/arxiv-metadata-oai-snapshot.json'
Мы будем использовать два компонента Dask для эффективной обработки большого JSON-файла arxiv.
- Dask Bag: он позволяет нам загружать JSON-файл блоками фиксированного размера и запускать некоторые функции предварительной обработки для каждой строки данных.
- Dask DataFrame: мы можем преобразовать dask bag в dask dataframe, чтобы получить доступ к API, похожим на pandas
Шаг 1: Загрузите JSON-файл в Dask bag
Давайте загрузим JSON-файл в dask bag, где каждый блок имеет размер 10 МБ. Вы можете настроить аргумент blocksize, чтобы контролировать, насколько большим должен быть каждый блок. Затем мы применяем функцию json.loads к каждой строке dask bag с помощью функции .map(), чтобы разобрать JSON-строку в словарь 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)
Изображение автора
Шаг 2: Напишите вспомогательные функции предварительной обработки
Из вывода мы видим, что каждая строка содержит несколько метаданных, связанных со статьей. Давайте напишем три вспомогательные функции, которые помогут нам предварительно обработать датасет.
v1_date(): эта функция нужна для извлечения даты, когда авторы загрузили первую версию статьи на arXiv. Мы преобразуем дату в UNIX-время и сохраним ее как новое поле в этой строке.text_col(): эта функция нужна для объединения полей “title” и “abstract” с использованием токена “[SEP]”, чтобы мы могли передать эти тексты в модель эмбеддингов SPECTRE. Подробнее о SPECTRE мы поговорим в следующем разделе.filters(): эта функция оставляет только строки, соответствующие некоторым критериям, таким как максимальная длина текста в различных столбцах и статьи в категории 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
)
Шаг 3: Запустите вспомогательные функции предварительной обработки для Dask bag
Мы можем легко использовать функции .map() и .filter(), чтобы запускать вспомогательные функции для каждой строки Dask bag, как показано ниже. Поскольку Dask поддерживает цепочки методов, мы используем эту возможность, чтобы оставить в нашем Dask bag только несколько важных столбцов и удалить остальные.
# 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)
Изображение автора
Шаг 4: Преобразуйте Dask Bag в Dask DataFrame
Последний шаг загрузки данных — преобразовать Dask Bag в Dask Dataframe, чтобы использовать API, похожие на pandas, для каждого блока или раздела данных.
# 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()
Изображение автора
Реализация приложения семантического поиска похожих научных статей с использованием векторной базы данных Milvus
Milvus — одна из самых популярных векторных баз данных с открытым исходным кодом, созданная для высокомасштабируемого и молниеносного поиска по сходству. В этой статье мы будем использовать Milvus Standalone, поскольку запускаем Milvus только на нашем локальном компьютере.
Шаг 1: Установите векторную базу данных Milvus локально
Установка векторной базы данных Milvus с помощью Docker очень проста, поэтому сначала нам нужно установить Docker и Docker Compose. Затем нам нужно лишь скачать docker-compose.yml и запустить docker-контейнеры, как показано в приведенном ниже фрагменте кода! Сайт milvus.io предоставляет множество других вариантов установки как Milvus standalone, так и Milvus Cluster; пожалуйста, ознакомьтесь с ними, если вам нужно установить его в Kubernetes-кластере или установить его офлайн.
# 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
Шаг 2: Создайте коллекцию Milvus
Теперь, когда сервер векторной базы данных Milvus запущен на нашем локальном компьютере, мы можем взаимодействовать с ним с помощью библиотеки pymilvus. Сначала импортируем необходимые модули и подключимся к серверу Milvus, работающему на localhost. Вы можете свободно изменить параметры alias и collection_name. Модель, которую мы используем для преобразования нашего текста в эмбеддинги, определяет значение параметра emb_dim. В случае SPECTRE эмбеддинги имеют размерность 768.
# Убедитесь, что сервер Milvus уже запущен
from pymilvus import connections, utility
from pymilvus import Collection, CollectionSchema, FieldSchema, DataType
# Подключение к серверу Milvus
connections.connect(alias="default", host="localhost", port="19530")
# Имя коллекции
collection_name = "arxiv"
# Размерность эмбеддинга
emb_dim = 768
# # Проверка существующей коллекции и удаление, если она существует
# if utility.has_collection(collection_name):
# print(utility.list_collections())
# utility.drop_collection(collection_name)
При желании вы можете проверить, присутствует ли коллекция, указанная в collection_name, на вашем сервере Milvus. В этом примере, если коллекция уже доступна, я удаляю ее. Но на производственном сервере вы бы этого не делали, а вместо этого пропустили бы приведенный ниже код создания коллекции.
Коллекция Milvus аналогична таблице в традиционной базе данных. Чтобы создать коллекцию для хранения данных, сначала нужно указать schema коллекции. В этом примере мы используем возможность Milvus 2.1 хранить строковые индексы и поля, чтобы хранить все необходимые метаданные, связанные с каждой статьей. Первичный ключ idx и другие поля categories, title, abstract имеют тип данных VARCHAR с разумными максимальными длинами, а embedding — это поле FLOAT_VECTOR, содержащее эмбеддинги размерности emb_dim. Milvus поддерживает широкий спектр типов данных, как показано на странице нашей документации.
# Создание схемы для коллекции
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 = [idx, categories, title, abstract, unix_time, embedding]
schema = CollectionSchema(
fields=fields, description="Semantic Similarity of Scientific Papers"
)
# Создание коллекции со схемой
collection = Collection(
name=collection_name, schema=schema, using="default", shards_num=10
)
После создания коллекции мы готовы загрузить в нее наши тексты и векторы.
Шаг 3: Итерируйте по разделам нашего dataframe Dask, преобразуйте тексты в эмбеддинги с помощью SPECTER и загрузите их в векторную базу данных Milvus
Сначала нам нужно преобразовать тексты в dataframe Dask в вектор эмбеддинга, чтобы выполнить поиск семантического сходства. Мой пост ниже рассказывает, как мы можем преобразовывать тексты в эмбеддинги. В частности, мы будем использовать модель SBERT Bi-Encoder под названием SPECTRE для преобразования научных статей в эмбеддинги.
SPECTER [статья] [Github]: Scientific Paper Embeddings using Citation-informed TransformERs — это модель для преобразования научных статей в эмбеддинги.
- Тексты Title и Abstract каждой статьи объединяются с токеном [SEP] и преобразуются в эмбеддинги с использованием токена [CLS] предобученной Transformer-модели (SciBERT).
- Используйте цитирования как прокси-сигнал связанности между документами. Если одна статья цитирует другую, мы можем предположить, что они связаны.
- Целевая функция обучения triplet loss: мы обучаем Transformer-модель так, чтобы статьи с общими цитированиями находились ближе в пространстве эмбеддингов.
- Другими словами, Positive paper — это статья, процитированная в Query paper, а Negative paper — это статья, не процитированная в Query paper. Случайно выбранные негативные примеры — это «простые» негативные примеры.
- Чтобы улучшить производительность, мы создаем «сложные» негативные примеры, используя статьи, которые НЕ цитируются в Query paper, но цитируются в Positive paper.
- Во время инференса нам нужны только Title и Abstract. Цитирования не требуются, поэтому SPECTER может создавать эмбеддинги даже для новых статей, у которых еще нет цитирований!
- SPECTER обеспечивает отличную производительность (лучше, чем SciBERT) в классификации тем, предсказании цитирований и рекомендации Scientific Papers.
Изображение автора с использованием скриншотов из open-sourced статьи SPECTER
Использовать предобученную модель SPECTRE просто с библиотекой Sentence Transformer. Мы можем скачать предобученную модель всего одной строкой кода, как показано ниже. Мы также пишем простую вспомогательную функцию для преобразования целого столбца текстов из раздела Dask dataframe в эмбеддинги.
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()
Нам нужно пройтись по разделам Dask dataframe, чтобы загрузить данные в нашу коллекцию Milvus. Во время каждой итерации мы загружаем в память только строки из этого раздела и добавляем данные из столбцов метаданных в переменную data. Мы можем использовать API dask .map_partitions() для применения генерации эмбеддингов к каждой строке в разделе и добавления результатов обратно в ту же переменную data. Наконец, мы можем загрузить данные в Milvus с помощью 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)
Пожалуйста, обратите внимание, что порядок столбцов, добавляемых в переменную data, должен соответствовать тому же порядку, что и переменная fields, которую мы определили при создании схемы!
Шаг 4: Создайте индекс Approximate Nearest Neighbors (ANN) для загруженных данных
После того как мы вставим все эмбеддинги в векторную базу данных Milvus, нам нужно создать ANN-индекс, чтобы ускорить поиск. В этом примере я использую тип индекса HNSW, один из самых быстрых и точных ANN-индексов. Посмотрите документацию Milvus, чтобы получить больше информации об индексе HNSW и его параметрах.
# 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)
Шаг 5: Запустите свои запросы Vector Similarity Search!
Наконец, данные в нашей коллекции Milvus готовы к запросам. Сначала мы должны загрузить коллекцию в память, чтобы выполнять к ней запросы.
# Load the collection into memory
collection = Collection(collection_name)
collection.load()
Далее я создал простую вспомогательную функцию, которая принимает query_text, преобразует его в embedding SPECTRE, выполняет ANN-поиск по коллекции Milvus и выводит результаты. Мы можем управлять качеством и скоростью поиска с помощью search_params, описанных на странице документации 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')}")
Теперь мы можем использовать вспомогательную функцию всего одной строкой кода, чтобы выполнить семантический поиск статей arXiv по всем ~640K статьям по Computer Science, хранящимся в нашей коллекции Milvus. Например, я ищу статьи, похожие на статью SimCSE, которую я подробно обсуждал в моем предыдущем посте. Топ-10 результатов вполне релевантны моему поисковому запросу, поскольку они в основном связаны с contrastive learning sentence embeddings! Еще более впечатляет то, что весь поиск занял всего 30 мс при запуске на моем ноутбуке, что вполне укладывается в типичные требования большинства приложений!
# 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)
Изображение автора
Если нам больше не нужно выполнять запросы, мы можем освободить коллекцию, чтобы освободить память нашей машины. Удаление коллекции из памяти не приводит к потере данных, поскольку она по-прежнему хранится на нашем диске и может быть снова загружена при необходимости.
# Release the collection from memory when it's not needed anymore
collection.release()
Если вы хотите остановить сервер Milvus и удалить все данные с диска, вы можете следовать инструкциям по остановке Milvus. Осторожно! Эта операция необратима и удалит все данные в вашем кластере Milvus.
Заключение
В этом посте мы реализовали ультрамасштабируемый сервис семантического поиска научных статей с использованием эмбеддингов SPECTRE и векторной базы данных Milvus за несколько простых шагов. Этот подход масштабируется в продакшене до сотен миллионов или даже миллиардов векторов. Мы протестировали поиск с помощью примера запроса по статье, который вернул 10 лучших результатов всего за 30 мс! Репутация Milvus как высокомасштабируемой и молниеносно быстрой базы данных для поиска векторного сходства вполне заслужена!
Для большего вдохновения по применениям Milvus перейдите к Milvus демо векторной базы данных demos и Bootcamp.
Читать далее
Cosmos World Foundation Model Platform for Physical AI
NVIDIA's Cosmos platform enables safe, digital twin training of GenAI models for physical applications, overcoming data scarcity and safety challenges.
Proactive Monitoring for Vector Database: Zilliz Cloud Integrates with Datadog
we're excited to announce Zilliz Cloud's integration with Datadog, enabling comprehensive monitoring and observability for your vectorDB deployments.
DeepRAG: Thinking to Retrieval Step by Step for Large Language Models
Discover DeepRAG, an advanced retrieval-augmented generation (RAG) model that improves LLM accuracy by retrieving only essential data through step-by-step reasoning.