visual-rag-toolkit 0.1.1__tar.gz

visual_rag_toolkit-0.1.1/.github/workflows/ci.yaml

@@ -0,0 +1,94 @@
+# Modern CI/CD with GitHub Actions
+# Replaces legacy travis.yml
+name: CI
+
+on:
+  push:
+    branches: [main, develop]
+  pull_request:
+    branches: [main]
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+      
+      - name: Install linting tools
+        run: pip install ruff black mypy
+      
+      - name: Run ruff
+        run: ruff check visual_rag/
+      
+      - name: Run black --check
+        run: black --check visual_rag/
+  
+  test:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, macos-latest]
+        python-version: ["3.9", "3.10", "3.11", "3.12"]
+    
+    steps:
+      - uses: actions/checkout@v4
+      
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e ".[dev]"
+      
+      - name: Run tests
+        run: pytest tests/ -v --cov=visual_rag --cov-report=xml
+      
+      - name: Upload coverage
+        uses: codecov/codecov-action@v4
+        if: matrix.os == 'ubuntu-latest' && matrix.python-version == '3.11'
+        with:
+          file: ./coverage.xml
+
+  # GPU tests (optional, needs self-hosted runner)
+  # test-gpu:
+  #   runs-on: [self-hosted, gpu]
+  #   steps:
+  #     - uses: actions/checkout@v4
+  #     - name: Run GPU tests
+  #       run: pytest tests/gpu/ -v
+
+  build:
+    runs-on: ubuntu-latest
+    needs: [lint, test]
+    
+    steps:
+      - uses: actions/checkout@v4
+      
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+      
+      - name: Install build tools
+        run: pip install build twine
+      
+      - name: Build package
+        run: python -m build
+      
+      - name: Check distribution
+        run: twine check dist/*
+      
+      - name: Upload artifacts
+        uses: actions/upload-artifact@v4
+        with:
+          name: dist
+          path: dist/

visual_rag_toolkit-0.1.1/.github/workflows/publish_pypi.yaml

@@ -0,0 +1,42 @@
+name: Publish to PyPI
+
+on:
+  push:
+    tags:
+      - "v*"
+  workflow_dispatch: {}
+
+permissions:
+  contents: read
+  id-token: write # required for PyPI trusted publishing (OIDC)
+
+jobs:
+  build-and-publish:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Install build tools
+        run: |
+          python -m pip install --upgrade pip
+          pip install build twine
+
+      - name: Build package
+        run: python -m build
+
+      - name: Check distribution
+        run: twine check dist/*
+
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          # Uses trusted publishing (no API token) once configured on PyPI:
+          # PyPI → Your project → Settings → Publishing → Add GitHub publisher
+          skip-existing: true
+

visual_rag_toolkit-0.1.1/.gitignore

@@ -0,0 +1,67 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Virtual environments
+venv/
+env/
+ENV/
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Project specific
+checkpoints/
+results/
+*.pkl
+*.pickle
+
+# Environment variables
+.env
+.env.local
+
+# Test coverage
+.coverage
+htmlcov/
+.pytest_cache/
+
+# Jupyter
+.ipynb_checkpoints/
+
+# MyPy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+
+
+
+
+
+

visual_rag_toolkit-0.1.1/LICENSE

@@ -0,0 +1,22 @@
+MIT License
+
+Copyright (c) 2026 Ara Yeroyan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+

visual_rag_toolkit-0.1.1/PKG-INFO

@@ -0,0 +1,305 @@
+Metadata-Version: 2.4
+Name: visual-rag-toolkit
+Version: 0.1.1
+Summary: End-to-end visual document retrieval with ColPali, featuring two-stage pooling for scalable search
+Project-URL: Homepage, https://github.com/Ara-Yeroyan/visual-rag-toolkit
+Project-URL: Documentation, https://github.com/Ara-Yeroyan/visual-rag-toolkit#readme
+Project-URL: Repository, https://github.com/Ara-Yeroyan/visual-rag-toolkit
+Project-URL: Issues, https://github.com/Ara-Yeroyan/visual-rag-toolkit/issues
+Author: Visual RAG Team
+License: MIT License
+        
+        Copyright (c) 2026 Ara Yeroyan
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+License-File: LICENSE
+Keywords: colbert,colpali,document-retrieval,late-interaction,multimodal-rag,pdf-processing,qdrant,visual-rag,visual-search
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Image Processing
+Requires-Python: >=3.9
+Requires-Dist: numpy>=1.21.0
+Requires-Dist: pillow>=9.0.0
+Requires-Dist: python-dotenv>=0.19.0
+Requires-Dist: pyyaml>=6.0
+Requires-Dist: torch>=2.0.0
+Requires-Dist: tqdm>=4.60.0
+Provides-Extra: all
+Requires-Dist: altair>=5.0.0; extra == 'all'
+Requires-Dist: cloudinary>=1.30.0; extra == 'all'
+Requires-Dist: colpali-engine>=0.3.0; extra == 'all'
+Requires-Dist: httpx>=0.24.0; extra == 'all'
+Requires-Dist: pandas>=2.0.0; extra == 'all'
+Requires-Dist: pdf2image>=1.16.0; extra == 'all'
+Requires-Dist: pypdf>=3.0.0; extra == 'all'
+Requires-Dist: qdrant-client>=1.7.0; extra == 'all'
+Requires-Dist: streamlit>=1.25.0; extra == 'all'
+Requires-Dist: transformers>=4.35.0; extra == 'all'
+Provides-Extra: cloudinary
+Requires-Dist: cloudinary>=1.30.0; extra == 'cloudinary'
+Provides-Extra: dev
+Requires-Dist: black>=23.0.0; extra == 'dev'
+Requires-Dist: mypy>=1.0.0; extra == 'dev'
+Requires-Dist: pytest-cov>=4.0.0; extra == 'dev'
+Requires-Dist: pytest>=7.0.0; extra == 'dev'
+Requires-Dist: ruff>=0.1.0; extra == 'dev'
+Provides-Extra: embedding
+Requires-Dist: colpali-engine>=0.3.0; extra == 'embedding'
+Requires-Dist: transformers>=4.35.0; extra == 'embedding'
+Provides-Extra: pdf
+Requires-Dist: pdf2image>=1.16.0; extra == 'pdf'
+Requires-Dist: pypdf>=3.0.0; extra == 'pdf'
+Provides-Extra: qdrant
+Requires-Dist: qdrant-client>=1.7.0; extra == 'qdrant'
+Provides-Extra: ui
+Requires-Dist: altair>=5.0.0; extra == 'ui'
+Requires-Dist: httpx>=0.24.0; extra == 'ui'
+Requires-Dist: pandas>=2.0.0; extra == 'ui'
+Requires-Dist: streamlit>=1.25.0; extra == 'ui'
+Description-Content-Type: text/markdown
+
+# Visual RAG Toolkit
+
+[![PyPI version](https://badge.fury.io/py/visual-rag-toolkit.svg)](https://badge.fury.io/py/visual-rag-toolkit)
+[![CI](https://github.com/Ara-Yeroyan/visual-rag-toolkit/actions/workflows/ci.yaml/badge.svg)](https://github.com/Ara-Yeroyan/visual-rag-toolkit/actions/workflows/ci.yaml)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![Python 3.9+](https://img.shields.io/badge/python-3.9+-blue.svg)](https://www.python.org/downloads/)
+
+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
+
+```bash
+# Core package (minimal dependencies)
+pip install visual-rag-toolkit
+
+# With specific features
+pip install visual-rag-toolkit[embedding]    # ColSmol/ColPali embedding support
+pip install visual-rag-toolkit[pdf]          # PDF processing
+pip install visual-rag-toolkit[qdrant]       # Vector database
+pip install visual-rag-toolkit[cloudinary]   # Image CDN
+pip install visual-rag-toolkit[ui]           # Streamlit demo dependencies
+
+# 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)
+
+```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",  # or: tokens_vs_tiles / pooled_query_vs_tiles / pooled_query_vs_global
+)
+
+for r in results[:3]:
+    print(r["id"], r["score_final"])
+```
+
+### 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 SIGIR_QDRANT_URL="https://your-cluster.qdrant.io"
+export SIGIR_QDRANT_KEY="your-api-key"
+
+# Backwards-compatible fallbacks (also supported)
+export QDRANT_URL="https://your-cluster.qdrant.io"
+export QDRANT_API_KEY="your-api-key"
+
+export VISUALRAG_MODEL="vidore/colSmol-500M"
+
+# 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 \
+  --prefetch-k 256 \
+  --top-k 100 \
+  --evaluation-scope union
+```
+
+More commands (including multi-stage variants and cropping configs) live in:
+- `benchmarks/vidore_tatdqa_test/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 Two-Stage 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
+