Integrate with Sentence Transformers via MultiVectorEncoder

1_Dense/config.json

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

1_Dense/model.safetensors

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

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,5 @@
+{
+  "skiplist_words": [],
+  "skiplist_tasks": [],
+  "keep_only_token_ids": null
+}

README.md

@@ -20,6 +20,7 @@ tags:
   - mistral
   - mteb
   - vidore
+  - sentence-transformers
 base_model: mistralai/Ministral-3B-Instruct
 pipeline_tag: image-text-to-text
 datasets:
@@ -132,6 +133,42 @@ Traditional OCR-based retrieval **loses layout, tables, and visual context**. Ou
 
 ## Installation & Usage
 
+### Sentence Transformers
+
+This model can be used with [Sentence Transformers](https://www.sbert.net/) as a multi-vector (ColBERT-style late interaction) retriever via the `MultiVectorEncoder`:
+
+```bash
+pip install "sentence-transformers[image]>=6.0.0"
+```
+
+```python
+from sentence_transformers import MultiVectorEncoder
+
+model = MultiVectorEncoder("VAGOsolutions/SauerkrautLM-ColMinistral3-3b-v0.1")
+
+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",
+]
+
+query_embeddings = model.encode_query(queries)
+image_embeddings = model.encode_document(images)
+print(query_embeddings[0].shape, image_embeddings[0].shape)
+# torch.Size([24, 128]) torch.Size([427, 128])
+
+# Diagonal should have higher scores
+scores = model.similarity(query_embeddings, image_embeddings)
+print(scores)
+# tensor([[24.0039, 24.0000],
+#         [19.2227, 19.5508]], device='cuda:0')
+```
+
+### SauerkrautLM ColPali
+
 > ⚠️ **Important**: Install our package first (requires transformers 5.0.0+):
 
 ```bash
@@ -139,7 +176,6 @@ pip install "sauerkrautlm-colpali[ministral]"
 # Or: pip install git+https://github.com/VAGOsolutions/sauerkrautlm-colpali && pip install transformers>=5.0.0rc0
 ```
 
-## Usage
 
 ```python
 import torch

additional_chat_templates/sentence_transformers.jinja

@@ -0,0 +1,17 @@
+{%- for message in messages -%}
+{%- set ns = namespace(text='') -%}
+{%- if message['content'] is string -%}
+{%- set ns.text = message['content'] -%}
+{%- else -%}
+{%- for item in message['content'] -%}
+{%- if 'text' in item -%}
+{%- set ns.text = ns.text + item.text -%}
+{%- endif -%}
+{%- endfor -%}
+{%- endif -%}
+{%- if task is defined and task == 'query' -%}
+{{- ns.text + '</s></s></s></s></s></s></s></s></s></s>' -}}
+{%- else -%}
+{{- '<s><s>[SYSTEM_PROMPT]# HOW YOU SHOULD THINK AND ANSWER\n\nFirst draft your thinking process (inner monologue) until you arrive at a response. Format your response using Markdown, and use LaTeX for any mathematical equations. Write both your thoughts and the response in the same language as the input.\n\nYour thinking process must follow the template below:[THINK]Your thoughts or/and draft, like working through an exercise on scratch paper. Be as casual and as long as you want until you are confident to generate the response to the user.[/THINK]Here, provide a self-contained response.[/SYSTEM_PROMPT][INST][IMG][/INST]' -}}
+{%- endif -%}
+{%- endfor -%}

config_sentence_transformers.json

@@ -0,0 +1,9 @@
+{
+  "model_type": "MultiVectorEncoder",
+  "similarity_fn_name": "maxsim",
+  "prompts": {},
+  "default_prompt_name": null,
+  "__version__": {
+    "sentence_transformers": "6.0.0"
+  }
+}

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.base.modules.normalize.Normalize"
+  },
+  {
+    "idx": 3,
+    "name": "3",
+    "path": "3_MultiVectorMask",
+    "type": "sentence_transformers.multi_vector_encoder.modules.multi_vector_mask.MultiVectorMask"
+  }
+]

processor_config.json

@@ -3,7 +3,7 @@
   "image_end_token": "[IMG_END]",
   "image_token": "[IMG]",
   "patch_size": 14,
-  "processor_class": "ColMinistral3Processor",
+  "processor_class": "PixtralProcessor",
   "spatial_merge_size": 2,
   "image_processor": {
     "crop_size": null,

sentence_bert_config.json

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

tokenizer_config.json

@@ -2017,5 +2017,6 @@
   "processor_class": "PixtralProcessor",
   "tokenizer_class": "LlamaTokenizer",
   "unk_token": "<unk>",
-  "use_default_system_prompt": false
+  "use_default_system_prompt": false,
+  "padding_side": "left"
 }