wiki-ann

Dataset for Testing Vector DB search with AND queries

Description

This repository provides a dataset built on top of the 35 million english embeddings from https://huggingface.co/datasets/Cohere/wikipedia-22-12-en-embeddings and 5000 randomly sampled embeddings from https://huggingface.co/datasets/Cohere/wikipedia-22-12-simple-embeddings.

From each of the corresponding passages to these embeddings, we identify words found in the 4000 most common words across all passages, and use them as a set of keyword filters for filtered ANNS.

In addition, we provide a random 1M point slice of the dataset. This is to facilitate easier prototyping of algorithms.

Setup

In order to get the original set of base embeddings, we include a combine_base_vecs.py script in each dataset subfolder. Simply run the script to recover the embeddings (be warned: this requires a lot of RAM!). For the full 35M dataset, you will also have to use cat to combine the labels files, which have been split.

The final embeddings file starts with 4 bytes containing the number of points, then another 4 bytes containing the embedding dimension, then the embedding data itself, stored sequentially.

For the base embeddings' labels, the final text file contains a comma-separated list of numbers on each line; for example, line 5 corresponds to the 5th embedding's labels. Each number corresponds to a unique label. For the query embeddings' labels, the final text file contains a &-separated list of labels. For this dataset, each query can only come with two labels.