Integrate with Sentence Transformers via MultiVectorEncoder

#1
by tomaarsen HF Staff - opened

Hello!

I recently released the new MultiVectorEncoder class in Sentence Transformers v6.0, and I think the SauerkrautLM Col* 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 matching PRs for the four ColQwen3 checkpoints and ColMinistral3, which use different backbones and so have their own descriptions. This one is extra fun as the smallest of the family and the first LFM2-VL late interaction model I have seen.

Note that this one is a bit different than the others because its current config.json is a bit lacking. I currently fix it via sentence_bert_config.json its config_kwargs, but we can also place these changes directly in config.json if you prefer. Let me know! cc @DavidGF

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 VAGOsolutions/SauerkrautLM-ColLFM2-450M-v0.1 with Sentence Transformers as a multi-vector (ColBERT-style late interaction) retriever via MultiVectorEncoder.
  • Add a Sentence Transformers usage section to the model card, plus the sentence-transformers tag.

Details

The integration is config-only: no new modeling code, and the weights are untouched. The checkpoint's config.json declares a flat custom collfm2 model type without the sub-configs transformers needs, so sentence_bert_config.json supplies the structural lfm2-vl configuration of the LiquidAI/LFM2-VL-450M base at load time (architecture fields only, nothing about precision or identity), plus a key mapping stripping the model. prefix from the checkpoint keys. The module pipeline Transformer -> Dense -> Normalize -> MultiVectorMask mirrors the reference forward step for step, where 1_Dense/ carries the custom_text_proj projection (1024 -> 128) extracted unchanged from the checkpoint. A named chat template in additional_chat_templates/sentence_transformers.jinja reproduces the reference prompt formats: image documents render the <image>Describe the image. user turn, and text queries tokenize as the BOS token plus the raw text, with no augmentation tokens. That last part is unusual for the family, but it is what this repository's own process_texts does (unlike the ColQwen and ColMinistral siblings, whose processors append ten augmentation tokens), so the template matches the training-time format. The shipped processor_config.json already names the stock Lfm2VlProcessor, so nothing to change there.

Verified against the sauerkrautlm-colpali fork's own ColLFM2 pipeline end to end, its processor included rather than only its model on shared batches: query token ids are bit-identical, image token ids and pixel tensors are bit-identical, per-token embeddings match at cosine 0.9999+ on queries, and the MaxSim scores match within 0.012 in bfloat16 (a couple of near-zero-norm page vectors show lower cosine, a bfloat16 artifact with no ranking effect). For reference, on a 400-query NanoViDoRe v3 evaluation this integration scores 0.4249 nDCG@10, close to the 0.4357 reported for this model on the ViDoRe v3 board (our evaluation is a 400-query subset, so the two are not measured identically).

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

model = MultiVectorEncoder("VAGOsolutions/SauerkrautLM-ColLFM2-450M-v0.1")

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",
]

query_embeddings = model.encode_query(queries)
image_embeddings = model.encode_document(images)
print(query_embeddings[0].shape, image_embeddings[0].shape)
# torch.Size([14, 128]) torch.Size([1792, 128])

# Diagonal should have higher scores
scores = model.similarity(query_embeddings, image_embeddings)
print(scores)
# tensor([[13.5820, 13.4766],
#         [ 9.2461,  9.5703]], device='cuda:0')

To try this before merging, pass revision="refs/pr/1" 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
DavidGF changed pull request status to merged

Sign up or log in to comment