visual-rag-toolkit 0.1.2

End-to-end visual document retrieval with ColPali, featuring two-stage pooling for scalable search

Visual RAG Toolkit

Note:

• The PyPI badge shows “not found” until the first release is published.
• The CI badge requires the GitHub repo to be public (GitHub does not serve Actions badges for private repos).

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

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 support behind a single VisualEmbedder interface.
• 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

# 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 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

🚀 Quick Start

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

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",  # or: tokens_vs_tiles / pooled_query_vs_tiles / 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.

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["SIGIR_QDRANT_URL"]  # or QDRANT_URL
QDRANT_KEY = os.getenv("SIGIR_QDRANT_KEY", "")  # or 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",
)
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:
    pipeline.process_pdf(
        pdf_path,
        skip_existing=True,
        upload_to_cloudinary=False,
        upload_to_qdrant=True,
    )

CLI equivalent:

export SIGIR_QDRANT_URL="https://YOUR_QDRANT"
export SIGIR_QDRANT_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)

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:

# 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):

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)

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:

# 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 \
  --prefetch-k 256 \
  --top-k 100 \
  --evaluation-scope union

More commands (including multi-stage variants and cropping configs) live in:

• examples/COMMANDS.md

🔧 Development

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:

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

📝 License

MIT License - see LICENSE for details.

🙏 Acknowledgments

• Qdrant - Vector database with multi-vector support
• ColPali - Visual document retrieval models
• ViDoRe - Benchmark dataset