MODE: Mixture of Document Experts for RAG

September 3, 2025 Β· View on GitHub

Project Overview

MODE (Mixture of Document Experts) is an advanced framework that improves Retrieval-Augmented Generation (RAG) by integrating external knowledge retrieval with a mixture of specialized expert models.

Key features of MODE include:

  • Hierarchical Clustering: Organizes documents into semantically meaningful clusters.
  • Expert Models: Assigns specialized models to different document clusters for targeted expertise.
  • Centroid-Based Retrieval: Selects representative documents efficiently to enhance retrieval relevance.

By combining these techniques, MODE delivers more accurate document retrieval and synthesis for query-based applications, improving answer quality while reducing retrieval noise. MODE is particularly well-suited for small to medium-sized document collections or datasets.

πŸ“„ arxiv: https://arxiv.org/abs/2509.00100
πŸ“„ Docs: https://mode-rag.readthedocs.io/en/latest/
🌐 Website: https://mode-rag.netlify.app/

πŸ“‚ Project Structure

.
β”œβ”€β”€ benchmarking              # Evaluation and benchmarking scripts
β”‚Β Β  β”œβ”€β”€ data.py
β”‚Β Β  β”œβ”€β”€ eval
β”‚Β Β  β”‚Β Β  β”œβ”€β”€ db
β”‚Β Β  β”‚Β Β  └── logs
β”‚Β Β  β”‚Β Β      β”œβ”€β”€ ours
β”‚Β Β  β”‚Β Β      └── traditional_rag
β”‚Β Β  β”œβ”€β”€ evaluate.py
β”‚Β Β  β”œβ”€β”€ metric_to_json.py
β”‚Β Β  β”œβ”€β”€ mode.py
β”‚Β Β  └── traditional_rag.py
β”œβ”€β”€ README.md
β”œβ”€β”€ requirements.txt     
β”œβ”€β”€ src                         
β”‚Β Β  β”œβ”€β”€ __init__.py
β”‚Β Β  β”œβ”€β”€ inference                # inference (retrieval + generation) 
β”‚Β Β  β”‚Β Β  β”œβ”€β”€ __init__.py
β”‚Β Β  β”‚Β Β  β”œβ”€β”€ find_cluster.py
β”‚Β Β  β”‚Β Β  β”œβ”€β”€ main.py
β”‚Β Β  β”‚Β Β  β”œβ”€β”€ model.py
β”‚Β Β  β”‚Β Β  β”œβ”€β”€ search.py
β”‚Β Β  β”‚Β Β  └── types.py
β”‚Β Β  β”œβ”€β”€ ingestion                # Data ingestion & clustering         
β”‚Β Β  β”‚Β Β  β”œβ”€β”€ __init__.py
β”‚Β Β  β”‚Β Β  β”œβ”€β”€ centroid.py
β”‚Β Β  β”‚Β Β  β”œβ”€β”€ cluster.py
β”‚Β Β  β”‚Β Β  └── main.py
β”‚Β Β  └── utils                    # (chunking, embeddings, data loading)
β”‚Β Β      β”œβ”€β”€ __init__.py
β”‚Β Β      β”œβ”€β”€ chunker.py
β”‚Β Β      β”œβ”€β”€ data.py
β”‚Β Β      └── embedding.py
└── test
    β”œβ”€β”€ inference_test.py
    β”œβ”€β”€ ingestion_test.py
    └── test.py

Quick start

Installation

git clone https://github.com/rahulanand1103/mode.git
cd mode
pip install -r requirements.txt
pip install mode_rag
import os

## set ENV variables
os.environ["OPENAI_API_KEY"] = "your-api-key"

1. Ingestion Code

This is a sample using RecursiveCharacterTextSplitter and EmbeddingGenerator. You can use your own chunking/embedding logic. Main inputs to ModeIngestion are chunks and embeddings:

# ========================================
# πŸ“„ Sample Code: 
# ========================================
#
# 1. Loading pdf using PyPDFLoader
# 2. create chunking using `RecursiveCharacterTextSplitter`.
# 3. for embedding we are using langchain_huggingface.
# This is a sample using `RecursiveCharacterTextSplitter` and `EmbeddingGenerator`.
# You can use your **own chunking/embedding** logic.
# Main inputs to `ModeIngestion` are `chunks` and `embeddings`:


## requirements
# pip install langchain_huggingface==0.1.2
# pip install langchain_community==0.3.4
# pip install pypdf==5.1.0


import os
import json

os.environ["TOKENIZERS_PARALLELISM"] = "false"


from mode_rag import ModeIngestion, EmbeddingGenerator
import os
import json

## Pdf reader
from langchain_community.document_loaders import PyPDFLoader

loader = PyPDFLoader("https://arxiv.org/pdf/1706.03762")
docs = loader.load()

print("downloaded the files")

from langchain.text_splitter import RecursiveCharacterTextSplitter

print("Chunking the pdf:doc")
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
documents = text_splitter.split_documents(docs)
chunks = []
for doc in documents:
    chunks.append(doc.page_content)

print("doing embedding")
embed_gen = EmbeddingGenerator()
embeddings = embed_gen.generate_embeddings(chunks)
print("embedding done")
main_processor = ModeIngestion(
    chunks=chunks,
    embedding=embeddings,
    persist_directory="attention",
)
main_processor.process_data(parallel=False)

2. Inference Code

This is a sample using ModeInference and EmbeddingGenerator. You can use your own embedding method. Main inputs to ModeInference.invoke are query, query_embedding, and prompts:

# ========================================
# πŸ“„ Sample Code:
# ========================================
#
# 1. Load clustered data (`ModeInference`).
# 2. Generate query embedding (replaceable with your `embedding.py`).
# 3. Retrieve context and synthesize response with `ModelPrompt`.

import os
import json
import sys

os.environ["TOKENIZERS_PARALLELISM"] = "false"


from mode_rag import (
    EmbeddingGenerator,
    ModeInference,
    ModelPrompt,
)


main_processor = ModeInference(
    persist_directory="attention",
)

print("====start======")
# Create a PromptManager instance

query = "What are the key mathematical operations involved in computing self-attention?"

embed_gen = EmbeddingGenerator()
embedding = embed_gen.generate_embedding(query)

prompts = ModelPrompt(
    ref_sys_prompt="Use the following pieces of context to answer the user's question. \nIf you don't know the answer, just return you don't know.",
    ref_usr_prompt="context: ",
    syn_sys_prompt="You have been provided with a set of responses from various models to the latest user query. Your task is to synthesize these responses into a single, high-quality response. It is crucial to critically evaluate the information provided in these responses, recognizing that some of it may be biased or incorrect. Your response should not simply replicate the given answers but should offer a refined, accurate, and comprehensive reply to the instruction. Ensure your response is well-structured, coherent, and adheres to the highest standards of accuracy and reliability.\nResponses from models:",
    syn_usr_prompt="responses:",
)

response = main_processor.invoke(
    query,
    embedding,
    prompts,
    model_input={"temperature": 0.3, "model": "openai/gpt-4o-mini"},
    top_n_model=2,
)
print(response)

πŸ§ͺ Running Sample Scripts (Same as Above)

The same ingestion and inference logic is provided as ready-to-run test scripts inside the test/ folder.

You can quickly test MODE without writing any code!

Run Ingestion Test

cd test
python ingestion_test.py

Run Inference Test

cd test
python inference_test.py

Note:
These scripts (test/ingestion_test.py and test/inference_test.py) use the same examples shown above.


Benchmarking

Run experiments on different datasets using mode.py and traditional_rag.py.

Setup

cd benchmarking
pip install -r bench-requirements.txt

Run Benchmarks

HotpotQA

Mode:

python mode.py --dataset hotpotqa --chunks 100 --num_questions 100 --top_n_model 2

Traditional RAG:

python traditional_rag.py --dataset hotpotqa --chunks 100 --num_questions 100

SQuAD

Mode:

python mode.py --dataset squad --chunks 100 --num_questions 100 --top_n_model 2

Traditional RAG:

python traditional_rag.py --dataset squad --chunks 100 --num_questions 100

Notes:

  • --dataset must be either hotpotqa or squad.
  • --top_n_model is only used in mode.py.
  • Customize --chunks and --num_questions as needed.

πŸ“Š Benchmark Results

View Logs

MODE

DatasetNo. ChunkNo. QuestionTop n ModelGPT AccuracyGPT F1 ScoreBERT PrecisionBERT RecallBERT F1 Score
HotpotQA10010010.800.88890.80590.82760.8154
HotpotQA10010020.700.82350.74270.76120.7493
HotpotQA20010010.750.85710.80480.75820.7745
HotpotQA20010020.800.88890.77460.79100.7811
HotpotQA50010010.78430.87910.77770.75810.7613
HotpotQA50010020.80390.89130.72080.75070.7320
SQuAD10010010.780.87640.78810.79390.7852
SQuAD10010020.890.94180.78050.82410.7993
SQuAD20010010.720.83720.74490.73800.7336
SQuAD20010020.780.87640.74290.78280.7595
SQuAD50010010.710.83040.74950.74730.7408
SQuAD50010020.820.90110.76600.80470.7825

Traditional RAG

View Logs

DatasetNo. ChunksGPT AccuracyGPT F1 ScoreBERT PrecisionBERT F1 Score
HotpotQA1000.700.820.230.29
HotpotQA2000.700.820.370.40
HotpotQA5000.720.840.250.29
SQuAD1000.880.940.460.51
SQuAD2000.870.930.460.51
SQuAD5000.860.920.460.51

Contributing

We welcome contributions! Here’s how you can help:

  • Report Bugs: Submit issues on GitHub.
  • Suggest Features: Open an issue with your ideas.
  • Code Contributions: Fork, make changes, and submit a pull request.
  • Documentation: Update and enhance our docs.

License

This project is licensed under the MIT License.