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

Integrate with Sentence Transformers via MultiVectorEncoder

#8
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.3-hf with Sentence Transformers as a multi-vector (ColBERT-style late interaction) retriever via MultiVectorEncoder, and restore the Query: prefix 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 prefix: transformers' ColPaliProcessor defaults to query_prefix="Question: " (the pre-#125 value), but this checkpoint (November 2024) was trained with Query: , so out of the box every query gets the wrong prefix. Rather than add a processor_config.json (which would change AutoProcessor output for every existing user), sentence_bert_config.json sets the Query: prefix through processor_kwargs, so only Sentence Transformers loads are affected. Happy to switch to a processor_config.json if you would rather fix it for everyone.

It also fixes a stray max_length: 50 query cap and sets do_convert_rgb (a grayscale page otherwise failed to batch with colour pages). Verified in fp32: the Sentence Transformers embeddings and MaxSim scores match a plain ColPaliForRetrieval plus ColPaliProcessor baseline with max absolute difference 0.0.

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

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

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))
# (28, 128) (1030, 128)

print(model.similarity(query_embeddings, document_embeddings))
# tensor([[22.3359, 19.8555, 19.6582, 19.0928],
#         [ 5.8828, 13.3398,  6.1621,  6.8135]])
  • Tom Aarsen
tomaarsen changed pull request status to open
manu changed pull request status to merged

Sign up or log in to comment