Integrate with Sentence Transformers via MultiVectorEncoder (#5)

- Integrate with Sentence Transformers via MultiVectorEncoder (9541b60607528dc5d1f475e72f7afd7b4b22ad12)


Co-authored-by: Tom Aarsen <tomaarsen@users.noreply.huggingface.co>

1_Dense/config.json

@@ -0,0 +1,9 @@
+{
+    "in_features": 768,
+    "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:e7ec5433ad5f7366368a1fdecb0245e9948f2b7ab4a60ceefe20413d8fddb6ce
+size 393888

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

@@ -8,6 +8,8 @@ tags:
 - colpali
 - vidore-experimental
 - vidore
+- sentence-transformers
+- multi-vector
 pipeline_tag: visual-document-retrieval
 ---
 
@@ -42,6 +44,48 @@ For more information about ModernVBERT, please check the [arXiv](https://arxiv.o
 
 ## Usage
 
+### Sentence Transformers
+
+ColModernVBERT can be loaded as a multi-vector (ColBERT-style late interaction) retriever with [Sentence Transformers](https://www.sbert.net/) via the `MultiVectorEncoder`, exposing the familiar `encode_query` / `encode_document` / `similarity` API.
+
+```bash
+pip install "sentence-transformers[image]>=6.0.0"
+```
+
+```python
+from sentence_transformers import MultiVectorEncoder
+
+model = MultiVectorEncoder("ModernVBERT/colmodernvbert")
+
+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:    (26, 128)
+# Document 0 shape: (1149, 128)
+
+scores = model.similarity(query_embeddings, document_embeddings)
+print(scores)
+# tensor([[16.7778, 10.3712, 11.8420,  9.0534],
+#         [ 7.3722, 12.0618,  8.1477,  7.9563]])
+```
+
+> [!NOTE]
+> `sentence_transformers.multi_vector_encoder.interpretability.get_n_patches` raises `NotImplementedError` for this model: like other Idefics3-style split-image processors, each page is split into sub-patch token blocks plus a global patch, so the token grid is not a simple rectangle.
+
+### ColPali Engine
+
 **🏎️ If your GPU supports it, we recommend using ModernVBERT with Flash Attention 2 to achieve the highest GPU throughput. To do so, install Flash Attention 2 as follows, then use the model as normal:**
 
 For now, the branch for using colmdernvbert is not yet merged in the official colpali repo, you need to clone the repo and checkout on the right branch to use it.

adapter_config.json

@@ -24,7 +24,7 @@
   "r": 32,
   "rank_pattern": {},
   "revision": null,
-  "target_modules": "(.*(model.text_model).*(Wo|Wqkv|Wi).*$|.*(custom_text_proj).*$)",
+  "target_modules": "(.*(text_model).*(Wo|Wqkv|Wi).*$|.*(custom_text_proj).*$)",
   "task_type": "FEATURE_EXTRACTION",
   "trainable_token_indices": null,
   "use_dora": false,

additional_chat_templates/sentence_transformers.jinja

@@ -0,0 +1,23 @@
+{%- for message in messages -%}
+    {%- if message['content'] is string -%}
+        {%- if task is defined and task == 'query' -%}
+            {{- message['content'] -}}
+            {%- for _ in range(10) -%}{{- '<end_of_utterance>' -}}{%- endfor -%}
+        {%- else -%}
+            {{- message['content'] -}}
+        {%- endif -%}
+    {%- else -%}
+        {%- for content in message['content'] -%}
+            {%- if content['type'] == 'image' -%}
+                {{- '<|begin_of_text|>User:<image>Describe the image.<end_of_utterance>\nAssistant:' -}}
+            {%- elif content['type'] == 'text' -%}
+                {%- if task is defined and task == 'query' -%}
+                    {{- content['text'] -}}
+                    {%- for _ in range(10) -%}{{- '<end_of_utterance>' -}}{%- endfor -%}
+                {%- else -%}
+                    {{- content['text'] -}}
+                {%- endif -%}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- 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

@@ -19,7 +19,7 @@
   "max_image_size": {
     "longest_edge": 512
   },
-  "processor_class": "ColModernVBertProcessor",
+  "processor_class": "Idefics3Processor",
   "resample": 1,
   "rescale_factor": 0.00392156862745098,
   "size": {

processor_config.json

@@ -1,4 +1,4 @@
 {
   "image_seq_len": 64,
-  "processor_class": "ColModernVBertProcessor"
+  "processor_class": "Idefics3Processor"
 }

sentence_bert_config.json

@@ -0,0 +1,30 @@
+{
+    "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",
+    "unpad_inputs": false,
+    "model_kwargs": {
+        "key_mapping": {
+            "^model\\.": ""
+        }
+    },
+    "processing_kwargs": {
+        "chat_template": {
+            "chat_template": "sentence_transformers"
+        }
+    }
+}

tokenizer_config.json

@@ -1271,7 +1271,7 @@
   "pad_token": "[PAD]",
   "pad_token_type_id": 0,
   "padding_side": "left",
-  "processor_class": "ColModernVBertProcessor",
+  "processor_class": "Idefics3Processor",
   "sep_token": "[SEP]",
   "stride": 0,
   "tokenizer_class": "PreTrainedTokenizerFast",