This repository provides the official code for "Direct Retrieval-augmented Optimization," a framework that synergizes knowledge selection and language model generation through a collaborative, jointly optimized loop. It's designed for researchers and practitioners working on knowledge-grounded LLMs seeking improved factuality, relevance, and performance.

How It Works

The Direct Retrieval-augmented Optimization (DRO) framework integrates a knowledge selector and a generator into a single feedback loop. This allows them to learn from each other, optimizing jointly for better performance. The process involves an E-step for permutation sampling and a subsequent M-step for maximization, effectively fine-tuning both components in tandem.

Quick Start & Requirements

• Installation: Create a Conda environment (conda create -n rag python=3.10), activate it (conda activate dro), and install dependencies (pip install -r requirements.txt, pip install pytrec_eval).
• Prerequisites: Python 3.10, Conda, pytrec_eval. Optional: wandb for logging, vllm for enhanced inference (requires setting VLLM_WORKER_MULTIPROC_METHOD=spawn).
• Data: Requires downloading various QA datasets (Nature Question, HotpotQA, MuSiQue, 2WikiMultihopQA, Wizard of Wikipedia) and a retrieval corpus (Wikipedia, compatible with ColBERT or DPR). Pre-processed datasets are available via Google Drive links.
• Models: Requires specifying paths for a ranker model and a generator LLM.

Highlighted Details

• Supports multiple datasets and benchmarks for open-domain QA, multi-hop QA, and knowledge-grounded dialogue.
• Employs a two-stage training process: warm-up supervised fine-tuning for both selector and generator, followed by the DRO optimization loop.
• Evaluation uses EM and F1 metrics via KILT evaluation scripts.
• Provides pre-processed data for evaluation and training across several benchmarks.

Maintenance & Community

The project acknowledges prior work from RankGPT and Llama-Factory. Further community engagement details (Discord/Slack, roadmap) are not explicitly provided in the README.

Licensing & Compatibility

The README does not specify a license. Compatibility for commercial use or closed-source linking is not mentioned.

Limitations & Caveats

The README indicates that the paper is "coming soon" and lists "add the command for evaluation" as a TODO item, suggesting the project may still be under active development or refinement. Specific hardware requirements beyond standard Python environments are not detailed.