Integrate with Sentence Transformers via MultiVectorEncoder

Hello!

Heads up, this PR was AI-generated and human-reviewed. 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).

Here's a summary of the changes as reported by my agent:

## Pull Request overview
* Integrate `vidore/colpali` with [Sentence Transformers](https://www.sbert.net/) 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` module loads the adapter directly onto the PaliGemma backbone through a small `key_mapping` that strips `colpali-engine`'s `model.` wrapper prefix, so no custom modeling code or `trust_remote_code` is needed, only `transformers>=5.15.0` (which ships huggingface/transformers#46766) and `peft`. The frozen `custom_text_proj` (2048 to 128) ships pre-merged as a roughly 1 MB `1_Dense` module. The trained weights are untouched and the existing `colpali-engine` usage keeps working unchanged.

On the query format: this checkpoint predates a tagged `colpali-engine` release and the revision it records is not in the `illuin-tech/colpali` history, so I gave it the August 2024 `Question: ` query format of its near-contemporaries. Current `colpali-engine` no longer sends that format: 0.3.4 changed the prefix from `Question: ` to `Query: ` (illuin-tech/colpali#125), 0.3.11 dropped the trailing newline (illuin-tech/colpali#280), and 0.3.13 dropped the prefix entirely (illuin-tech/colpali#339). This configuration reproduces the training-time format, so its embeddings differ slightly from current `colpali-engine` output, and the README flags this next to the `colpali-engine` snippet. On a ViDoRe v1 check, reproducing the training-time format improved nDCG@5 over the current `colpali-engine` format in 6 of 6 checkpoint x dataset cells measured.

```bash
pip install "sentence-transformers[image] @ git+https://github.com/huggingface/sentence-transformers.git"
```

```python
from sentence_transformers import MultiVectorEncoder

model = MultiVectorEncoder("vidore/colpali", revision="refs/pr/N")

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))
# (23, 128) (1030, 128)

print(model.similarity(query_embeddings, document_embeddings))
# tensor([[17.3789, 17.1055, 15.4727, 15.4082],
# [ 8.3750, 12.3047, 8.5898, 9.0957]])
```

- Tom Aarsen

1_Dense/config.json

@@ -0,0 +1,9 @@
+{
+    "in_features": 2048,
+    "out_features": 128,
+    "bias": true,
+    "activation_function": "torch.nn.modules.linear.Identity",
+    "module_input_name": "token_embeddings",
+    "module_output_name": "token_embeddings",
+    "use_residual": false
+}

1_Dense/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a9b0ae57a26f3f576a8652b1827896ee6e6385674a4309009f05c3f186d8a2d0
+size 1049248

2_Normalize/config.json

@@ -0,0 +1,4 @@
+{
+    "module_input_name": "token_embeddings",
+    "module_output_name": "token_embeddings"
+}

3_MultiVectorMask/config.json

@@ -0,0 +1,3 @@
+{
+    "skiplist_words": []
+}

README.md

@@ -7,6 +7,8 @@ language:
 tags:
 - colpali
 - vidore
+- sentence-transformers
+- multi-vector
 new_version: vidore/colpali-v1.1
 datasets:
 - vidore/colpali_train_set
@@ -46,8 +48,58 @@ We train on an 8 GPU setup with data parallelism, a learning rate of 5e-5 with l
 
 ## Usage
 
-### For best performance, newer models are available (vidore/colpali-v1.2)
+### Using Sentence Transformers
 
+ColPali can be used as a multi-vector (ColBERT-style late interaction) retriever directly with Sentence Transformers via the `MultiVectorEncoder`.
+
+```bash
+pip install "sentence-transformers[image]>=6.0.0"
+```
+
+```python
+from sentence_transformers import MultiVectorEncoder
+
+model = MultiVectorEncoder("tomaarsen/colpali-st")
+
+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, convert_to_tensor=True)
+document_embeddings = model.encode_document(images, convert_to_tensor=True)
+print(f"Query 0 shape:    {tuple(query_embeddings[0].shape)}")
+print(f"Document 0 shape: {tuple(document_embeddings[0].shape)}")
+# Query 0 shape:    (23, 128)
+# Document 0 shape: (1030, 128)
+
+# MaxSim late-interaction scoring (rows = queries, columns = images)
+scores = model.similarity(query_embeddings, document_embeddings)
+print(scores)
+# tensor([[17.3789, 17.1055, 15.4727, 15.4082],
+#         [ 8.3750, 12.3047,  8.5898,  9.0957]])
+```
+
+### Using ColPali Engine
+
+> [!WARNING]
+> Note: current `colpali-engine` no longer sends the query prefix and trailing newline that this
+> checkpoint was trained with. The trailing newline went in 0.3.11 (illuin-tech/colpali#280) and the prefix in 0.3.13 (illuin-tech/colpali#339). The Sentence Transformers
+> configuration in this repository reproduces the original training-time format, so its embeddings differ
+> slightly from current `colpali-engine` output.
+> Release 0.3.4 had already changed the prefix from `Question: ` to `Query: ` (illuin-tech/colpali#125),
+> which this checkpoint predates.
+> The Sentence Transformers configuration also sends `token_type_ids` to the model, which on
+> `transformers` 5.x is what makes PaliGemma build an explicit attention mask at all. Without it no
+> mask is materialized and the shorter queries in a batch attend to their own padding.
+
+> For best performance, newer models are available (vidore/colpali-v1.2)
 
 ```bash
 # This model checkpoint is compatible with version 0.1.1, but not more recent versions of the inference lib

chat_template.jinja

@@ -0,0 +1,10 @@
+{%- for message in messages -%}
+{%- set images = message['content'] | selectattr('type', 'equalto', 'image') | list -%}
+{%- set texts = message['content'] | selectattr('type', 'equalto', 'text') | map(attribute='text') | list -%}
+{%- if images -%}
+{%- for _ in images -%}{{- '<image>' -}}{%- endfor -%}
+{{- texts[0] if texts else 'Describe the image.' -}}
+{%- else -%}
+{{ bos_token }}Question: {{ texts[0] }}{{ '<unused0>' * 5 }}
+{% endif %}
+{% endfor %}

config_sentence_transformers.json

@@ -0,0 +1,18 @@
+{
+    "__version__": {
+        "sentence_transformers": "5.7.0"
+    },
+    "default_prompt_name": null,
+    "model_type": "MultiVectorEncoder",
+    "requirements": {
+        "transformers": {
+            "specifier": ">=5.15",
+            "reason": "Older versions ignore the key_mapping, which silently randomizes the adapter weights."
+        }
+    },
+    "prompts": {
+        "document": "",
+        "query": ""
+    },
+    "similarity_fn_name": null
+}

modules.json

@@ -0,0 +1,26 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "sentence_transformers.base.modules.transformer.Transformer"
+  },
+  {
+    "idx": 1,
+    "name": "1",
+    "path": "1_Dense",
+    "type": "sentence_transformers.base.modules.dense.Dense"
+  },
+  {
+    "idx": 2,
+    "name": "2",
+    "path": "2_Normalize",
+    "type": "sentence_transformers.sentence_transformer.modules.normalize.Normalize"
+  },
+  {
+    "idx": 3,
+    "name": "3",
+    "path": "3_MultiVectorMask",
+    "type": "sentence_transformers.multi_vector_encoder.modules.multi_vector_mask.MultiVectorMask"
+  }
+]

preprocessor_config.json

@@ -1,40 +1,40 @@
-{
-  "_valid_processor_keys": [
-    "images",
-    "do_resize",
-    "size",
-    "resample",
-    "do_rescale",
-    "rescale_factor",
-    "do_normalize",
-    "image_mean",
-    "image_std",
-    "return_tensors",
-    "data_format",
-    "input_data_format",
-    "do_convert_rgb"
-  ],
-  "do_convert_rgb": null,
-  "do_normalize": true,
-  "do_rescale": true,
-  "do_resize": true,
-  "image_mean": [
-    0.5,
-    0.5,
-    0.5
-  ],
-  "image_processor_type": "SiglipImageProcessor",
-  "image_seq_length": 1024,
-  "image_std": [
-    0.5,
-    0.5,
-    0.5
-  ],
-  "processor_class": "PaliGemmaProcessor",
-  "resample": 3,
-  "rescale_factor": 0.00392156862745098,
-  "size": {
-    "height": 448,
-    "width": 448
-  }
-}
+{
+  "_valid_processor_keys": [
+    "images",
+    "do_resize",
+    "size",
+    "resample",
+    "do_rescale",
+    "rescale_factor",
+    "do_normalize",
+    "image_mean",
+    "image_std",
+    "return_tensors",
+    "data_format",
+    "input_data_format",
+    "do_convert_rgb"
+  ],
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "image_processor_type": "SiglipImageProcessor",
+  "image_seq_length": 1024,
+  "image_std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "processor_class": "PaliGemmaProcessor",
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "height": 448,
+    "width": 448
+  }
+}

sentence_bert_config.json

@@ -0,0 +1,31 @@
+{
+    "transformer_task": "feature-extraction",
+    "modality_config": {
+        "text": {
+            "method": "forward",
+            "method_output_name": "last_hidden_state"
+        },
+        "image": {
+            "method": "forward",
+            "method_output_name": "last_hidden_state"
+        },
+        "message": {
+            "method": "forward",
+            "method_output_name": "last_hidden_state",
+            "format": "structured"
+        }
+    },
+    "module_output_name": "token_embeddings",
+    "model_kwargs": {
+        "key_mapping": {
+            "^model\\.": ""
+        }
+    },
+    "processor_kwargs": {
+        "model_input_names": [
+            "input_ids",
+            "attention_mask",
+            "token_type_ids"
+        ]
+    }
+}