Integrate with Sentence Transformers and implement replace_image_token

1_MultiVectorMask/config.json

@@ -0,0 +1,7 @@
+{
+    "skiplist_words": [],
+    "skiplist_tasks": [
+        "document"
+    ],
+    "keep_only_token_ids": null
+}

README.md

@@ -23,6 +23,8 @@ tags:
   - colbert
   - colqwen3_5
   - multilingual-embedding
+  - sentence-transformers
+  - multi-vector
 ---
 
 # webAI-Official/webAI-ColVec1.1-4b
@@ -121,6 +123,81 @@ checkpoints.
 
 ## 💻 Usage
 
+### Using Sentence Transformers
+
+The checkpoint loads as a `MultiVectorEncoder`, which is available from
+Sentence Transformers v6.0.0:
+
+```bash
+pip install "sentence-transformers[image]>=6.0.0"
+```
+
+`encode_query` and `encode_document` apply the query and document prompt
+formats, the ten query-augmentation tokens, and the L2-normalized
+640-dimensional projection. `similarity` computes the MaxSim score matrix.
+
+```python
+from io import BytesIO
+
+import requests
+from PIL import Image
+from sentence_transformers import MultiVectorEncoder
+
+model = MultiVectorEncoder("webAI-Official/webAI-ColVec1.1-4b", trust_remote_code=True)
+
+queries = [
+    "When was the United States Declaration of Independence proclaimed?",
+    "Who printed the edition of Romeo and Juliet?",
+]
+document_urls = [
+    "https://upload.wikimedia.org/wikipedia/commons/8/89/US-original-Declaration-1776.jpg",
+    "https://upload.wikimedia.org/wikipedia/commons/thumb/4/4c/Romeoandjuliet1597.jpg/500px-Romeoandjuliet1597.jpg",
+]
+documents = [
+    Image.open(BytesIO(requests.get(url, headers={"User-Agent": "Mozilla/5.0"}, timeout=30).content))
+    for url in document_urls
+]
+
+query_embeddings = model.encode_query(queries)
+document_embeddings = model.encode_document(documents)
+print(query_embeddings[0].shape, document_embeddings[0].shape)
+# (27, 640) (523, 640)
+
+scores = model.similarity(query_embeddings, document_embeddings)
+print(scores)
+# tensor([[23.3935,  5.7660],
+#         [ 4.8504, 23.1851]])
+print("Best document per query:", scores.argmax(dim=1))
+# Best document per query: tensor([0, 1])
+```
+
+Documents may also be given as file paths or URLs instead of `PIL.Image`
+objects. The Wikimedia URLs above are fetched with a browser `User-Agent`
+because Wikimedia rejects the default one. Both encode methods accept a
+`batch_size`, and the loading options are forwarded through `model_kwargs`
+(`dtype`, `attn_implementation`, `device_map`) and `processor_kwargs`
+(`max_num_visual_tokens`):
+
+```python
+model = MultiVectorEncoder(
+    "webAI-Official/webAI-ColVec1.1-4b",
+    trust_remote_code=True,
+    model_kwargs={"attn_implementation": "sdpa", "device_map": "cuda:0"},
+    processor_kwargs={"max_num_visual_tokens": 1024},
+)
+```
+
+The scores above come from the plain load, which uses the checkpoint's
+`bfloat16` weights and SDPA. MaxSim is accumulated in `float32` rather
+than in the embedding dtype, matching the scoring path behind the reported
+evaluation, so a `bfloat16` `score_retrieval` call on the same embeddings
+returns coarser values. The scores still move in the second decimal place
+across PyTorch and Transformers builds. Text passed to `encode_document`
+is rendered as a query, because the model defines no text-document format,
+and a warning says so.
+
+### Using transformers
+
 The processor provides the current retrieval API:
 
 - `process_images(images)` prepares one or more document images.

additional_chat_templates/sentence_transformers.jinja

@@ -0,0 +1,19 @@
+{%- for message in messages -%}
+{%- set ns = namespace(has_image=false, text='') -%}
+{%- if message['content'] is string -%}
+{%- set ns.text = message['content'] -%}
+{%- else -%}
+{%- for item in message['content'] -%}
+{%- if 'image' in item or 'image_url' in item or item.type == 'image' -%}
+{%- set ns.has_image = true -%}
+{%- elif 'text' in item -%}
+{%- set ns.text = ns.text + item.text -%}
+{%- endif -%}
+{%- endfor -%}
+{%- endif -%}
+{%- if ns.has_image -%}
+{{- '<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|><|endoftext|>' -}}
+{%- else -%}
+{{- '<|im_start|>user\nQuery: ' + ns.text + '<|endoftext|>' * 10 + '<|im_end|><|endoftext|>' -}}
+{%- 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,14 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "sentence_transformers.base.modules.transformer.Transformer"
+  },
+  {
+    "idx": 1,
+    "name": "1",
+    "path": "1_MultiVectorMask",
+    "type": "sentence_transformers.multi_vector_encoder.modules.multi_vector_mask.MultiVectorMask"
+  }
+]

processing_colqwen35_bidirection.py

@@ -166,6 +166,17 @@ class ColQwen35BidirectionProcessor(ProcessorMixin):
         if cur_min_pixels > max_pixels:
             self.image_processor.min_pixels = max_pixels
 
+    def replace_image_token(self, image_inputs: dict, image_idx: int, **kwargs) -> str:
+        """Expand one ``<|image_pad|>`` placeholder into its per-image token run.
+
+        ``ProcessorMixin.__call__`` delegates placeholder expansion here, and
+        ``apply_chat_template`` calls ``__call__``, so without this both raise
+        ``NotImplementedError`` for image inputs.
+        """
+        merge_length = self.image_processor.merge_size ** 2
+        num_image_tokens = image_inputs["image_grid_thw"][image_idx].prod() // merge_length
+        return self.image_token * num_image_tokens
+
     @classmethod
     def from_pretrained(
         cls,

sentence_bert_config.json

@@ -0,0 +1,27 @@
+{
+    "transformer_task": "retrieval",
+    "modality_config": {
+        "text": {
+            "method": "forward",
+            "method_output_name": null
+        },
+        "image": {
+            "method": "forward",
+            "method_output_name": null
+        },
+        "message": {
+            "method": "forward",
+            "method_output_name": null,
+            "format": "structured"
+        }
+    },
+    "module_output_name": "token_embeddings",
+    "processing_kwargs": {
+        "chat_template": {
+            "chat_template": "sentence_transformers"
+        },
+        "text": {
+            "return_mm_token_type_ids": true
+        }
+    }
+}