# MR-RAG: Multimodal Relevance-Aware Retrieval-Augmented Generation for Medical Visual Question Answering

This repository contains the source code for our paper **"MR-RAG: Multimodal Relevance-Aware Retrieval-Augmented Generation for Medical Visual Question Answering"**, submitted to AAAI-2026.

We provide both the **retrieval module** and the **generation module**, with scripts and data to reproduce key experimental results.


## 📁 Directory Structure

```
MR-RAG/
├── README.md                  # Project overview and usage instructions
├── requirements.txt           # Python dependencies
│
├── retrieve/                  # 🔍 Retrieval module
│   └── src/
│       └── retrieve.py        # Main script for retrieval
│
├── llava/                     # 🧠 Generation module
│   └── eval/
│       └── model_vqa.py       # Main script for multimodal generation
│
├── transformers/              # 🤖 Customized transformers
├── utils/                     # ⚙️  Utility functions and hooks
│
├── data/                      # 📦 Sample input data
│   └── test/                  # Example test set
│
├── scripts/                   # 🛠️  Shell scripts for running experiments
│   ├── aug_inference.sh       # Generation with IFA module
│   └── inference.sh           # Generation without IFA module
│
├── answers/                   # 📝 Output predictions (example results)
│   └── example_output.jsonl   # Saved results for sample inputs
```



## 🔧 Installation

We recommend using **conda** to create a clean and reproducible environment:

```bash
# Step 1: Create a new conda environment
conda create -n mrrag python=3.10 -y

# Step 2: Activate the environment
conda activate mrrag

# Step 3: Install required Python packages
pip install -r requirements.txt
```


## 🚀 How to Run

Follow the steps below to reproduce our pipeline:


### 🔍 1. Run the Retrieval Module

```bash
python retrieve/src/retrieve.py
```

This will execute the retrieval pipeline and store the retrieved texts for each **question-image pair** into the folder:`/output_retrieve/`



### 💡 2. Run the Generation Module

```bash
bash scripts/aug_inference.sh
```

This will perform **multimodal answer generation** using the modified LLaVA model located at:`llava/eval/model_vqa_harvard.py`. The answer predictions stored in: `answers/example_output.jsonl`





## 🧠 Description of model_vqa.py

The script `llava/eval/model_vqa.py` serves as the main entry point for the generation module. It performs multimodal visual question answering (VQA) using a modified LLaVA model. Key components and logic include:

- **Model and Tokenizer Loading**: Loads a pretrained LLaVA model along with tokenizer, image processor, and conversation templates.
- **Hook Injection**: Uses a custom `AttentionHook` class (defined in `utils/hook.py`) to register forward hooks into specific transformer layers. These hooks allow relevance-aware augmentation by dynamically modulating attention based on retrieved text tokens.
- **Input Construction**: For each test instance, constructs a multi-turn dialogue prompt that incorporates both the question and its associated retrieved reports, following the LLaVA input format.
- **Image Processing**: Loads and processes the fundus image corresponding to each question using the provided `image_processor`.
- **Autoregressive Generation**: Generates answers using the `generate()` method, with optional sampling or beam search parameters.
- **Output Formatting**: Saves each prediction to a JSONL file containing metadata such as `question_id`, `answer`, `gt_answer`, and `image`.

The script supports flexible configurations via command-line arguments, such as setting the augmentation range (`--aug-layer1`, `--aug-layer2`), number of top-K retrieved documents (`--real-k`), and generation parameters (e.g., `--temperature`, `--num_beams`).

An example execution is provided in `scripts/aug_inference.sh`.



## 📄 Example Input & Output for generation
**Input: **Located in `output_retrieve/example_retrieval.jsonl`, each sample includes question, image path, answer and reference reports.
**Output: **Example predictions saved in `answers/example_output.jsonl`.




## 📜 License
This code is released for academic use only.

