AutoFaiss

Automatically create Faiss knn indices with the most optimal similarity search parameters.

It selects the best indexing parameters to achieve the highest recalls given memory and query speed constraints.

Doc and posts and notebooks

Using faiss efficient indices, binary search, and heuristics, Autofaiss makes it possible to automatically build in 3 hours a large (200 million vectors, 1TB) KNN index in a low amount of memory (15 GB) with latency in milliseconds (10ms).

Get started by running this colab notebook, then check the full documentation.
Get some insights on the automatic index selection function with this colab notebook.

Then you can check our multimodal search example (using OpenAI Clip model).

Read the medium post to learn more about it!

Installation

To install run pip install autofaiss

It's probably best to create a virtual env:

python -m venv .venv/autofaiss_env
source .venv/autofaiss_env/bin/activate
pip install -U pip
pip install autofaiss

Using autofaiss in python

If you want to use autofaiss directly from python, check the API documentation and the examples

In particular you can use autofaiss with on memory or on disk embeddings collections:

Using in memory numpy arrays

If you have a few embeddings, you can use autofaiss with in memory numpy arrays:

from autofaiss import build_index
import numpy as np

embeddings = np.float32(np.random.rand(100, 512))
index, index_infos = build_index(embeddings, save_on_disk=False)

query = np.float32(np.random.rand(1, 512))
_, I = index.search(query, 1)
print(I)

Using numpy arrays saved as .npy files

If you have many embeddings file, it is preferred to save them on disk as .npy files then use autofaiss like this:

from autofaiss import build_index

build_index(embeddings="embeddings", index_path="my_index_folder/knn.index",
            index_infos_path="my_index_folder/index_infos.json", max_index_memory_usage="4G",
            current_memory_available="4G")

Memory-mapped indices

Faiss makes it possible to use memory-mapped indices. This is useful when you don't need a fast search time (>50ms) and still want to reduce the memory footprint to the minimum.

We provide the should_be_memory_mappable boolean in build_index function to generate memory-mapped indices only. Note: Only IVF indices can be memory-mapped in faiss, so the output index will be a IVF index.

To load an index in memory mapping mode, use the following code:

import faiss
faiss.read_index("my_index_folder/knn.index", faiss.IO_FLAG_MMAP | faiss.IO_FLAG_READ_ONLY)

You can have a look to the examples to see how to use it.

Technical note: You can create a direct map on IVF indices with index.make_direct_map() (or directly from the build_index function by passing the make_direct_map boolean). Doing this speeds up a lot the .reconstruct() method, function that gives you the value of one of your vector given its rank. However, this mapping will be stored in RAM... We advise you to create your own direct map in a memory-mapped numpy array and then call .reconstruct_from_offset() with your custom direct_map.

Using autofaiss with pyspark

Autofaiss allows you to build indices with Spark for the following two use cases:

• To build a big index in a distributed way
• Given a partitioned dataset of embeddings, building one index per partition in parallel and in a distributed way.

Prerequisities:

1. Install pyspark: pip install pyspark.
2. Prepare your embeddings files (partitioned or not).
3. Create a Spark session before calling autofaiss. If no Spark session exists, a default session will be creaed with a minimum configuration.

Creating a big index in a distributed way

See distributed_autofaiss.md for a complete guide.

It is possible to generate an index that would require more memory than what's available. To do so, you can control the number of index splits that will compose your index with nb_indices_to_keep. For example, if nb_indices_to_keep is 10 and index_path is knn.index, the final index will be decomposed into 10 smaller indexes:

• knn.index01
• knn.index02
• knn.index03
• ...
• knn.index10

A concrete example shows how to produce N indices and how to use them.

Creating partitioned indexes

Given a partitioned dataset of embeddings, it is possible to create one index per partition by calling the method build_partitioned_indexes.

See this example that shows how to create partitioned indexes.

Using the command line

Create embeddings

import os
import numpy as np
embeddings = np.random.rand(1000, 100)
os.mkdir("embeddings")
np.save("embeddings/part1.npy", embeddings)
os.mkdir("my_index_folder")

Generate a Knn index

autofaiss build_index --embeddings="embeddings" --index_path="my_index_folder/knn.index" --index_infos_path="my_index_folder/index_infos.json" --metric_type="ip"

Try the index

import faiss
import glob
import numpy as np

my_index = faiss.read_index(glob.glob("my_index_folder/*.index")[0])

query_vector = np.float32(np.random.rand(1, 100))
k = 5
distances, indices = my_index.search(query_vector, k)

print(list(zip(distances[0], indices[0])))

How are indices selected ?

To understand better why indices are selected and what are their characteristics, check the index selection demo

Command quick overview

Quick description of the autofaiss build_index command:

embeddings -> Source path of the embeddings in numpy.
index_path -> Destination path of the created index. index_infos_path -> Destination path of the index infos. save_on_disk -> Save the index on the disk. metric_type -> Similarity distance for the queries.

index_key -> (optional) Describe the index to build.
index_param -> (optional) Describe the hyperparameters of the index.
current_memory_available -> (optional) Describe the amount of memory available on the machine.
use_gpu -> (optional) Whether to use GPU or not (not tested).

Command details

The autofaiss build_index command takes the following parameters:

Install from source

First, create a virtual env and install dependencies:

python3 -m venv .env
source .env/bin/activate
make install

python -m pytest -x -s -v tests -k "test_get_optimal_hyperparameters" to run a specific test