LongRAG Framework

You can skip manual corpus preparation by loading the pre-processed retrieval corpora for NQ and HotpotQA directly from Hugging Face using the datasets library.

from datasets import load_dataset
corpus_nq = load_dataset("TIGER-Lab/LongRAG", "nq_corpus")
corpus_hotpotqa = load_dataset("TIGER-Lab/LongRAG", "hotpot_qa_corpus")

LongRAG uses the Tevatron toolkit for dense retrieval. The default base embedding model is bge-large-en-v1.5. You can use the provided sample script to encode your corpus, but ensure you update the parameters for your local dataset path and adjust the GPU count as needed.

sh scripts/run_retrieve_tevatron.sh

To install LongRAG, clone the repository and install the dependencies via pip using the provided requirements.txt file.

git clone https://github.com/TIGER-AI-Lab/LongRAG.git
cd LongRAG
pip install -r requirements.txt

The Long Reader uses long-context LLMs (like GPT-4o or Gemini-1.5-Pro) to extract answers from retrieved units.

Important Setup: You must manually configure your API keys and configurations in the inference utility files:

• For GPT series: utils/gpt_inference.py
• For Gemini series: utils/gemini_inference.py

Supported Models:

• GPT-4o
• GPT-4-Turbo
• Gemini-1.5-Pro
• Claude-3-Opus

Execution:

1. Create an output directory: mkdir -p exp/
2. Run the evaluation script: sh scripts/run_eval_qa.sh
3. Or run via Python directly for specific subsets:

python eval/eval_qa.py \
  --test_data_name "nq" \
  --test_data_split "subset_100" \
  --output_file_path "./exp/nq_gpt4o_100.json" \
  --reader_model "GPT-4o"

If you choose to build your own corpus, the process involves three main steps:

1. Clean Wikipedia raw data: Use the provided script to clean Wikipedia database dumps using WikiExtractor.
2. Preprocess Wikipedia data: Run the processing script to generate metadata and text pickle files (e.g., doc_dict.pickle, doc_size.pickle).
3. Group documents: Construct long retrieval units (over 4K tokens) by grouping related documents to reduce corpus size and improve information completeness.

# 1. Clean Wikipedia
sh scripts/extract_and_clean_wiki_dump.sh

# 2. Preprocess
sh scripts/process_wiki_page.sh

# 3. Group documents
sh scripts/group_documents.sh

The output of the Long Reader is a JSON file where each row represents a test case. The key field for the final prediction is short_ans.

{
    "query_id": "383", 
    "question": "how many episodes of touching evil are there", 
    "answers": ["16"], 
    "long_ans": "16 episodes.", 
    "short_ans": "16", 
    "is_exact_match": 1, 
    "is_substring_match": 1, 
    "is_retrieval": 1
}