Visual Document Retrieval
Transformers
Safetensors
ColPali
sentence-transformers
English
pretraining
multi-vector

Integrate with Sentence Transformers via MultiVectorEncoder

#2
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-v1.2-hf with Sentence Transformers as a multi-vector (ColBERT-style late interaction) retriever via MultiVectorEncoder, and restore the five-token <unused0> query augmentation buffer this checkpoint was trained with.

Details

This adds a Sentence Transformers loading path so the checkpoint works with the familiar model.encode_query(...) / model.encode_document(...) / model.similarity(...) API. The pipeline is just Transformer(retrieval) -> MultiVectorMask, because ColPaliForRetrieval.forward already applies the projection, L2 normalization and padding zeroing, so there is no 1_Dense or 2_Normalize to ship. Sentence Transformers recognises architectures: ["ColPaliForRetrieval"] and builds the pipeline automatically, so no custom modeling code or trust_remote_code is needed.

The substantive fix is the query augmentation buffer: this checkpoint was trained with Question: {query} plus five <unused0> tokens (23 tokens for the README query), but transformers' ColPaliProcessor hardcodes ten <pad> tokens, so out of the box it renders the later-checkpoint format instead. sentence_bert_config.json carries a processing_kwargs text suffix that restores the trained five-token render exactly. The query prefix needs no change, since it already defaults to the Question: this checkpoint used. This is scoped to Sentence Transformers, so nothing changes for existing users.

It also corrects processor_class (it named PaliGemmaProcessor, which has no process_queries or process_images, so it now names ColPaliProcessor) and applies two smaller config fixes (a stray max_length: 50 query cap and do_convert_rgb). Verified in fp32: the query token ids match the training-era render, and MaxSim scores agree with the colpali-engine-format vidore/colpali-v1.2-merged to 4.8e-6.

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

model = MultiVectorEncoder("vidore/colpali-v1.2-hf", revision="refs/pr/2")

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
manu changed pull request status to merged

Sign up or log in to comment