Fix transformers 5.x loading and add Sentence Transformers support (#4)

- Fix transformers 5.x loading and add Sentence Transformers support (1b87f5df5fe780bcf0da3f35460441f9ac708750)


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

1_MultiVectorMask/config.json

@@ -0,0 +1 @@
+{}

README.md

@@ -3,6 +3,8 @@ license: apache-2.0
 license_name: apache-2.0
 license_link: https://www.apache.org/licenses/LICENSE-2.0
 tags:
+- sentence-transformers
+- multi-vector
 - text
 - image
 - video
@@ -156,7 +158,47 @@ pip install transformers pillow requests
 pip install flash-attn --no-build-isolation
 ```
 
-### Inference Code
+### Using Sentence Transformers
+
+`tomoro-colqwen3-embed-4b` can be used as a multi-vector (ColBERT-style late interaction) retriever directly with Sentence Transformers via the `MultiVectorEncoder`.
+
+```bash
+pip install "sentence-transformers[image]>=6.0.0"
+```
+
+```python
+from sentence_transformers import MultiVectorEncoder
+
+model = MultiVectorEncoder("TomoroAI/tomoro-colqwen3-embed-4b", trust_remote_code=True)
+
+queries = [
+    "What is the variable represented on the y-axis of the graph?",
+    "Total outlay is maximum in which year?",
+]
+documents = [
+    f"https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/doc{i}.jpg"
+    for i in range(1, 5)
+]
+
+query_embeddings = model.encode_query(queries, convert_to_tensor=True)
+document_embeddings = model.encode_document(documents, 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:    (23, 320)
+# Document 0 shape: (1251, 320)
+
+# MaxSim late-interaction scoring (rows = queries, columns = images)
+scores = model.similarity(query_embeddings, document_embeddings)
+print(scores)
+# tensor([[12.8291,  9.0850,  6.4121,  5.8818],
+#         [ 4.5928, 10.7617,  4.7812,  5.3145]])
+```
+
+> [!NOTE]
+> Pages are tiled adaptively, so document embeddings vary in length (1251 tokens for the first three
+> example pages, 1271 for the fourth). MaxSim handles that, and `model.similarity` masks the padding.
+
+### Using Transformers
 
 ```python
 import torch
@@ -243,7 +285,7 @@ scores = processor.score_multi_vector(query_embeddings, doc_embeddings)
 print(scores)
 ```
 
-### 🎞️ Lightweight Video Retrieval
+#### 🎞️ Lightweight Video Retrieval
 
 ColQwen3 generalizes to short videos while learning from image-text retrieval task. This minimal example samples a clip with `torchvision`, encodes queries and frames, then pools frame embeddings with a per-dimension max before MaxSim scoring.
 

config_sentence_transformers.json

@@ -0,0 +1,11 @@
+{
+  "__version__": {
+    "sentence_transformers": "6.0.0"
+  },
+  "model_type": "MultiVectorEncoder",
+  "similarity_fn_name": "maxsim",
+  "prompts": {
+    "query": "",
+    "document": ""
+  }
+}

configuration_colqwen3.py

@@ -32,7 +32,7 @@ class ColQwen3Config(PretrainedConfig):
     """Configuration for ColQwen3 retrieval model."""
 
     model_type = "colqwen3"
-    sub_configs: dict[str, Any] = {"vision_config": Qwen3VLVisionConfig, "text_config": Qwen3VLTextConfig}
+    sub_configs = {"vision_config": Qwen3VLVisionConfig, "text_config": Qwen3VLTextConfig}
 
     def __init__(
         self,
@@ -70,9 +70,11 @@ class ColQwen3Config(PretrainedConfig):
         if embed_dim <= 0:
             raise ValueError(f"`embed_dim` must be positive, got {embed_dim}.")
 
-        super().__init__(**kwargs)
+        # transformers 5.x runs config validators inside PretrainedConfig.__init__, and
+        # validate_token_ids reaches get_text_config(), so the sub-configs must exist first.
         self.vision_config = vision_config
         self.text_config = text_config
+        super().__init__(**kwargs)
         self.embed_dim = embed_dim
         self.padding_side = padding_side
         self.initializer_range = initializer_range

modeling_colqwen3.py

@@ -173,6 +173,7 @@ class ColQwen3(ColQwen3PreTrainedModel):
         cache_position: Optional[torch.LongTensor] = None,
         pixel_values_videos: Optional[torch.Tensor] = None,
         video_grid_thw: Optional[torch.LongTensor] = None,
+        mm_token_type_ids: Optional[torch.LongTensor] = None,
     ) -> ColQwen3ForRetrievalOutput:
         r"""
         image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
@@ -201,6 +202,8 @@ class ColQwen3(ColQwen3PreTrainedModel):
             image_grid_thw=image_grid_thw,
             video_grid_thw=video_grid_thw,
             cache_position=cache_position,
+            # transformers 5.x Qwen3-VL needs this to compute multimodal RoPE.
+            **({"mm_token_type_ids": mm_token_type_ids} if mm_token_type_ids is not None else {}),
         )
 
         vlm_hidden_states = vlm_output.hidden_states if output_hidden_states else None
@@ -237,7 +240,10 @@ class ColQwen3(ColQwen3PreTrainedModel):
     def set_output_embeddings(self, new_embeddings):
         self.vlm.set_output_embeddings(new_embeddings)
 
-    def tie_weights(self):
+    def tie_weights(self, *args, **kwargs):
+        # transformers 5.x calls this with missing_keys/recompute_mapping during loading. Those refer
+        # to the wrapper's key namespace, so they must not be forwarded to the inner VLM, whose params
+        # are still meta at that point (torch.equal then fails on meta tensors).
         return self.vlm.tie_weights()
 
     def resize_token_embeddings(

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

sentence_bert_config.json

@@ -0,0 +1,22 @@
+{
+    "transformer_task": "retrieval",
+    "modality_config": {
+        "text": {
+            "method": "forward",
+            "method_output_name": "embeddings"
+        },
+        "image": {
+            "method": "forward",
+            "method_output_name": "embeddings"
+        }
+    },
+    "module_output_name": "token_embeddings",
+    "processor_kwargs": {
+        "max_num_visual_tokens": 1280
+    },
+    "processing_kwargs": {
+        "text": {
+            "return_mm_token_type_ids": true
+        }
+    }
+}