Integrate with Sentence Transformers via MultiVectorEncoder

1_Dense/config.json

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

1_Dense/model.safetensors

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

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,3 @@
+{
+    "skiplist_words": []
+}

README.md

@@ -8,6 +8,8 @@ tags:
 - colsmolvlm
 - vidore-experimental
 - vidore
+- sentence-transformers
+- multi-vector
 pipeline_tag: visual-document-retrieval
 ---
 # ColSmolVLM-Instruct-256M: Visual Retriever based on SmolVLM-Instruct-250M with ColBERT strategy
@@ -45,6 +47,53 @@ We train on a 4 GPU setup with data parallelism, a learning rate of 5e-4 with li
 
 ## Usage
 
+### Using Sentence Transformers
+
+ColSmolVLM 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("vidore/colSmol-256M")
+
+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, convert_to_tensor=True)
+document_embeddings = model.encode_document(images, 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:    (27, 128)
+# Document 0 shape: (1135, 128)
+
+# MaxSim late-interaction scoring (rows = queries, columns = images)
+scores = model.similarity(query_embeddings, document_embeddings)
+print(scores)
+# tensor([[18.1855, 16.2119, 11.7363,  9.7974],
+#         [ 9.2637, 15.1357, 10.4395,  8.0791]])
+```
+
+### Using ColPali Engine
+
+> [!WARNING]
+> Note: current `colpali-engine` no longer sends the query prefix and trailing newline that this checkpoint
+> was trained with. The trailing newline went in 0.3.11 (illuin-tech/colpali#280) and the `"Query: "` prefix
+> in 0.3.13 (illuin-tech/colpali#339), and the image document prompt was rewritten in 0.3.9 and again in
+> 0.3.11. The Sentence Transformers configuration in this repository reproduces the original training-time
+> format, so its embeddings differ from current `colpali-engine` output.
+
 Make sure `colpali-engine` is installed from source or with a version superior to 0.3.5 (main branch from the repo currently).
 `transformers` version must be > 4.46.2.
 

adapter_config.json

@@ -22,7 +22,7 @@
   "r": 32,
   "rank_pattern": {},
   "revision": null,
-  "target_modules": "(.*(model.text_model).*(down_proj|gate_proj|up_proj|k_proj|q_proj|v_proj|o_proj).*$|.*(custom_text_proj).*$)",
+  "target_modules": "(.*(text_model).*(down_proj|gate_proj|up_proj|k_proj|q_proj|v_proj|o_proj).*$|.*(custom_text_proj).*$)",
   "task_type": "FEATURE_EXTRACTION",
   "use_dora": false,
   "use_rslora": false

additional_chat_templates/sentence_transformers.jinja

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

chat_template.jinja

@@ -0,0 +1,2 @@
+<|im_start|>{% for message in messages %}{{message['role'] | capitalize}}{% if message['content'][0]['type'] == 'image' %}{{':'}}{% else %}{{': '}}{% endif %}{% for line in message['content'] %}{% if line['type'] == 'text' %}{{line['text']}}{% elif line['type'] == 'image' %}{{ '<image>' }}{% endif %}{% endfor %}<end_of_utterance>
+{% endfor %}{% if add_generation_prompt %}{{ 'Assistant:' }}{% endif %}

chat_template.json

@@ -1,3 +0,0 @@
-{
-  "chat_template": "<|im_start|>{% for message in messages %}{{message['role'] | capitalize}}{% if message['content'][0]['type'] == 'image' %}{{':'}}{% else %}{{': '}}{% endif %}{% for line in message['content'] %}{% if line['type'] == 'text' %}{{line['text']}}{% elif line['type'] == 'image' %}{{ '<image>' }}{% endif %}{% endfor %}<end_of_utterance>\n{% endfor %}{% if add_generation_prompt %}{{ 'Assistant:' }}{% endif %}"
-}

config.json

@@ -1,266 +0,0 @@
-{
-  "_name_or_path": "models/SmolVLM-Instruct-250M",
-  "architectures": [
-    "ColIdefics3"
-  ],
-  "image_token_id": 49190,
-  "model_type": "idefics3",
-  "scale_factor": 4,
-  "text_config": {
-    "_attn_implementation_autoset": false,
-    "_flash_attn_2_enabled": true,
-    "_name_or_path": "None",
-    "add_cross_attention": false,
-    "architectures": [
-      "VLlama3ForCausalLM"
-    ],
-    "attention_bias": false,
-    "attention_dropout": 0.0,
-    "bad_words_ids": null,
-    "begin_suppress_tokens": null,
-    "bos_token_id": 1,
-    "chunk_size_feed_forward": 0,
-    "cross_attention_hidden_size": null,
-    "decoder_start_token_id": null,
-    "diversity_penalty": 0.0,
-    "do_sample": false,
-    "early_stopping": false,
-    "encoder_no_repeat_ngram_size": 0,
-    "eos_token_id": 2,
-    "exponential_decay_length_penalty": null,
-    "finetuning_task": null,
-    "forced_bos_token_id": null,
-    "forced_eos_token_id": null,
-    "head_dim": 64,
-    "hidden_act": "silu",
-    "hidden_size": 576,
-    "id2label": {
-      "0": "LABEL_0",
-      "1": "LABEL_1"
-    },
-    "initializer_range": 0.041666666666666664,
-    "intermediate_size": 1536,
-    "is_decoder": false,
-    "is_encoder_decoder": false,
-    "is_llama_config": true,
-    "label2id": {
-      "LABEL_0": 0,
-      "LABEL_1": 1
-    },
-    "length_penalty": 1.0,
-    "max_length": 20,
-    "max_position_embeddings": 8192,
-    "min_length": 0,
-    "mlp_bias": false,
-    "model_type": "llama",
-    "neftune_noise_alpha": 0.0,
-    "no_repeat_ngram_size": 0,
-    "num_attention_heads": 9,
-    "num_beam_groups": 1,
-    "num_beams": 1,
-    "num_hidden_layers": 30,
-    "num_key_value_heads": 3,
-    "num_return_sequences": 1,
-    "output_attentions": false,
-    "output_hidden_states": false,
-    "output_scores": false,
-    "pad_token_id": 2,
-    "perceiver_config": {
-      "_attn_implementation_autoset": false,
-      "_name_or_path": "",
-      "add_cross_attention": false,
-      "architectures": null,
-      "attention_dropout": 0.0,
-      "bad_words_ids": null,
-      "begin_suppress_tokens": null,
-      "bos_token_id": null,
-      "chunk_size_feed_forward": 0,
-      "cross_attention_hidden_size": null,
-      "decoder_start_token_id": null,
-      "diversity_penalty": 0.0,
-      "do_sample": false,
-      "early_stopping": false,
-      "encoder_no_repeat_ngram_size": 0,
-      "eos_token_id": null,
-      "exponential_decay_length_penalty": null,
-      "finetuning_task": null,
-      "forced_bos_token_id": null,
-      "forced_eos_token_id": null,
-      "hidden_act": "silu",
-      "id2label": {
-        "0": "LABEL_0",
-        "1": "LABEL_1"
-      },
-      "is_decoder": false,
-      "is_encoder_decoder": false,
-      "label2id": {
-        "LABEL_0": 0,
-        "LABEL_1": 1
-      },
-      "length_penalty": 1.0,
-      "max_length": 20,
-      "min_length": 0,
-      "model_type": "vllama3",
-      "no_repeat_ngram_size": 0,
-      "num_beam_groups": 1,
-      "num_beams": 1,
-      "num_key_value_heads": 1,
-      "num_return_sequences": 1,
-      "output_attentions": false,
-      "output_hidden_states": false,
-      "output_scores": false,
-      "pad_token_id": null,
-      "prefix": null,
-      "problem_type": null,
-      "pruned_heads": {},
-      "qk_layer_norms_perceiver": false,
-      "remove_invalid_values": false,
-      "repetition_penalty": 1.0,
-      "resampler_depth": 6,
-      "resampler_head_dim": 96,
-      "resampler_n_heads": 16,
-      "resampler_n_latents": 64,
-      "return_dict": true,
-      "return_dict_in_generate": false,
-      "sep_token_id": null,
-      "suppress_tokens": null,
-      "task_specific_params": null,
-      "temperature": 1.0,
-      "tf_legacy_loss": false,
-      "tie_encoder_decoder": false,
-      "tie_word_embeddings": true,
-      "tokenizer_class": null,
-      "top_k": 50,
-      "top_p": 1.0,
-      "torch_dtype": null,
-      "torchscript": false,
-      "transformers_version": "4.46.0",
-      "typical_p": 1.0,
-      "use_bfloat16": false
-    },
-    "pixel_shuffle_factor": 4,
-    "prefix": null,
-    "pretraining_tp": 1,
-    "problem_type": null,
-    "pruned_heads": {},
-    "qk_layer_norms": false,
-    "remove_invalid_values": false,
-    "repetition_penalty": 1.0,
-    "return_dict": true,
-    "return_dict_in_generate": false,
-    "rms_norm_eps": 1e-05,
-    "rope_interleaved": false,
-    "rope_scaling": null,
-    "rope_theta": 100000,
-    "sep_token_id": null,
-    "suppress_tokens": null,
-    "task_specific_params": null,
-    "temperature": 1.0,
-    "tf_legacy_loss": false,
-    "tie_encoder_decoder": false,
-    "tie_word_embeddings": false,
-    "tokenizer_class": null,
-    "top_k": 50,
-    "top_p": 1.0,
-    "torch_dtype": "bfloat16",
-    "torchscript": false,
-    "transformers.js_config": {
-      "kv_cache_dtype": {
-        "fp16": "float16",
-        "q4f16": "float16"
-      }
-    },
-    "typical_p": 1.0,
-    "use_bfloat16": false,
-    "use_cache": true,
-    "use_resampler": false,
-    "vocab_size": 49280
-  },
-  "tie_word_embeddings": false,
-  "torch_dtype": "bfloat16",
-  "transformers_version": "4.46.3",
-  "use_cache": true,
-  "vision_config": {
-    "_attn_implementation_autoset": false,
-    "_name_or_path": "",
-    "add_cross_attention": false,
-    "architectures": null,
-    "attention_dropout": 0.0,
-    "bad_words_ids": null,
-    "begin_suppress_tokens": null,
-    "bos_token_id": null,
-    "chunk_size_feed_forward": 0,
-    "cross_attention_hidden_size": null,
-    "decoder_start_token_id": null,
-    "diversity_penalty": 0.0,
-    "do_sample": false,
-    "early_stopping": false,
-    "encoder_no_repeat_ngram_size": 0,
-    "eos_token_id": null,
-    "exponential_decay_length_penalty": null,
-    "finetuning_task": null,
-    "forced_bos_token_id": null,
-    "forced_eos_token_id": null,
-    "hidden_act": "gelu_pytorch_tanh",
-    "hidden_size": 768,
-    "id2label": {
-      "0": "LABEL_0",
-      "1": "LABEL_1"
-    },
-    "image_size": 512,
-    "initializer_range": 0.02,
-    "intermediate_size": 3072,
-    "is_decoder": false,
-    "is_encoder_decoder": false,
-    "label2id": {
-      "LABEL_0": 0,
-      "LABEL_1": 1
-    },
-    "layer_norm_eps": 1e-06,
-    "length_penalty": 1.0,
-    "max_image_size": {
-      "longest_edge": 512
-    },
-    "max_length": 20,
-    "min_length": 0,
-    "model_type": "idefics3",
-    "no_repeat_ngram_size": 0,
-    "num_attention_heads": 12,
-    "num_beam_groups": 1,
-    "num_beams": 1,
-    "num_channels": 3,
-    "num_hidden_layers": 12,
-    "num_return_sequences": 1,
-    "output_attentions": false,
-    "output_hidden_states": false,
-    "output_scores": false,
-    "pad_token_id": null,
-    "patch_size": 16,
-    "prefix": null,
-    "problem_type": null,
-    "pruned_heads": {},
-    "remove_invalid_values": false,
-    "repetition_penalty": 1.0,
-    "return_dict": true,
-    "return_dict_in_generate": false,
-    "sep_token_id": null,
-    "size": {
-      "longest_edge": 2048
-    },
-    "suppress_tokens": null,
-    "task_specific_params": null,
-    "temperature": 1.0,
-    "tf_legacy_loss": false,
-    "tie_encoder_decoder": false,
-    "tie_word_embeddings": false,
-    "tokenizer_class": null,
-    "top_k": 50,
-    "top_p": 1.0,
-    "torch_dtype": null,
-    "torchscript": false,
-    "typical_p": 1.0,
-    "use_base_siglip": true,
-    "use_bfloat16": false
-  },
-  "vocab_size": 49280
-}

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

preprocessor_config.json

@@ -19,7 +19,7 @@
   "max_image_size": {
     "longest_edge": 512
   },
-  "processor_class": "ColIdefics3Processor",
+  "processor_class": "Idefics3Processor",
   "resample": 1,
   "rescale_factor": 0.00392156862745098,
   "size": {

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

tokenizer_config.json

@@ -1168,14 +1168,13 @@
     "<end_of_utterance>"
   ],
   "bos_token": "<|im_start|>",
-  "chat_template": "<|im_start|>{% for message in messages %}{{message['role'] | capitalize}}{% if message['content'][0]['type'] == 'image' %}{{':'}}{% else %}{{': '}}{% endif %}{% for line in message['content'] %}{% if line['type'] == 'text' %}{{line['text']}}{% elif line['type'] == 'image' %}{{ '<image>' }}{% endif %}{% endfor %}<end_of_utterance>\n{% endfor %}{% if add_generation_prompt %}{{ 'Assistant:' }}{% endif %}",
   "clean_up_tokenization_spaces": false,
   "eos_token": "<|im_end|>",
   "extra_special_tokens": {},
   "legacy": false,
   "model_max_length": 8192,
   "pad_token": "<|im_end|>",
-  "processor_class": "ColIdefics3Processor",
+  "processor_class": "Idefics3Processor",
   "tokenizer_class": "GPT2Tokenizer",
   "truncation_side": "left",
   "unk_token": "<|endoftext|>",