Visual Document Retrieval
ColPali
Safetensors
English
vidore

Integrate with Sentence Transformers via MultiVectorEncoder

#15
by tomaarsen HF Staff - opened

Hello!

The MultiVectorEncoder class ships in the next Sentence Transformers release, planned for around the 18th, so for now the install below pulls from source. I would love to feature this model in that release's blog post and documentation, especially once it loads without the revision pin (that is, once this PR is merged).

Heads up, this PR was AI-generated and human-reviewed. Here's a summary of the changes as reported by my agent:

Pull Request overview

  • Integrate vidore/colpali with Sentence Transformers as a multi-vector (ColBERT-style late interaction) retriever via MultiVectorEncoder.

Details

This adds a Sentence Transformers loading path on top of the existing LoRA adapter, exposing the usual model.encode_query(...) / model.encode_document(...) / model.similarity(...) API with MaxSim scoring. The stock Transformer module loads the adapter directly onto the PaliGemma backbone through a small key_mapping that strips colpali-engine's model. wrapper prefix, so no custom modeling code or trust_remote_code is needed, only transformers>=5.15.0 (which ships huggingface/transformers#46766) and peft. The frozen custom_text_proj (2048 to 128) ships pre-merged as a roughly 1 MB 1_Dense module. The trained weights are untouched and the existing colpali-engine usage keeps working unchanged.

On the query format: this checkpoint predates a tagged colpali-engine release and the revision it records is not in the illuin-tech/colpali history, so I gave it the August 2024 Question: query format of its near-contemporaries. Current colpali-engine no longer sends that format: 0.3.4 changed the prefix from Question: to Query: (illuin-tech/colpali#125), 0.3.11 dropped the trailing newline (illuin-tech/colpali#280), and 0.3.13 dropped the prefix entirely (illuin-tech/colpali#339). This configuration reproduces the training-time format, so its embeddings differ slightly from current colpali-engine output, and the README flags this next to the colpali-engine snippet. On a ViDoRe v1 check, reproducing the training-time format improved nDCG@5 over the current colpali-engine format in 6 of 6 checkpoint x dataset cells measured.

pip install "sentence-transformers[image] @ git+https://github.com/huggingface/sentence-transformers.git"
from sentence_transformers import MultiVectorEncoder

model = MultiVectorEncoder("vidore/colpali", revision="refs/pr/15")

queries = [
    "What is the variable represented on the y-axis of the graph?",
    "Total outlay is maximum in which year?",
]
documents = [
    f"https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/doc{i}.jpg"
    for i in range(1, 5)
]

query_embeddings = model.encode_query(queries, convert_to_tensor=True)
document_embeddings = model.encode_document(documents, convert_to_tensor=True)
print(tuple(query_embeddings[0].shape), tuple(document_embeddings[0].shape))
# (23, 128) (1030, 128)

print(model.similarity(query_embeddings, document_embeddings))
# tensor([[17.3789, 17.1055, 15.4727, 15.4082],
#         [ 8.3750, 12.3047,  8.5898,  9.0957]])
  • Tom Aarsen
Ready to merge
This branch is ready to get merged automatically.

Sign up or log in to comment