Integrate with Sentence Transformers via MultiVectorEncoder, fix transformers 5.x processor and modeling

1_MultiVectorMask/config.json

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

README.md

@@ -6,6 +6,8 @@ pipeline_tag: visual-document-retrieval
 library_name: transformers
 tags:
 - transformers
+- sentence-transformers
+- multi-vector
 - multimodal_embedding
 - embedding
 - colpali
@@ -30,6 +32,47 @@ The model is trained using a multi-stage strategy that combines large-scale text
 
 ## 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(
+    "OpenSearch-AI/Ops-Colqwen3-4B",
+    trust_remote_code=True,
+    model_kwargs={"dtype": "bfloat16"},
+)
+
+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([25, 2560]) torch.Size([1254, 2560])
+
+# Diagonal should have higher scores
+scores = model.similarity(query_embeddings, image_embeddings)
+print(scores)
+# tensor([[17.4668, 12.9785],
+#         [ 7.0088, 15.9492]], device='cuda:0')
+```
+
+### Transformers
+
+
 **Requirements**
 ```
 pillow

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

modeling_ops_colqwen3.py

@@ -54,7 +54,7 @@ class OpsColQwen3Model(OpsColQwen3PreTrainedModel):
             model.dims = dims
         return model
 
-    def forward(self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, pixel_values: Optional[torch.Tensor] = None, image_grid_thw: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor:
+    def forward(self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, pixel_values: Optional[torch.Tensor] = None, image_grid_thw: Optional[torch.Tensor] = None, mm_token_type_ids: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor:
         has_pixel_values = pixel_values is not None
 
         if has_pixel_values:
@@ -67,6 +67,11 @@ class OpsColQwen3Model(OpsColQwen3PreTrainedModel):
             unpadded = [pixel_sequence[: int(offset.item())] for pixel_sequence, offset in zip(pixel_values, offsets)]
             pixel_values = torch.cat(unpadded, dim=0) if unpadded else None
 
+        extra_kwargs = {}
+        if mm_token_type_ids is not None:
+            # Required for M-RoPE on recent transformers versions. Passed conditionally so
+            # older versions without the argument keep working.
+            extra_kwargs["mm_token_type_ids"] = mm_token_type_ids
         outputs = self.qwen3vl(
             input_ids=input_ids,
             attention_mask=attention_mask,
@@ -75,6 +80,7 @@ class OpsColQwen3Model(OpsColQwen3PreTrainedModel):
             use_cache=False,
             output_hidden_states=True,
             return_dict=True,
+            **extra_kwargs,
         )
 
         last_hidden_states = outputs.last_hidden_state

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

@@ -61,13 +61,44 @@ class OpsColQwen3Processor(Qwen3VLProcessor):
         if self.tokenizer is not None:
             self.tokenizer.padding_side = "left"
 
+    def __call__(self, images=None, text=None, audio=None, videos=None, **kwargs) -> Union[BatchFeature, BatchEncoding]:
+        """Standard processor interface with retrieval formatting.
+
+        Routes plain calls through process_images/process_queries so that generic pipelines
+        (for example Sentence Transformers) produce the same inputs as the dedicated methods:
+        images get the visual prompt and per-image pixel padding, text gets the query prefix
+        and augmentation tokens. `mm_token_type_ids` is requested so the Qwen3-VL M-RoPE
+        requirement of recent transformers versions is satisfied.
+        """
+        # The process_* methods set padding and return_tensors themselves: drop duplicates that
+        # generic callers pass through nested kwargs.
+        for nest in ("text_kwargs", "images_kwargs", "videos_kwargs", "audio_kwargs", "common_kwargs"):
+            sub = kwargs.get(nest)
+            if isinstance(sub, dict):
+                for key in ("padding", "return_tensors", "return_mm_token_type_ids"):
+                    sub.pop(key, None)
+        if images is not None:
+            image_list = images if isinstance(images, list) else [images]
+            flat_images = []
+            for item in image_list:
+                flat_images.extend(item if isinstance(item, list) else [item])
+            return self.process_images(flat_images, return_mm_token_type_ids=True, **kwargs)
+        if text is not None:
+            texts = text if isinstance(text, str) else list(text)
+            return self.process_queries([texts] if isinstance(texts, str) else texts, **kwargs)
+        raise ValueError("You have to specify at least one of `images` or `text`.")
+
+    def _raw_call(self, *args, **kwargs) -> Union[BatchFeature, BatchEncoding]:
+        """The inherited Qwen3VLProcessor call, used internally by the process_* methods."""
+        return super().__call__(*args, **kwargs)
+
     def process_images(self, images: List[Image.Image], return_tensors: str = "pt", **kwargs) -> Union[BatchFeature, BatchEncoding]:
         """
         Process a batch of PIL images for the model.
         """
         images = [image.convert("RGB") for image in images]
 
-        batch_doc = self(text=[self.visual_prompt_prefix] * len(images), images=images, padding="longest", return_tensors=return_tensors, **kwargs)
+        batch_doc = self._raw_call(text=[self.visual_prompt_prefix] * len(images), images=images, padding="longest", return_tensors=return_tensors, **kwargs)
 
         if batch_doc["pixel_values"].numel() == 0:
             return batch_doc
@@ -83,7 +114,7 @@ class OpsColQwen3Processor(Qwen3VLProcessor):
         Process a list of text queries.
         """
         processed_queries = [self.query_prefix + q + self.query_augmentation_token * 10 for q in queries]
-        return self(text=processed_queries, return_tensors=return_tensors, padding="longest", **kwargs)
+        return self._raw_call(text=processed_queries, return_tensors=return_tensors, padding="longest", **kwargs)
 
     @staticmethod
     def score_multi_vector(

sentence_bert_config.json

@@ -0,0 +1,14 @@
+{
+  "transformer_task": "retrieval",
+  "modality_config": {
+    "text": {
+      "method": "forward",
+      "method_output_name": null
+    },
+    "image": {
+      "method": "forward",
+      "method_output_name": null
+    }
+  },
+  "module_output_name": "token_embeddings"
+}