WARP: An Efficient Engine for Contextualized Multi-Vector Retrieval

---

WARP is an extremely fast and accurate retrieval engine based on Stanford's ColBERTv2/PLAID and Google DeepMind's XTR retrieval systems.

Installation

WARP requires Python 3.7+ and PyTorch 1.9+ and uses the Hugging Face Transformers library. In addition, WARP supports the following runtimes for inference:

• ONNX
• OpenVino
• Core ML (macOS only)

It is strongly recommended to create a conda environment using the commands below.

We include an environment file specifically for CPU-only environments (conda_env_cpu.yml) and one for GPU environments (conda_env.yml).

conda env create -f conda_env[_cpu].yml
conda activate warp

[!NOTE] While WARP's retrieval process is heavily optimized for CPU, it is still strongly recommended to use a GPU for index construction.

Environment Setup

To construct indexes and perform retrieval, define the following environment variables in a .env file in the repository root:

INDEX_ROOT=...
EXPERIMENT_ROOT=...

BEIR_COLLECTION_PATH=...
LOTTE_COLLECTION_PATH=...

• INDEX_ROOT: Specifies the on-disk location for indexes.
• EXPERIMENT_ROOT: Specifies the on-disk location for experiment files.
• BEIR_COLLECTION_PATH: Designates the path to the datasets of the BEIR Benchmark.
• LOTTE_COLLECTION_PATH: Specifies the path to the LoTTE dataset.

Dataset Setup

BEIR Benchmark

To download and extract a dataset from the BEIR Benchmark:

python utility/extract_collection.py -d ${dataset} -i "${BEIR_COLLECTION_PATH}" -s test

Replace ${dataset} with the desired dataset name as specified here.

LoTTE Dataset

1. Download the LoTTE dataset files from here.
2. Extract the files manually to the directory specified in LOTTE_COLLECTION_PATH.

[!NOTE] If you face any problems, please feel free to open a new issue.

Branches

• main: Stable branch with XTR/WARP.

Bugs

If you experience bugs, or have suggestions for improvements, please use the issue tracker to report them.

---

We provide code to reproduce the baseline evaluations for XTR and ColBERTv2/PLAID.

[!TIP] We provide scripts to reproduce all of these measurements on Google Cloud. The scripts can be found here.