Forward-Looking Active REtrieval augmented generation (FLARE)

This repository contains the code and data for the paper Active Retrieval Augmented Generation.

Overview

FLARE is a generic retrieval-augmented generation method that actively decides when and what to retrieve using a prediction of the upcoming sentence to anticipate future content and utilize it as the query to retrieve relevant documents if it contains low-confidence tokens.

Install environment with Conda

Create a conda env and follow setup.sh to install dependencies.

Quick start

Download Wikipedia dump

Download the Wikipedia dump from the DPR repository using the following command:

wget -O data/dpr/psgs_w100.tsv https://dl.fbaipublicfiles.com/dpr/wikipedia_split/psgs_w100.tsv.gz

Build index over the Wikipedia dump using Elasticsearch

wget -O elasticsearch-7.17.9.tar.gz https://artifacts.elastic.co/downloads/elasticsearch/elasticsearch-7.17.9-linux-x86_64.tar.gz  # download Elasticsearch
tar zxvf elasticsearch-7.17.9.tar.gz
pushd elasticsearch-7.17.9
nohup bin/elasticsearch &  # run Elasticsearch in background
popd
python prep.py --task build_elasticsearch --inp data/dpr/psgs_w100.tsv wikipedia_dpr  # build index

Setup OpenAI keys

Put OpenAI keys in the keys.sh file. Multiple keys can be used to accelerate experiments. Please avoid uploading your keys to Github by accident!

Run FLARE

Use the following command to run FLARE on the 2WikiMultihopQA dataset (500 examples) with text-davinci-003. Be careful, the experiment is relatively expensive because FLARE iteratively calls OpenAI APIs. To save credits, you can set debug=true to active the debugging mode which walks you through the process one example at a time, or you can decrease max_num_examples to run small-scale experiments.

./openai.sh 2wikihop configs/2wikihop_flare_config.json

my-basement/FLARE