Skip to content

maxbaluev/colpali-retrieve

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

1 Commit
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

colpali-retrieve

ColPali-style multimodal late-interaction retrieval. A text query is matched against document-page images by late interaction (MaxSim): the query is encoded to per-token vectors, each page to per-patch vectors, and relevance is the sum over query tokens of the best-matching image patch.

score(q, d) = Σ_i  max_j  sim(q_i, d_j)

This is the scoring rule from ColBERT (Khattab & Zaharia, 2020) generalised to images by ColPali (Faysse et al., 2024). This repo is a small, dependency-light (numpy-only) reference implementation of that scoring core, with a runnable synthetic demo and a documented path to a real ColPali/ColQwen encoder.

Why late interaction beats single-vector retrieval

A single-vector (bi-encoder) retriever pools an entire query — or an entire page — into one dense embedding and ranks by a single dot product. That pooling is lossy: a page that answers only part of a query, or whose answer lives in a small table cell or a figure axis, gets averaged away.

Late interaction keeps both sides as sets of vectors and defers the comparison to scoring time:

  • Each query token independently finds the document region it matches best (max_j). A query is rewarded only if every token finds a home somewhere.
  • For visual documents, ColPali encodes the page image directly into per-patch vectors — no OCR, no layout parsing. A query token like "revenue" can align to the exact chart patch that contains it.

That patch-level alignment is what lets ColPali outperform single-vector pipelines (and OCR-then-embed pipelines) on visually rich document retrieval benchmarks such as ViDoRe: tables, charts, figures, and scanned pages are matched by what they look like, region by region.

single-vector:   query ─► [one vec]      page ─► [one vec]      one dot product
late interaction: query ─► [tok,tok,tok]  page ─► [patch,…,patch]  MaxSim over all pairs

Install

pip install -e .

Core dependency is numpy only. For the real encoder path, install the optional extra (see below):

pip install -e ".[real-model]"   # pulls in colpali-engine

Usage — synthetic, runs offline (no model download)

python examples/retrieve.py

The example builds toy multi-vector embeddings and ranks three "pages" against the query "quarterly revenue chart". The scoring API:

import numpy as np
from colpali_retrieve import max_sim, rank

# query: per-token vectors (n_q, dim); document: per-patch vectors (n_d, dim)
query = np.array([[1.0, 0.0], [0.0, 1.0]])
page  = np.array([[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]])

max_sim(query, page)                 # -> 2.0  (both query tokens find an exact patch)
rank(query, [page_a, page_b, page_c])  # -> [(idx, score), ...] sorted desc

Usage — real ColPali / ColQwen encoder

The scoring code above is identical with real embeddings; only the encoder changes. Install the extra (pip install -e ".[real-model]") and feed the per-token / per-patch vectors into max_sim / rank:

import torch
from PIL import Image
from colpali_engine.models import ColQwen2, ColQwen2Processor
from colpali_retrieve import rank

model = ColQwen2.from_pretrained("vidore/colqwen2-v1.0", torch_dtype=torch.bfloat16)
processor = ColQwen2Processor.from_pretrained("vidore/colqwen2-v1.0")

# Encode page images -> per-patch vectors, and the query -> per-token vectors.
pages = [Image.open(p) for p in ("page1.png", "page2.png", "page3.png")]
with torch.no_grad():
    page_emb = model(**processor.process_images(pages).to(model.device))
    query_emb = model(**processor.process_queries(["quarterly revenue chart"]).to(model.device))

# Hand the multi-vector tensors (as numpy) to the late-interaction scorer.
docs = [p.float().cpu().numpy() for p in page_emb]
ranking = rank(query_emb[0].float().cpu().numpy(), docs)
print(ranking)  # [(page_index, maxsim_score), ...] sorted best-first

colpali-engine ships an optimised, batched MaxSim scorer for production use. The implementation here is the transparent reference: small, readable, and exact, so you can see precisely what late interaction computes.

API

Function Signature Description
cosine_similarity_matrix (query_tokens, doc_patches) -> np.ndarray Pairwise cosine matrix (n_q, n_d).
max_sim (query_tokens, doc_patches) -> float Late-interaction score Σ_i max_j sim(q_i, d_j).
rank (query_tokens, docs) -> list[(idx, score)] Score many documents, sorted descending (stable on ties).

Inputs are 2-D (n_vectors, dim) numpy arrays. Empty query / document / docs list yield 0.0 / 0.0 / []. Mismatched dimensions raise ValueError.

Tests

Fully offline — synthetic numpy vectors, no model or network:

pip install -e ".[test]"
pytest

References

  • O. Khattab and M. Zaharia. ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT. SIGIR 2020. arXiv:2004.12832
  • M. Faysse, H. Sibille, T. Wu, B. Omrani, G. Viaud, C. Hudelot, P. Colombo. ColPali: Efficient Document Retrieval with Vision Language Models. 2024. arXiv:2407.01449

License

MIT © 2026 Max Baluev. See LICENSE.

About

ColPali-style multimodal late-interaction retrieval (text → document-image patches via MaxSim) — numpy reference implementation

Topics

Resources

License

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages