Integrate with Sentence Transformers via MultiVectorEncoder

#5
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 ModernVBERT/colmodernvbert 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 loads the adapter onto a bare ModernVBert backbone via a small key_mapping that strips the model. wrapper prefix, with the target_modules regex relaxed from model.text_model to text_model so PEFT finds the modules on the bare backbone (still matching the colpali-engine wrapper). No custom modeling code or trust_remote_code is needed, only transformers>=5.15.0 (huggingface/transformers#46766) and peft. The 768-to-128 projection ships pre-merged as the 1_Dense module. The trained weights are untouched, the existing colpali-engine usage is unchanged, and the MaxSim scores match a colpali-engine fp32 baseline bit-for-bit (0.0).

One processor note: AutoProcessor natively resolves ColModernVBertProcessor, but that class hardcodes chat_template = None, so the apply_chat_template call Sentence Transformers makes would raise. ColModernVBertProcessor is a thin subclass of Idefics3Processor, which produces byte-identical token ids while honouring the chat template, so processor_class is flipped to Idefics3Processor. This only affects AutoProcessor.from_pretrained, so direct ColModernVBertProcessor use is unaffected.

On the chat template: the checkpoint's conversational chat_template.jinja is kept as the default, and the Sentence Transformers query/document prompt format is added as a named template at additional_chat_templates/sentence_transformers.jinja, selected through the ST config. apply_chat_template is unchanged, and colpali-engine formats its inputs directly so it uses neither.

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

model = MultiVectorEncoder("ModernVBERT/colmodernvbert", revision="refs/pr/5")

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))
# (26, 128) (1149, 128)

print(model.similarity(query_embeddings, document_embeddings))
# tensor([[16.7778, 10.3712, 11.8420,  9.0534],
#         [ 7.3722, 12.0618,  8.1477,  7.9563]])
  • Tom Aarsen
tomaarsen changed pull request status to open
QuentinJG changed pull request status to merged

Sign up or log in to comment