ColPali is a model based on a novel model architecture and training strategy based on Vision Language Models (VLMs) to efficiently index documents from their visual features.
It is a
PaliGemma-3B extension that generates
ColBERT- style multi-vector representations of text and images.
It was introduced in the paper
ColPali: Efficient Document Retrieval with Vision Language Models and first released in
this repository
This model is built iteratively starting from an off-the-shelf
SigLIP model.
We finetuned it to create
BiSigLIP and fed the patch-embeddings output by SigLIP to an LLM,
PaliGemma-3B to create
BiPali.
One benefit of inputting image patch embeddings through a language model is that they are natively mapped to a latent space similar to textual input (query).
This enables leveraging the
ColBERT strategy to compute interactions between text tokens and image patches, which enables a step-change improvement in performance compared to BiPali.
Our training dataset of 127,460 query-page pairs is comprised of train sets of openly available academic datasets (63%) and a synthetic dataset made up of pages from web-crawled PDF documents and augmented with VLM-generated (Claude-3 Sonnet) pseudo-questions (37%).
Our training set is fully English by design, enabling us to study zero-shot generalization to non-English languages. We explicitly verify no multi-page PDF document is used both
ViDoRe and in the train set to prevent evaluation contamination.
A validation set is created with 2% of the samples to tune hyperparameters.
All models are trained for 1 epoch on the train set. Unless specified otherwise, we train models in
bfloat16 format, use low-rank adapters (
LoRA)
with
alpha=32 and
r=32 on the transformer layers from the language model,
as well as the final randomly initialized projection layer, and use a
paged_adamw_8bit optimizer.
We train on an 8 GPU setup with data parallelism, a learning rate of 5e-5 with linear decay with 2.5% warmup steps, and a batch size of 32.
ColPali can be used as a multi-vector (ColBERT-style late interaction) retriever directly with Sentence Transformers via the MultiVectorEncoder.
1from sentence_transformers import MultiVectorEncoder
2
3model = MultiVectorEncoder("vidore/colpali")
4
5queries = [
6 "What is the variable represented on the y-axis of the graph?",
7 "Total outlay is maximum in which year?",
8]
9images = [
10 "https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/doc1.jpg",
11 "https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/doc2.jpg",
12 "https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/doc3.jpg",
13 "https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/doc4.jpg",
14]
15
16query_embeddings = model.encode_query(queries)
17document_embeddings = model.encode_document(images)
18print(f"Query 0 shape: {tuple(query_embeddings[0].shape)}")
19print(f"Document 0 shape: {tuple(document_embeddings[0].shape)}")
20# Query 0 shape: (23, 128)
21# Document 0 shape: (1030, 128)
22
23# MaxSim late-interaction scoring (rows = queries, columns = images)
24scores = model.similarity(query_embeddings, document_embeddings)
25print(scores)
26# tensor([[17.3789, 17.1055, 15.4727, 15.4082],
27# [ 8.3750, 12.3047, 8.5898, 9.0957]])
1# This model checkpoint is compatible with version 0.1.1, but not more recent versions of the inference lib
2pip install colpali_engine==0.1.1
1import torch
2import typer
3from torch.utils.data import DataLoader
4from tqdm import tqdm
5from transformers import AutoProcessor
6from PIL import Image
7
8from colpali_engine.models.paligemma_colbert_architecture import ColPali
9from colpali_engine.trainer.retrieval_evaluator import CustomEvaluator
10from colpali_engine.utils.colpali_processing_utils import process_images, process_queries
11from colpali_engine.utils.image_from_page_utils import load_from_dataset
12
13
14def main() -> None:
15 """Example script to run inference with ColPali"""
16
17 # Load model
18 model_name = "vidore/colpali"
19 model = ColPali.from_pretrained("vidore/colpaligemma-3b-mix-448-base", torch_dtype=torch.bfloat16, device_map="cuda").eval()
20 model.load_adapter(model_name)
21 processor = AutoProcessor.from_pretrained(model_name)
22
23 # select images -> load_from_pdf(<pdf_path>), load_from_image_urls(["<url_1>"]), load_from_dataset(<path>)
24 images = load_from_dataset("vidore/docvqa_test_subsampled")
25 queries = ["From which university does James V. Fiorca come ?", "Who is the japanese prime minister?"]
26
27 # run inference - docs
28 dataloader = DataLoader(
29 images,
30 batch_size=4,
31 shuffle=False,
32 collate_fn=lambda x: process_images(processor, x),
33 )
34 ds = []
35 for batch_doc in tqdm(dataloader):
36 with torch.no_grad():
37 batch_doc = {k: v.to(model.device) for k, v in batch_doc.items()}
38 embeddings_doc = model(**batch_doc)
39 ds.extend(list(torch.unbind(embeddings_doc.to("cpu"))))
40
41 # run inference - queries
42 dataloader = DataLoader(
43 queries,
44 batch_size=4,
45 shuffle=False,
46 collate_fn=lambda x: process_queries(processor, x, Image.new("RGB", (448, 448), (255, 255, 255))),
47 )
48
49 qs = []
50 for batch_query in dataloader:
51 with torch.no_grad():
52 batch_query = {k: v.to(model.device) for k, v in batch_query.items()}
53 embeddings_query = model(**batch_query)
54 qs.extend(list(torch.unbind(embeddings_query.to("cpu"))))
55
56 # run evaluation
57 retriever_evaluator = CustomEvaluator(is_multi_vector=True)
58 scores = retriever_evaluator.evaluate(qs, ds)
59 print(scores.argmax(axis=1))
60
61
62if __name__ == "__main__":
63 typer.run(main)
64
ColPali's vision language backbone model (PaliGemma) is under
gemma license as specified in its
model card. The adapters attached to the model are under MIT license.
If you use any datasets or models from this organization in your research, please cite the original dataset as follows:
1@misc{faysse2024colpaliefficientdocumentretrieval,
2 title={ColPali: Efficient Document Retrieval with Vision Language Models},
3 author={Manuel Faysse and Hugues Sibille and Tony Wu and Bilel Omrani and Gautier Viaud and Céline Hudelot and Pierre Colombo},
4 year={2024},
5 eprint={2407.01449},
6 archivePrefix={arXiv},
7 primaryClass={cs.IR},
8 url={https://arxiv.org/abs/2407.01449},
9}