Visual Document Retrieval
ColPali
Safetensors
sentence-transformers
English
paligemma
vidore-experimental
multi-vector

Integrate with Sentence Transformers via MultiVectorEncoder

#6
by tomaarsen HF Staff - opened
Vidore org
•
edited Aug 14

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-v1.2-merged 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 merged checkpoint. The PaliGemma backbone and the custom_text_proj head already live in the weights, so a stock Transformer loads it directly, with no trust_remote_code, no peft, and no colpali-engine at inference. A small key_mapping strips colpali-engine's model. wrapper prefix, and the projection head ships as the roughly 1 MB 1_Dense module. The trained weights are untouched and the existing colpali-engine usage is unchanged.

On the query format: this checkpoint was trained with the Question: query prefix and a trailing newline, verified token-identical against the matching-era colpali-engine build. 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-v1.2-merged", revision="refs/pr/6")

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([[11.1816, 10.7031,  8.3398,  5.6191],
#         [ 5.5391,  9.4785,  4.7588,  6.1982]])
  • Tom Aarsen
tomaarsen changed pull request status to open
QuentinJG changed pull request status to merged

Sign up or log in to comment