Загрузка и поиск по наборам данных Hugging Face с помощью Vector DB#

Платформа Hugging Face предоставляет возможность делиться и использовать модели машинного обучения и наборы данных. Platform V Vector DB (далее - Vector DB) также публикует наборы данных совместно с вложениями, которые можно использовать для практики работы с Vector DB и создания приложений на основе семантического поиска.

arxiv-titles-instructorxl-embeddings#

Набор данных содержит вложения, созданные только из названий статей. Каждый вектор имеет полезную нагрузку с названием, использованным для его создания, а также DOI (идентификатор цифрового объекта).

{
    "title": "Nash Social Welfare for Indivisible Items under Separable, Piecewise-Linear Concave Utilities",
    "DOI": "1612.05191"
}

Загрузка набора данных так же проста, как использование функции load_dataset из библиотеки datasets:

from datasets import load_dataset

dataset = load_dataset("Qdrant/arxiv-titles-instructorxl-embeddings")

Набор данных занимает более 16 ГБ, поэтому загрузка может занять некоторое время.

В наборе данных содержится 2 250 000 векторов. Вот как можно проверить список функций в наборе данных:

dataset.features

Потоковая передача набора данных#

Потоковая передача набора данных позволяет работать с набором данных без необходимости его скачивания. Данные передаются потоком при итерации по набору данных. Подробнее о Hugging Face описано в оригинальной документации.

from datasets import load_dataset

dataset = load_dataset(
    "Qdrant/arxiv-titles-instructorxl-embeddings", split="train", streaming=True
)

Загрузка набора данных в Vector DB#

Набор данных может быть загружен в Vector DB с использованием Python SDK. Вложения уже предварительно вычислены, поэтому можно сохранить их в коллекции, которую необходимо создать:

from qdrant_client import QdrantClient, models

client = QdrantClient("http://localhost:6333")

client.create_collection(
    collection_name="arxiv-titles-instructorxl-embeddings",
    vectors_config=models.VectorParams(
        size=768,
        distance=models.Distance.COSINE,
    ),
)

Всегда полезно использовать пакетную обработку при загрузке большого набора данных. Понадобится вспомогательная функция для разделения набора данных на пакеты:

from itertools import islice

def batched(iterable, n):
    iterator = iter(iterable)
    while batch := list(islice(iterator, n)):
        yield batch

Пользователи Python версии 3.12+ могут использовать функцию batched из пакета itertools вместо нее.

Независимо от используемой версии Python можно использовать метод upsert, чтобы загрузить набор данных партиями в Vector DB:

batch_size = 100

for batch in batched(dataset, batch_size):
    ids = [point.pop("id") for point in batch]
    vectors = [point.pop("vector") for point in batch]

    client.upsert(
        collection_name="arxiv-titles-instructorxl-embeddings",
        points=models.Batch(
            ids=ids,
            vectors=vectors,
            payloads=batch,
        ),
    )