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-hard-v1.1 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 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-hard-v1.1", 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([[13.4004, 13.0762,  9.6904,  9.7334],
#         [ 4.9082,  9.5684,  5.9062,  6.7139]])
  • Tom Aarsen
tomaarsen changed pull request status to open
manu changed pull request status to merged

Sign up or log in to comment