Embedding Model RAG Benchmark

This repository provides a framework for benchmarking different embedding models in a Retrieval-Augmented Generation (RAG) system using ChromaDB and Maximal Marginal Relevance (MMR).

Models Being Benchmarked

The system benchmarks the following open-source embedding models:

1. all-MiniLM-L6-v2 - A lightweight sentence transformer model
2. all-mpnet-base-v2 - A more powerful sentence transformer model
3. E5-large - A state-of-the-art embedding model from Microsoft
4. Instructor-Large - An instruction-tuned embedding model
5. Contriever - Facebook's retrieval-focused embedding model

Project Structure

embedding-model-rag-poc/
├── src/
│   ├── embedding_models/  # Embedding model implementations
│   ├── vector_db/         # ChromaDB integration with MMR
│   ├── data/              # Dataset handling
│   ├── evaluation/        # Benchmarking tools
│   └── rag/               # RAG system implementation
├── run_benchmark.py       # Main script to run benchmarks
├── create_rag_app.py      # Create a RAG application with the best model
├── interactive_rag_demo.py # Interactive demo for the RAG system
└── requirements.txt       # Project dependencies

Getting Started

Installation

1. Clone the repository:

git clone https://github.com/yourusername/embedding-model-rag-poc.git
cd embedding-model-rag-poc

2. Install dependencies:

pip install -r requirements.txt

3. Set up environment variables (for LLM API keys):

# Create a .env file with your API keys
echo "OPENAI_API_KEY=your_openai_key" > .env

Running the Benchmark

To run the benchmark with default settings:

python run_benchmark.py --create-sample

This will:

1. Create a sample dataset
2. Initialize each embedding model
3. Create vector stores in ChromaDB
4. Run retrieval tests with different MMR settings
5. Generate benchmark results

Creating a RAG Application

After benchmarking, create a RAG application using the best model:

python create_rag_app.py --from-benchmark ./benchmark_results --documents-dir ./your_documents

Interactive Demo

Test your RAG system interactively:

python interactive_rag_demo.py

Customization Options

Benchmark Options

• --dataset-name: Use a HuggingFace dataset
• --dataset-path: Use a local dataset
• --models: Specify which models to benchmark
• --k-values: List of k values to test for retrieval
• --lambda-values: List of lambda values to test for MMR

RAG Application Options

• --embedding-model: Manually specify an embedding model
• --llm-provider: Choose LLM provider (openai, anthropic, etc.)
• --llm-model: Specify LLM model name
• --mmr-lambda: Set MMR diversity parameter

Evaluation Metrics

The benchmark evaluates models on:

• NDCG@k: Normalized Discounted Cumulative Gain
• Precision@k: Precision at k retrieved documents
• Recall@k: Recall at k retrieved documents
• Diversity: Diversity of retrieved documents (for MMR)
• Retrieval Time: Time taken to retrieve documents

License

MIT License