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 rest of the family: the 4b, 2b and 1.7b-Turbo point back here for the details, while ColMinistral3 and ColLFM2 have their own descriptions since those two use different backbones and needed different handling.

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-ColQwen3-8b-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.
  • Add a processor_config.json so that AutoProcessor resolves the vision processor on plain transformers too.

Details

The integration is config-only: no new modeling code, and the weights are untouched. The module pipeline Transformer -> Dense -> Normalize -> MultiVectorMask mirrors the ColQwen3 forward step for step (last hidden state -> custom_text_proj -> L2 normalize -> mask by attention_mask). The custom_text_proj layer is extracted unchanged from the checkpoint into 1_Dense/. A named chat template in additional_chat_templates/sentence_transformers.jinja reproduces the processor's prompt formats: plain text renders the query format (Query: plus the text plus the ten <|endoftext|> augmentation tokens), and image inputs render the visual document format, so no manual prefixing is needed. Since config.json declares the custom colqwen3 model type, sentence_bert_config.json remaps it to the underlying qwen3_vl architecture at load time. tokenizer_config.json gains padding_side: left, matching what the reference processor sets in code.

The new processor_config.json (naming the stock Qwen3VLProcessor) doubles as a fix for plain transformers users: the repository ships no processor config, so AutoProcessor.from_pretrained currently returns a bare tokenizer and silently loses the vision modality. For context on why I think this integration is valuable here: the stock colpali-engine ColQwen3 class hardcodes a 320-dim projection and reads the vision tower's hidden size, so as far as I can tell it silently reinitializes the projection for these checkpoints, and the sauerkrautlm-colpali fork's processor does not construct on transformers 5.15. Happy to be corrected if I missed a working path.

On the query format: the sauerkrautlm-colpali package's shared base processor carries query_prefix = "Query: " (the pre-0.3.13 colpali-engine lineage it was forked from), the ColQwen3 subclass does not override it, and the model card's own usage calls process_queries, which applies it. The chat template therefore renders Query: plus the text plus the ten augmentation tokens, and its token ids are bit-identical to that formula under this repository's tokenizer.

Verified against sauerkrautlm-colpali itself. Preprocessing was compared against a live ColQwen3Processor in a dedicated environment on transformers 4.57.1 (the floor your pyproject.toml declares), where it constructs normally: query and document token ids match exactly, and the pixel values are bit-identical once the layouts are aligned (colpali stacks per image, transformers concatenates). The embeddings were then compared against the fork's ColQwen3 model class on text queries and image documents: per-token cosine similarity 1.0000, MaxSim scores within 0.004. For reference, on a 400-query NanoViDoRe v3 evaluation this integration scores 0.5819 nDCG@10, close to the 0.5916 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-ColQwen3-8b-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([25, 128]) torch.Size([1251, 128])

# Diagonal should have higher scores
scores = model.similarity(query_embeddings, image_embeddings)
print(scores)
# tensor([[14.9121,  9.1426],
#         [ 5.8672, 15.3125]], 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