Visual Document Retrieval
PEFT
Safetensors
ColPali
vidore
multimodal_embedding
multilingual_embedding
Text-to-Visual Document (T→VD) retrieval

Integrate with Sentence Transformers via MultiVectorEncoder

#3
by tomaarsen HF Staff - opened

Hello!

I recently released the new MultiVectorEncoder class in Sentence Transformers v6.0, and I think the ColNomic models would be great fits for it. You can read more about the release here: https://huggingface.co/blog/multi-vector-encoder. I would also love to feature this model in the accompanying documentation. I am opening the same integration for colnomic-embed-multimodal-3b.

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 nomic-ai/colnomic-embed-multimodal-7b with Sentence Transformers as a multi-vector (ColBERT-style late interaction) retriever via MultiVectorEncoder.

Details

This mirrors the integration recently merged into vidore/colqwen2.5-v0.2, adapted for this adapter repository. The integration is config-only: no new modeling code, and the adapter weights are untouched. Sentence Transformers loads the base model plus this adapter directly. The module pipeline Transformer -> Dense -> Normalize -> MultiVectorMask mirrors the reference forward step for step (last hidden state -> custom_text_proj -> L2 normalize -> mask by attention_mask), where 1_Dense/ carries the custom_text_proj projection (3584 -> 128) with the LoRA weights merged in, exactly what the peft-wrapped reference computes. A named chat template in additional_chat_templates/sentence_transformers.jinja reproduces the processor's prompt formats: plain text renders the query format with the ten augmentation tokens, and image inputs render the visual document format. The query branch renders the training-time Query: prefix, see below.

sentence_bert_config.json maps the adapter's ^model. keys onto the language_model. submodule of the transformers 5.x layout and declares a structured requirements floor of transformers>=5.15, the first version that applies such mappings on adapter loads. The legacy chat_template.json is also converted to chat_template.jinja: on transformers 5.x the legacy file conflicts with additional_chat_templates/, so the conversion doubles as a fix for plain transformers usage. The README gains a Sentence Transformers usage section, the query prefix warning described below, and the sentence-transformers tag.

On the query prefix, one thing is worth flagging on its own. This repository was uploaded in March 2025 and its card installs colpali-engine from source. ColQwen2_5_Processor prefixed every query with "Query: " from the first ColQwen2.5 commit onward (0.3.11 moved the value to the shared base class unchanged, illuin-tech/colpali#280), and the prefix was dropped in 0.3.13 (illuin-tech/colpali#339), months after this model was trained, so a current install no longer reproduces the training-time query format. The chat template here ships the training-time format, matching what was done for vidore/colqwen2.5-v0.2 and its siblings, and the README gains a warning in the ColPali Engine section pointing out the difference and how to restore it on that path (processor.query_prefix = "Query: ").

One note on precision: the nomic-ai/colqwen2.5-7B-base config declares no dtype, so a plain load runs in fp32 (roughly 28 GB). The snippet below passes model_kwargs={"dtype": "bfloat16"} to match the model card.

Verified against the repository's own colpali-engine pipeline (base plus adapter), with processor.query_prefix = "Query: " restored on the reference so the comparison is against the training-time behaviour: the per-token embeddings match with cosine similarity 0.99996 and the MaxSim scores differ by at most 0.028. For reference, on a 400-query NanoViDoRe v3 subset this integration scores 0.5942 nDCG@10, above the model's 0.5688 ViDoRe v3 leaderboard entry, consistent with the restored training-time query prefix helping.

pip install "sentence-transformers[image]>=6.0.0"
from sentence_transformers import MultiVectorEncoder

model = MultiVectorEncoder(
    "nomic-ai/colnomic-embed-multimodal-7b",
    model_kwargs={"dtype": "bfloat16"},
)

queries = [
    "What is the variable represented on the y-axis of the graph?",
    "Total outlay is maximum in which year?",
]
documents = [
    "https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/doc1.jpg",
    "https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/doc2.jpg",
    "https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/doc3.jpg",
    "https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/doc4.jpg",
]

query_embeddings = model.encode_query(queries)
document_embeddings = model.encode_document(documents)
print(query_embeddings[0].shape, document_embeddings[0].shape)
# torch.Size([25, 128]) torch.Size([755, 128])

scores = model.similarity(query_embeddings, document_embeddings)
print(scores)
# tensor([[13.4023, 11.7891, 11.3535, 10.5449],
#         [ 9.1855, 14.4668,  9.3027,  8.6133]], device='cuda:0')
print("Best document per query:", scores.argmax(dim=1))
# Best document per query: tensor([0, 1], device='cuda:0')

To try this before merging, pass revision="refs/pr/3" to MultiVectorEncoder.

Happy to tweak anything you'd like changed. Please let me know if you have any questions or feedback!

  • Tom Aarsen
tomaarsen changed pull request status to open
Ready to merge
This branch is ready to get merged automatically.

Sign up or log in to comment