Vector Search on Hugging Face with the Hub as Backend

Datasets on the Hugging Face Hub rely on parquet files. We can interact with these files using DuckDB as a fast in-memory database system. One of DuckDB's features is vector similarity search which can be used with or without an index.

Install dependencies

Create embeddings for the dataset

First, we need to create embeddings for the dataset to search over. We will use the sentence-transformers library to create embeddings for the dataset.

Now, let's load the ai-blueprint/fineweb-bbc-news dataset from the Hub.

We can now create embeddings for the dataset. Normally, we might want to chunk our data into smaller batches to avoid losing precision, but for this example, we will just create embeddings for the full text of the dataset.

We can now upload our dataset with embeddings back to the Hub.

Vector Search the Hugging Face Hub

We can now perform vector search on the dataset using duckdb. When doing so, we can either use an index or not. Searching without an index is slower but more precise, whereas searching with an index is faster but less precise.

Without an index

To search without an index, we can use the duckdb library to connect to the dataset and perform a vector search. This is a slow operation, but normally works quick enough for small datasets up to let's say 100k rows. Meaning querying our dataset will be somewhat slower.

With an index

This approach creates a local copy of the dataset and uses this to create an index. This has some minor overhead but it will significantly speed up the search once you've created it.

Now we can perform a vector search with the index, which return the results instantly.

The query reduces from 30 seconds to sub-second response times and does not require you to deploy a heavy-weight vector search engine, while storage is handled by the Hub.

Conclusion

We have seen how to perform vector search on the Hub using duckdb. For small datasets <100k rows, we can perform vector search without an index using the Hub as a vector search backend, but for larger datasets, we should create an index with the vss extension while doing local search and using the Hub as a storage backend.

Learn more

• Vector Search on Hugging Face
• Vector Search Indexing with DuckDB