Integrate with Sentence Transformers via MultiVectorEncoder

#2
by tomaarsen HF Staff - opened

Hello!

The MultiVectorEncoder class ships in the next Sentence Transformers release, 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).

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 vidore/colSmol-256M with Sentence Transformers as a multi-vector (ColBERT-style late interaction) retriever via MultiVectorEncoder.

Details

This adds a Sentence Transformers loading path on top of the existing LoRA adapter, exposing the usual model.encode_query(...) / model.encode_document(...) / model.similarity(...) API with MaxSim scoring. The stock Transformer loads the adapter onto a bare Idefics3 / SmolVLM backbone. This needed removing a shipped config.json that made transformers treat the repo as a full checkpoint (so it now falls back to the adapter's base model), and relaxing the target_modules regex from model.text_model to text_model so PEFT finds the modules on the bare backbone (still matching the colpali-engine wrapper, verified bit-identical), with a small key_mapping stripping the model. wrapper prefix. No custom modeling code or trust_remote_code is needed, only transformers>=5.15.0 (huggingface/transformers#46766) and peft. The projection head ships as the 1_Dense module, the trained weights are untouched, the existing colpali-engine usage is unchanged, and the MaxSim scores match a colpali-engine fp32 baseline within floating point tolerance (embeddings bit-exact, max score difference about 1e-6).

On the format: colpali-engine's ColIdefics3Processor has moved on since this checkpoint's era (the trailing newline went in 0.3.11 / illuin-tech/colpali#280, the Query: prefix in 0.3.13 / illuin-tech/colpali#339, and the image-document prompt was rewritten in 0.3.9 and 0.3.11). This configuration keeps the training-time formats, so its embeddings differ from a current colpali-engine install, and the README flags this next to the colpali-engine snippet.

On the chat template: the checkpoint's conversational template is preserved as chat_template.jinja (rewritten from the original chat_template.json), and the Sentence Transformers query/document prompt format is added as a named template at additional_chat_templates/sentence_transformers.jinja, selected through the ST config. apply_chat_template is unchanged, and colpali-engine formats its inputs directly so it reads neither file.

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

model = MultiVectorEncoder("vidore/colSmol-256M", revision="refs/pr/2")

queries = [
    "What is the variable represented on the y-axis of the graph?",
    "Total outlay is maximum in which year?",
]
documents = [
    f"https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/doc{i}.jpg"
    for i in range(1, 5)
]

query_embeddings = model.encode_query(queries, convert_to_tensor=True)
document_embeddings = model.encode_document(documents, convert_to_tensor=True)
print(tuple(query_embeddings[0].shape), tuple(document_embeddings[0].shape))
# (27, 128) (1135, 128)

print(model.similarity(query_embeddings, document_embeddings))
# tensor([[18.1855, 16.2119, 11.7363,  9.7974],
#         [ 9.2637, 15.1357, 10.4395,  8.0791]])
  • Tom Aarsen
tomaarsen changed pull request status to open
QuentinJG changed pull request status to merged

Sign up or log in to comment