README.md

---
title: Visual Rag Toolkit
emoji: 👀
colorFrom: green
colorTo: red
sdk: docker
pinned: false
license: mit
short_description: Efficient Multi Vector Visual RAG Toolkit
---


# Visual RAG Toolkit

[![PyPI](https://img.shields.io/pypi/v/visual-rag-toolkit)](https://pypi.org/project/visual-rag-toolkit/)
[![Python](https://img.shields.io/pypi/pyversions/visual-rag-toolkit)](https://pypi.org/project/visual-rag-toolkit/)
[![License](https://img.shields.io/pypi/l/visual-rag-toolkit)](LICENSE)
[![Demo](https://img.shields.io/badge/Demo-Hugging%20Face-yellow)](https://huggingface.co/spaces/Yeroyan/visual-rag-toolkit)

End-to-end visual document retrieval toolkit featuring **fast multi-stage retrieval** (prefetch with pooled vectors + exact MaxSim reranking).

**[Try the Live Demo](https://huggingface.co/spaces/Yeroyan/visual-rag-toolkit)** - Upload PDFs, index to Qdrant, and query with visual retrieval.

**[Watch the Tutorial](https://drive.google.com/file/d/1SmpVJicsvyZ-awlwYtSkLw6BCSO7IJTw/view)** - Video walkthrough of the toolkit in action.

This repo contains:
- a **Python package** (`visual_rag`)
- a **Streamlit demo app** (`demo/`)
- **benchmark & evaluation scripts** for ViDoRe v2 (`benchmarks/`)

## 🎯 Key Features

- **Modular**: PDF → images, embedding, Qdrant indexing, retrieval can be used independently.
- **Multi-stage retrieval**: two-stage and three-stage retrieval modes built for Qdrant named vectors.
- **Model-aware embedding**: ColSmol, ColPali, and ColQwen2/2.5 support behind a single `VisualEmbedder` interface.
- **Configurable pooling**: adaptive mean-pooling cap for ColQwen2.5 (`--max-mean-pool-vectors`), and experimental pooling stored as Qdrant named vectors (`experimental_pooling` (ColQwen Gaussian alias), `experimental_pooling_gaussian`, `experimental_pooling_triangular`, `experimental_pooling_{k}` (ColPali), `experimental_pooling_2d` (ColSmol)).
- **Single-stage ablations**: direct search modes over experimental pooled vectors (tokens-vs-doc and pooled-query-vs-doc) for fast storage-reduction experiments.
- **Token hygiene**: query special-token filtering by default for more stable MaxSim behavior.
- **Practical pipelines**: robust indexing, retries, optional Cloudinary image URLs, evaluation reporting.

## 📦 Installation

```bash
# Core package (minimal dependencies)
pip install visual-rag-toolkit

# With specific features
pip install visual-rag-toolkit[ui]           # Streamlit demo dependencies
pip install visual-rag-toolkit[qdrant]       # Vector database
pip install visual-rag-toolkit[embedding]    # ColSmol/ColPali/ColQwen2(.5) embedding support
pip install visual-rag-toolkit[cloudinary]   # Image CDN

# All dependencies
pip install visual-rag-toolkit[all]
```

### System dependencies (PDF)

`pdf2image` requires Poppler.

- macOS: `brew install poppler`
- Ubuntu/Debian: `sudo apt-get update && sudo apt-get install -y poppler-utils`

**ColQwen2.5 note:** `vidore/colqwen2.5-v0.2` requires `transformers>=4.45.0` and `colpali-engine>=0.3.7` (installing ColPali/ColQwen from source may be required for the latest processors).

## 🚀 Quick Start

### Minimal: embed a query and run two-stage search (server-side)

```python
from qdrant_client import QdrantClient
from visual_rag import VisualEmbedder, TwoStageRetriever

client = QdrantClient(url="https://YOUR_QDRANT", api_key="YOUR_KEY")
collection_name = "your_collection"

# Embed query tokens
embedder = VisualEmbedder(model_name="vidore/colpali-v1.3")
q = embedder.embed_query("What is the budget allocation?")

# Fast path: all stages computed in Qdrant (prefetch + exact rerank)
retriever = TwoStageRetriever(client, collection_name)
results = retriever.search_server_side(
    query_embedding=q,
    top_k=10,
    prefetch_k=256,
    stage1_mode="tokens_vs_experimental_pooling",  # or: tokens_vs_standard_pooling / pooled_query_vs_standard_pooling / pooled_query_vs_global
)

for r in results[:3]:
    print(r["id"], r["score_final"])
```

### End-to-end: ingest PDFs (with cropping) → index in Qdrant

This is the "SDK-style" pipeline: PDF → images → optional crop → embed → store vectors + payload in Qdrant.

```python
import os
from pathlib import Path

import numpy as np
import torch

from visual_rag import VisualEmbedder
from visual_rag.indexing import ProcessingPipeline, QdrantIndexer

QDRANT_URL = os.environ["QDRANT_URL"]
QDRANT_KEY = os.getenv("QDRANT_API_KEY", "")

collection = "my_visual_docs"

embedder = VisualEmbedder(
    model_name="vidore/colSmol-500M",
    torch_dtype=torch.float16,
    output_dtype=np.float16,
    batch_size=8,
)

indexer = QdrantIndexer(
    url=QDRANT_URL,
    api_key=QDRANT_KEY,
    collection_name=collection,
    prefer_grpc=True,
    vector_datatype="float16",
)

# Creates collection + required payload indexes (e.g., "filename" for skip_existing)
indexer.create_collection(force_recreate=False)

pipeline = ProcessingPipeline(
    embedder=embedder,
    indexer=indexer,
    embedding_strategy="all",  # store full tokens + pooled vectors in one pass
    crop_empty=True,
    crop_empty_percentage_to_remove=0.99,  # kept for traceability
    crop_empty_remove_page_number=True,
    crop_empty_preserve_border_px=1,
    crop_empty_uniform_rowcol_std_threshold=3.0,
)

pdfs = [Path("docs/a.pdf"), Path("docs/b.pdf")]
for pdf_path in pdfs:
    result = pipeline.process_pdf(
        pdf_path,
        skip_existing=True,  # Skip pages already in Qdrant (uses filename index)
        upload_to_cloudinary=False,
        upload_to_qdrant=True,
    )
    # Logs automatically shown:
    # [10:23:45] 📚 Processing PDF: a.pdf
    # [10:23:45] 🖼️ Converting PDF to images...
    # [10:23:46]    ✅ Converted 12 pages
    # [10:23:46] 📦 Processing pages 1-8/12
    # [10:23:46] 🤖 Generating embeddings for 8 pages...
    # [10:23:48] 📤 Uploading batch of 8 pages...
    # [10:23:48]    ✅ Uploaded 8 points to Qdrant
    # [10:23:48] 📦 Processing pages 9-12/12
    # [10:23:48] 🤖 Generating embeddings for 4 pages...
    # [10:23:50] 📤 Uploading batch of 4 pages...
    # [10:23:50]    ✅ Uploaded 4 points to Qdrant
    # [10:23:50] ✅ Completed a.pdf: 12 uploaded, 0 skipped, 0 failed
```

CLI equivalent:

```bash
export QDRANT_URL="https://YOUR_QDRANT"
export QDRANT_API_KEY="YOUR_KEY"

visual-rag process \
  --reports-dir ./docs \
  --collection my_visual_docs \
  --model vidore/colSmol-500M \
  --strategy all \
  --batch-size 8 \
  --qdrant-vector-dtype float16 \
  --prefer-grpc \
  --crop-empty \
  --crop-empty-remove-page-number
```

### Process a PDF into images (no embedding, no vector DB)

```python
from pathlib import Path
from visual_rag import PDFProcessor

processor = PDFProcessor(dpi=140)
images, texts = processor.process_pdf(Path("report.pdf"))
print(len(images), "pages")
```

## 🔬 Multi-stage Retrieval (Two-stage / Three-stage)

Traditional ColBERT-style MaxSim scoring compares all query tokens vs all document tokens, which becomes expensive at scale.

**Our approach:**

```
Stage 1: Fast prefetch with tile-level pooled vectors
         ├── Pool each tile (64 patches) → num_tiles vectors
         ├── Use HNSW index for O(log N) retrieval  
         └── Retrieve top-K candidates (e.g., 200)

Stage 2: Exact MaxSim reranking on candidates
         ├── Load full multi-vector embeddings
         ├── Compute exact ColBERT MaxSim scores
         └── Return top-k results (e.g., 10)
```

Three-stage extends this with an additional "cheap prefetch" stage before stage 2.

## 📁 Package Structure

```
visual-rag-toolkit/
├── visual_rag/              # Import as: from visual_rag import ...
│   ├── embedding/           # VisualEmbedder, pooling functions
│   ├── indexing/            # PDFProcessor, QdrantIndexer, CloudinaryUploader
│   ├── retrieval/           # TwoStageRetriever
│   ├── visualization/       # Saliency maps
│   ├── cli/                 # Command-line: visual-rag process/search
│   └── config.py            # load_config, get, get_section
│
├── benchmarks/              # ViDoRe evaluation scripts
└── examples/                # Usage examples
```

## ⚙️ Configuration

Configure via environment variables or YAML:

```bash

# Qdrant credentials (preferred names used by the demo + scripts)
export QDRANT_URL="https://your-cluster.qdrant.io"
export QDRANT_API_KEY="your-api-key"

# Special token handling (default: filter them out)
export VISUALRAG_INCLUDE_SPECIAL_TOKENS=true  # Include special tokens
```

Or use a config file (`visual_rag.yaml`):

```yaml
model:
  name: "vidore/colSmol-500M"
  batch_size: 4
  
qdrant:
  url: "https://your-cluster.qdrant.io"
  collection: "my_documents"
  
search:
  strategy: "two_stage"  # or "multi_vector", "pooled"
  prefetch_k: 200
  top_k: 10
```

## 🖥️ Demo (Streamlit)

```bash
pip install "visual-rag-toolkit[ui,qdrant,embedding,pdf]"

# Option A: from Python
python -c "import visual_rag; visual_rag.demo()"

# Option B: CLI launcher
visual-rag-demo
```

## 📊 Benchmark Evaluation

Run ViDoRe benchmark evaluation:

```bash
# Example: evaluate a collection against ViDoRe BEIR datasets in Qdrant
python -m benchmarks.vidore_beir_qdrant.run_qdrant_beir \
  --datasets vidore/esg_reports_v2 vidore/biomedical_lectures_v2 \
  --collection YOUR_COLLECTION \
  --mode two_stage \
  --stage1-mode tokens_vs_experimental_pooling \
  --prefetch-k 256 \
  --top-k 100 \
  --evaluation-scope union
```

More commands (including multi-stage variants and cropping configs) live in:
- `examples/COMMANDS.md`

## 🔧 Development

```bash
git clone https://github.com/Ara-Yeroyan/visual-rag-toolkit
cd visual-rag-toolkit
pip install -e ".[dev]"
pytest tests/ -v
```

## 📄 Citation

If you use this toolkit in your research, please cite:

```bibtex
@software{visual_rag_toolkit,
  title = {Visual RAG Toolkit: Scalable Visual Document Retrieval with 1D Convolutional Pooling},
  author = {Ara Yeroyan},
  year = {2026},
  url = {https://github.com/Ara-Yeroyan/visual-rag-toolkit}
}
```

## 📝 License

MIT License - see [LICENSE](LICENSE) for details.

## 🙏 Acknowledgments

- [Qdrant](https://qdrant.tech/) - Vector database with multi-vector support
- [ColPali](https://github.com/illuin-tech/colpali) - Visual document retrieval models
- [ViDoRe](https://huggingface.co/spaces/vidore/vidore-leaderboard) - Benchmark dataset