Integrate with Sentence Transformers via MultiVectorEncoder

1_Dense/config.json

@@ -0,0 +1,8 @@
+{
+  "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"
+}

1_Dense/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:75e5f5783ba184514e1bf9657296dbe12d78853233f16b4098271350738a235a
+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,5 @@
+{
+  "skiplist_words": [],
+  "skiplist_tasks": [],
+  "keep_only_token_ids": null
+}

README.md

@@ -20,6 +20,8 @@ tags:
 - multimodal_embedding
 - multilingual_embedding
 - Text-to-Visual Document (T→VD) retrieval
+- sentence-transformers
+- multi-vector
 library_name: peft
 pipeline_tag: visual-document-retrieval
 ---
@@ -85,6 +87,68 @@ KeyError: 'qwen2_5_vl'
 
 ## Usage
 
+### Using Sentence Transformers
+
+This checkpoint can be used as a multi-vector (ColBERT-style late interaction) retriever with Sentence Transformers via the `MultiVectorEncoder`:
+
+```bash
+pip install "sentence-transformers[image]>=6.0.0"
+```
+
+```python
+from sentence_transformers import MultiVectorEncoder
+
+model = MultiVectorEncoder("Metric-AI/ColQwen2.5-3b-multilingual-v1.0")
+
+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)
+document_embeddings = model.encode_document(images)
+print(f"Query 0 shape:    {tuple(query_embeddings[0].shape)}")
+print(f"Document 0 shape: {tuple(document_embeddings[0].shape)}")
+# Query 0 shape:    (25, 128)
+# Document 0 shape: (4115, 128)
+
+# MaxSim late-interaction scoring (rows = queries, columns = images)
+scores = model.similarity(query_embeddings, document_embeddings)
+print(scores)
+# tensor([[15.6797, 12.8027, 12.3555, 12.2031],
+#         [ 8.9121, 15.5352,  9.7109,  8.1387]])
+```
+
+`encode_query` applies the `Query: ` prefix and the ten `<|endoftext|>` augmentation tokens, `encode_document` applies the visual prompt, and both apply the 128-dimensional projection, the L2 normalization and the padding mask. Documents can be file paths, URLs or `PIL.Image` objects, and both methods accept a `batch_size`.
+
+The scores above come from the plain load, which uses the base checkpoint's `bfloat16` weights. Loading options are forwarded through `model_kwargs`:
+
+```python
+model = MultiVectorEncoder(
+    "Metric-AI/ColQwen2.5-3b-multilingual-v1.0",
+    model_kwargs={"dtype": "float32", "attn_implementation": "sdpa", "device_map": "cuda:0"},
+)
+```
+
+Note that `preprocessor_config.json` in this repository keeps the stock Qwen2.5-VL `max_pixels` of 12845056, so a full page turns into roughly 4000 to 5000 visual tokens rather than the 768 patches mentioned above. Pass `processor_kwargs={"max_pixels": 768 * 28 * 28}` to cap it, which is the Sentence Transformers equivalent of `ColQwen2_5_Processor.from_pretrained(..., max_num_visual_tokens=768)`.
+
+### Using ColPali Engine
+
+> [!WARNING]
+> Current `colpali-engine` no longer sends the `Query: ` prefix that this checkpoint was trained
+> with. It was dropped from `ColQwen2_5_Processor` in 0.3.11
+> ([illuin-tech/colpali#280](https://github.com/illuin-tech/colpali/pull/280)). The Sentence
+> Transformers configuration in this repository reproduces the original training-time format, so
+> its query embeddings differ slightly from current `colpali-engine` output. Install
+> `colpali-engine<0.3.11`, or set `processor.query_prefix = "Query: "` after loading, to get the
+> training-time format from the snippet below.
+
 ```python
 import torch
 from PIL import Image

additional_chat_templates/sentence_transformers.jinja

@@ -0,0 +1,22 @@
+{%- if task is defined and task == 'query' -%}
+    {%- for message in messages -%}
+        {%- for content in message['content'] -%}
+            {%- if content['type'] == 'text' -%}
+                {{- 'Query: ' + content['text'] -}}
+                {%- for _ in range(10) -%}
+                    {{- '<|endoftext|>' -}}
+                {%- endfor -%}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- endfor -%}
+{%- else -%}
+    {%- for message in messages -%}
+        {%- for content in message['content'] -%}
+            {%- if content['type'] == 'image' -%}
+                {{- '<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|><|endoftext|>' -}}
+            {%- elif content['type'] == 'text' -%}
+                {{- content['text'] -}}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- endfor -%}
+{%- endif -%}

chat_template.jinja

@@ -0,0 +1,7 @@
+{% set image_count = namespace(value=0) %}{% set video_count = namespace(value=0) %}{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|im_start|>system
+You are a helpful assistant.<|im_end|>
+{% endif %}<|im_start|>{{ message['role'] }}
+{% if message['content'] is string %}{{ message['content'] }}<|im_end|>
+{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>
+{% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant
+{% endif %}

chat_template.json

@@ -1,3 +0,0 @@
-{
-  "chat_template": "{% set image_count = namespace(value=0) %}{% set video_count = namespace(value=0) %}{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n{% endif %}<|im_start|>{{ message['role'] }}\n{% if message['content'] is string %}{{ message['content'] }}<|im_end|>\n{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>\n{% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}"
-}

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.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

@@ -0,0 +1,3 @@
+{
+  "processor_class": "Qwen2_5_VLProcessor"
+}

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\\.": "language_model."
+    }
+  },
+  "processing_kwargs": {
+    "chat_template": {
+      "chat_template": "sentence_transformers"
+    }
+  }
+}

tokenizer_config.json

@@ -202,6 +202,7 @@
   "extra_special_tokens": {},
   "model_max_length": 131072,
   "pad_token": "<|endoftext|>",
+  "padding_side": "left",
   "processor_class": "ColQwen2_5Processor",
   "split_special_tokens": false,
   "tokenizer_class": "Qwen2Tokenizer",