Integrate with Sentence Transformers via MultiVectorEncoder

#2
by tomaarsen HF Staff - opened

Hello!

A few days ago I released the new MultiVectorEncoder class in Sentence Transformers v6.0, and the multilingual ColQwen2.5 models look like great fits for it. If you're interested, 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.

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 Metric-AI/ColQwen2.5-3b-multilingual-v1.0 with Sentence Transformers as a multi-vector (ColBERT-style late interaction) retriever via MultiVectorEncoder.

Details

The integration is config-only: no new modeling code, and adapter_model.safetensors is untouched. The module pipeline is Transformer(feature-extraction) -> Dense(2048 -> 128) -> Normalize -> MultiVectorMask, which mirrors ColQwen2_5.forward step for step. sentence_bert_config.json carries model_kwargs.key_mapping of ^model\. to language_model., the same remap ColQwen2_5._checkpoint_conversion_mapping performs, so AutoModel resolves a bare Qwen2_5_VLModel from Metric-AI/colqwen2.5-3b-base and the LoRA loads on top of it through the usual adapter path. custom_text_proj sits outside that prefix and is reported unexpected by the backbone, so it ships separately as 1_Dense/model.safetensors: the base projection with the LoRA delta merged in (lora_alpha / r is 1.0), in float32. colpali-engine and peft are therefore not needed at inference.

additional_chat_templates/sentence_transformers.jinja reproduces the ColQwen2_5_Processor formats, so encode_query and encode_document produce the same token ids as process_queries and process_images without any manual prefixing. The query branch uses the training-time format, "Query: " + query followed by ten <|endoftext|> tokens. That prefix was part of ColQwen2_5_Processor from the first ColQwen2.5 commit through 0.3.12 (0.3.11 moved the value to the shared base class unchanged, illuin-tech/colpali#280) and was removed in 0.3.13 (illuin-tech/colpali#339), well after this checkpoint was trained, so current colpali-engine sends a slightly different query than the model saw during training. The README notes this in the ColPali Engine section. If you would rather the configuration follow the current library instead, that is a one-line change to the template and I am happy to make it.

Three smaller points. The legacy chat_template.json is converted to chat_template.jinja with identical content, because transformers 5.x refuses to load a repository that has both a legacy chat_template.json and an additional_chat_templates/ directory. A minimal processor_config.json names the stock Qwen2_5_VLProcessor, so AutoProcessor resolves without colpali-engine importable. And tokenizer_config.json gains "padding_side": "left", which is what ColQwen2_5_Processor.__init__ sets in code, so left padding now survives a plain AutoTokenizer load too.

Verified against your own pipeline (ColQwen2_5 plus ColQwen2_5_Processor with the training-time prefix restored) in float32: the query and document token ids are identical, every embedding matches at a minimum per-token cosine of 1.000000 with a maximum absolute difference of 1.8e-07, and the MaxSim score matrix matches at 9.5e-07. At the default bfloat16 load and full page resolution the same comparison gives a minimum per-token cosine of 0.99998 and a maximum score difference of 0.011, which is bfloat16 rounding. For reference, on a 400-query NanoViDoRe v3 subset this integration scores 0.5763 nDCG@10.

Also worth flagging, unrelated to this PR: preprocessor_config.json carries the stock Qwen2.5-VL max_pixels of 12845056 rather than a value matching the "768 image patches at most" in the model card, so a full page becomes roughly 4000 to 5000 visual tokens for anyone who does not pass max_num_visual_tokens. The configuration here is left exactly as shipped so both libraries behave identically, and the README explains how to cap it.

Added files:

  • modules.json, sentence_bert_config.json, config_sentence_transformers.json: the module pipeline, the transformer and processing options, and the model metadata.
  • 1_Dense/config.json and 1_Dense/model.safetensors: the LoRA-merged custom_text_proj projection.
  • 2_Normalize/config.json, 3_MultiVectorMask/config.json: L2 normalization and the padding mask.
  • chat_template.jinja: the existing chat template, converted from chat_template.json.
  • additional_chat_templates/sentence_transformers.jinja: the query and document prompt formats.
  • processor_config.json: names Qwen2_5_VLProcessor.

Modified files:

  • tokenizer_config.json: adds "padding_side": "left".
  • README.md: adds the sentence-transformers and multi-vector tags and a Sentence Transformers usage section.

Removed files:

  • chat_template.json: replaced by chat_template.jinja.
pip install "sentence-transformers[image]>=6.0.0"
from sentence_transformers import MultiVectorEncoder

model = MultiVectorEncoder("Metric-AI/ColQwen2.5-3b-multilingual-v1.0")

queries = [
    "What is the variable represented on the y-axis of the graph?",
    "Total outlay is maximum in which year?",
]
images = [
    "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(images)
print(f"Query 0 shape:    {tuple(query_embeddings[0].shape)}")
print(f"Document 0 shape: {tuple(document_embeddings[0].shape)}")
# Query 0 shape:    (25, 128)
# Document 0 shape: (4115, 128)

# MaxSim late-interaction scoring (rows = queries, columns = images)
scores = model.similarity(query_embeddings, document_embeddings)
print(scores)
# tensor([[15.6797, 12.8027, 12.3555, 12.2031],
#         [ 8.9121, 15.5352,  9.7109,  8.1387]])

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

Note that none of the existing behaviour changes. The colpali-engine path keeps working exactly as before, this only adds a second way to run the model.

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