OPEN-RAG: Experimentation and Expansion

Assignment 3 — FAST-NUCES | NLP Course 2026 | Dr. Zohair Ahmed

Team: Ayesha Shabbir · Hifza Hussain · Saad bin Amir

---

Overview

This project reproduces and extends OPEN-RAG — an adaptive retrieval-augmented generation framework that teaches a language model to decide when to retrieve external knowledge rather than always retrieving or never retrieving.

The model (shayekh/openrag_llama2_7b_8x135m) is a LLaMA-2-7B fine-tuned with a Mixture-of-Experts adapter. At inference time it produces a retrieval probability P(retrieve). If P(retrieve) ≥ threshold γ, the model fetches external passages; otherwise it answers from parametric memory.

Assignment 2 reproduced the baseline on 2WikiMultihopQA.
Assignment 3 (this repo) adds:

• A new dataset evaluation: HotpotQA
• A proposed method: Threshold calibration via confidence stratification analysis
• Cross-dataset comparison and full visualisations

---

Results

OPEN-RAG Performance: Paper vs Reproduced

Threshold Sweep (2WikiMultihopQA)

Finding: For multi-hop questions a lower threshold (more aggressive retrieval) consistently helps. The model's default threshold of 0.5 is conservative for this task.

Retrieval Confidence Stratification (2WikiMultihopQA)

Samples where the model is confident about retrieving achieve substantially higher EM (54% vs 36%), validating the adaptive retrieval design.

---

Repository Structure

project-root/
├── README.md
├── requirements.txt
├── train.py               ← evaluation runner (calls OPEN-RAG repo scripts)
├── inference.py           ← single-question inference CLI
├── config.yaml            ← all experiment parameters
│
├── data/
│   └── sample_data.csv    ← 7 representative QA samples with metadata
│
├── notebooks/
│   └── 01_inference_demo.ipynb   ← live Gradio demo (run on Colab)
│
├── src/
│   ├── model.py           ← model loading utilities
│   ├── dataset.py         ← dataset loading helpers
│   └── utils.py           ← EM metric, threshold sweep, confidence analysis
│
├── results/
│   ├── baseline_metrics.json    ← 2WikiMultihopQA baseline (A2)
│   ├── improved_metrics.json    ← HotpotQA + threshold sweep (A3)
│   └── training_log.csv         ← all experiment runs in one table
│
└── checkpoints/
    └── README.md          ← model weights are on HuggingFace (see below)

---

Setup

1. Clone this repo and the OPEN-RAG repo

git clone <this-repo-url>
cd <this-repo>

git clone https://github.com/ShayekhBinIslam/openrag.git

2. Install dependencies

pip install -r requirements.txt

GPU required. The model needs ~4–5 GB VRAM in 4-bit mode. We ran all experiments on a Google Colab A100.

3. Model weights

No local weights are stored. The model downloads automatically from HuggingFace on first use:

Model: shayekh/openrag_llama2_7b_8x135m
URL  : https://huggingface.co/shayekh/openrag_llama2_7b_8x135m

---

Usage

Run evaluation

# Baseline (2WikiMultihopQA, threshold=0.5)
python train.py --task 2wikimultihopqa --threshold 0.5 --max_samples 101

# HotpotQA
python train.py --task hotpotqa --threshold 0.5 --max_samples 100

# Threshold sweep
python train.py --task 2wikimultihopqa --threshold 0.3 --max_samples 100
python train.py --task 2wikimultihopqa --threshold 0.7 --max_samples 100

Single-question inference

python inference.py --question "Were Scott Derrickson and Ed Wood of the same nationality?"

# Interactive mode
python inference.py --interactive

Live demo

Open notebooks/01_inference_demo.ipynb in Google Colab and run all cells. The notebook launches a Gradio interface with pre-computed results and real inference capability.

---

Proposed Method: Confidence-Stratified Threshold Selection

Motivation: The default threshold γ=0.5 is a single flat value applied to all questions. Our analysis shows that the model's retrieval confidence distribution varies across question types and datasets.

Method: Instead of a fixed threshold, we propose selecting γ based on the variance of P(retrieve) across a validation set:

• High-variance datasets (mixed questions) → lower γ (retrieve more)
• Low-variance datasets (model consistently confident) → higher γ

Evidence from our experiments:

• Post-hoc fine-grained sweep (γ ∈ [0.1, 0.9] step 0.05) peaks at γ=0.75 with 54.36% EM
• Samples with P(retrieve) > 0.7 achieve 54% EM vs 36% for low-confidence samples
• Lowering γ from 0.5 to 0.3 improves EM from 48.5% to 52.0% (+3.5pp)

---

References

• Shi, H. et al. (2024). OPEN-RAG: Enhanced Retrieval-Augmented Reasoning with Open-Source Large Language Models. EMNLP 2024 Findings.
• HuggingFace model: shayekh/openrag_llama2_7b_8x135m
• Original codebase: ShayekhBinIslam/openrag