Integrate with Sentence Transformers via MultiVectorEncoder

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("vidore/colpali")
+
+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": "6.0.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

@@ -14,7 +14,7 @@
     "input_data_format",
     "do_convert_rgb"
   ],
-  "do_convert_rgb": null,
+  "do_convert_rgb": true,
   "do_normalize": true,
   "do_rescale": true,
   "do_resize": true,

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"
+        ]
+    }
+}

tokenizer_config.json

@@ -10969,7 +10969,6 @@
   "bos_token": "<bos>",
   "clean_up_tokenization_spaces": false,
   "eos_token": "<eos>",
-  "max_length": 50,
   "model_max_length": 1000000000000000019884624838656,
   "pad_token": "<pad>",
   "processor_class": "PaliGemmaProcessor",