Load Hugging Face Datasets into Qdrant

Hugging Face provides a platform for sharing and using ML models and datasets. Qdrant also publishes datasets along with the embeddings that you can use to practice with Qdrant and build your applications based on semantic search. Please let us know if you’d like to see a specific dataset!

Loading your own dataset changes the parts worth getting right. The Load a dataset into Qdrant prompt has an agent pick streaming, batch size, and indexing settings from your record count and shape.

Help me load my dataset into Qdrant. First, read https://skills.qdrant.tech/qdrant-performance-optimization/indexing-performance-optimization/SKILL.md

Before writing any upload code, ask me where the data lives, how many records it contains, whether it already has embeddings, and which fields I plan to filter on.

If the data has precomputed vectors, confirm their dimensions and intended distance metric. Don’t assume they match my needs. If they don’t, explain what re-embedding would cost.

Help me decide whether to stream the data or download it first. Choose a batch size based on the size of my records, and explain the choice.

Show me what to configure before the bulk load, including the indexing threshold and any payload indexes to create up front. Explain which settings to restore afterward and when to restore them.

Finish with checks I can use to confirm that all records were loaded and the collection is ready to query. Make clear how to tell the difference between an upload that has finished and a collection that has finished indexing.

arxiv-titles-instructorxl-embeddings

This dataset contains embeddings generated from the paper titles only. Each vector has a payload with the title used to create it, along with the DOI (Digital Object Identifier).

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

You can find a detailed description of the dataset in the Practice Datasets section. If you prefer loading the dataset from a Qdrant snapshot, it also linked there.

Loading the dataset is as simple as using the load_dataset function from the datasets library:

from datasets import load_dataset

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

The dataset contains 2,250,000 vectors. This is how you can check the list of the features in the dataset:

dataset.features

Streaming the dataset

Dataset streaming lets you work with a dataset without downloading it. The data is streamed as you iterate over the dataset. You can read more about it in the Hugging Face documentation.

from datasets import load_dataset

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

Loading the dataset into Qdrant

You can load the dataset into Qdrant using the Python SDK. The embeddings are already precomputed, so you can store them in a collection, that we’re going to create in a second:

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,
    ),
)

It is always a good idea to use batching, while loading a large dataset, so let’s do that. We are going to need a helper function to split the dataset into batches:

from itertools import islice

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

If you are a happy user of Python 3.12+, you can use the batched function from the itertools package instead.

No matter what Python version you are using, you can use the upsert method to load the dataset, batch by batch, into Qdrant:

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,
        ),
    )

Your collection is ready to be used for search! Please let us know using Discord if you would like to see more datasets published on Hugging Face hub.