Integrate with Sentence Transformers via MultiVectorEncoder
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/colpaliwith Sentence Transformers as a multi-vector (ColBERT-style late interaction) retriever viaMultiVectorEncoder.
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
Hey @tomaarsen thanks a lot for all those PRs, I'm verifying them one by one to check everything works right but as far as I checked it looks perfect :) (it actually underlined a few inconsistencies in the original colpali-engine repo...)
Looking forward to train our next Multimodal Multivector models using ST 😁
@QuentinJG I'm glad it's been working well for you so far! I see you've merged the majority of them, I believe only these are still open:
- https://huggingface.co/vidore/colpali-hard-v1.1/discussions/2
- https://huggingface.co/vidore/colpali-v1.2-hf/discussions/2
- https://huggingface.co/vidore/colpali-v1.3-hf/discussions/8
- https://huggingface.co/vidore/colqwen-omni-v0.1/discussions/1
The last one is especially interesting, as it's the only audio and video-capable model that I'm planning on supporting on day-0. Do let me know if you encounter issues with any of them!
Edit: I see @manu has just merged those 🤗
- Tom Aarsen