Integrate with Sentence Transformers via MultiVectorEncoder

#1
by tomaarsen HF Staff - opened

Hello @xiaoxiaoshadiao and team!

Congratulations on the big release! I'm quite a fan that you're keeping the dimensionality lower here again, there was a bit of a trend to grow the dimensionality (at huge index costs, obviously), and I think it's important to keep it low to keep these models viable. For context, Sentence Transformers is releasing a MultiVectorEncoder class, 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). Your release here is great timing in my opinion!

Heads up, the PR text below was AI-generated and human-reviewed. Here's a summary of the changes as reported by my agent:

Pull Request overview

  • Integrate tencent/EVIE-Preview-4.5B with Sentence Transformers as a multi-vector (ColBERT-style late interaction) retriever via MultiVectorEncoder.

Details

The integration is config-only: no modeling file, no trust_remote_code, and config.json and model.safetensors are untouched. The module pipeline Transformer(feature-extraction) -> Dense(2560->128) -> Normalize -> MultiVectorMask mirrors ColQwen3_5.forward step for step (last hidden state -> custom_text_proj -> L2 normalize -> mask by attention_mask). 1_Dense/model.safetensors is custom_text_proj lifted out of the checkpoint, since AutoModel resolves a bare Qwen3_5Model and would otherwise report those tensors as unexpected. Nothing is retrained.

Bidirectional attention is set through sentence_bert_config.json's config_kwargs (text_config.is_causal=false), which transformers threads into the attention interface, so the 8 full-attention layers become bidirectional while the 24 GatedDeltaNet layers stay recurrent, exactly as enable_bidirectional_attention() intends. A new chat_template.jinja reproduces the two ColQwen3_5Processor prompt formats, and processor_class moves to Qwen3VLProcessor so AutoProcessor resolves without importing colpali-engine. The existing ColQwen3_5 / infer.py / reproduce.py path is untouched and loads the same weights and config as before.

Verified against colpali-engine 0.3.17 in float32: token ids are bit-identical for queries and images, per-item embeddings match at max |diff| 0.0, and the batched MaxSim matrix matches at 0.0.

Unrelated to this PR: model.enable_bidirectional_attention() in your Quick Start does not exist in any released colpali-engine (through 0.3.17), so that snippet raises AttributeError today. Without it the rankings hold but the top-hit MaxSim margins drop by roughly 1.1, so it is worth wiring up.

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

model = MultiVectorEncoder("tencent/EVIE-Preview-4.5B", revision="refs/pr/1")

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([23, 128]) torch.Size([755, 128])

scores = model.similarity(query_embeddings, document_embeddings)
print(scores)
# tensor([[17.2949, 10.7598,  7.9268,  7.3613],
#         [ 6.5547, 13.3711,  6.2627,  6.1104]])
  • 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