Add Sentence Transformers support and ship bidirectional attention fix

Applies the MultiVectorEncoder integration from PR #1 by @tomaarsen , verified against the ColPali path (identical shapes, MaxSim within bfloat16 noise).

Also ships bidirectional.py: released colpali-engine (through 0.3.17) exposes no enable_bidirectional_attention, so infer.py and reproduce.py raised AttributeError on a plain install, and running causal shifts top-hit MaxSim by about 1.1.

1_Dense/config.json

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

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

@@ -24,6 +24,7 @@ tags:
 - document-retrieval
 - multimodal
 - state-of-the-art
+- sentence-transformers
 base_model:
 - Qwen/Qwen3.5-4B
 datasets:
@@ -181,6 +182,8 @@ Queries in **English, French, German, Italian, Spanish, Portuguese and Chinese**
 
 ## ⚡ Quick Start
 
+ColPali Engine is the reference path — every number on this card comes from it. A [Sentence Transformers](#sentence-transformers) path is also available for late-interaction pipelines already built on that API.
+
 ### Installation
 
 ```bash
@@ -194,6 +197,8 @@ import torch
 from PIL import Image
 from colpali_engine.models import ColQwen3_5, ColQwen3_5Processor
 
+from bidirectional import enable_bidirectional_attention
+
 model_id = "tencent/EVIE-Preview-4.5B"
 
 # 1. Load model and enable bidirectional attention
@@ -203,7 +208,7 @@ model = ColQwen3_5.from_pretrained(
     device_map="cuda",
     attn_implementation="flash_attention_2",
 ).eval()
-model.enable_bidirectional_attention()
+enable_bidirectional_attention(model)
 
 # 2. Load processor
 processor = ColQwen3_5Processor.from_pretrained(model_id)
@@ -225,7 +230,7 @@ scores = processor.score(query_embeddings, image_embeddings)
 print("Late-interaction retrieval scores:", scores)
 ```
 
-> ⚠️ Call `model.enable_bidirectional_attention()` and reset `model.rope_deltas = None` before every query forward pass. Both are required to reach the scores above.
+> ⚠️ Apply `enable_bidirectional_attention(model)` once after loading, and reset `model.rope_deltas = None` before every query forward pass. Both are required to reach the scores above. [`bidirectional.py`](bidirectional.py) ships with this repository because released `colpali-engine` builds ColQwen3.5 with causal masks.
 
 ### CLI
 
@@ -233,6 +238,42 @@ print("Late-interaction retrieval scores:", scores)
 python infer.py --query "Quarterly revenue report" --image page_1.png --image page_2.png
 ```
 
+### Sentence Transformers
+
+EVIE also loads as a [Sentence Transformers](https://www.sbert.net/) `MultiVectorEncoder`, exposing the familiar `encode_query` / `encode_document` / `similarity` API with MaxSim scoring built in. Bidirectional attention is baked into the shipped configuration, so no extra call is needed.
+
+`MultiVectorEncoder` requires Sentence Transformers 6.0.0, which is not on PyPI yet — install from source until it is released:
+
+```bash
+pip install "sentence-transformers[image] @ git+https://github.com/huggingface/sentence-transformers.git"
+```
+
+```python
+from sentence_transformers import MultiVectorEncoder
+
+model = MultiVectorEncoder("tencent/EVIE-Preview-4.5B")
+
+queries = ["What key insights are presented on this page?"]
+documents = ["document_page.png"]
+
+query_embeddings = model.encode_query(queries)
+document_embeddings = model.encode_document(documents)
+
+scores = model.similarity(query_embeddings, document_embeddings)
+```
+
+Documents may be file paths, URLs or `PIL.Image` objects. Text passed to `encode_document` is rendered as a query, since this model has no separate text-document format.
+
+The default page budget is the 768-token tier. To score the 1,792-token tier, raise the pixel budget through `processor_kwargs`:
+
+```python
+model = MultiVectorEncoder(
+    "tencent/EVIE-Preview-4.5B",
+    model_kwargs={"attn_implementation": "flash_attention_2", "device_map": "cuda:0"},
+    processor_kwargs={"size": {"longest_edge": 1792 * 32 * 32, "shortest_edge": 65536}},
+)
+```
+
 ---
 
 ## 🔬 Reproducing

bidirectional.py

@@ -0,0 +1,28 @@
+"""Bidirectional attention for ColQwen3.5 retrieval.
+
+Released `colpali-engine` (through 0.3.17) builds ColQwen3.5 with the causal
+Qwen3.5 masks it inherits from the generative backbone. EVIE was trained and
+evaluated with the full-attention layers encoder-ized, so the checkpoint must be
+switched before it reproduces the reported scores.
+
+Qwen3.5 interleaves GatedDeltaNet (`linear_attention`) and `full_attention`
+layers. Only the full-attention layers are flipped here; the recurrent layers
+are order-dependent by construction and are left untouched.
+"""
+
+from typing import Any
+
+_ATTENTION_CLASSES = ("Qwen3_5Attention", "Qwen3Attention")
+
+
+def enable_bidirectional_attention(model: Any) -> None:
+    """Encoder-ize the full-attention layers of a ColQwen3.5 model, in place."""
+    config = getattr(model, "config", None)
+    for cfg in (config, getattr(config, "text_config", None)):
+        # `create_causal_mask` falls back to `create_bidirectional_mask` on this flag.
+        if cfg is not None:
+            cfg.is_causal = False
+
+    for module in model.modules():
+        if module.__class__.__name__ in _ATTENTION_CLASSES and hasattr(module, "is_causal"):
+            module.is_causal = False

chat_template.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 -%}
+{{- ns.text + '<|endoftext|>' * 10 -}}
+{%- endif -%}
+{%- endfor -%}

config_sentence_transformers.json

@@ -0,0 +1,9 @@
+{
+    "__version__": {
+        "sentence_transformers": "6.0.0"
+    },
+    "model_type": "MultiVectorEncoder",
+    "similarity_fn_name": "maxsim",
+    "prompts": {},
+    "default_prompt_name": null
+}

infer.py

@@ -8,6 +8,8 @@ from PIL import Image
 
 from colpali_engine.models import ColQwen3_5, ColQwen3_5Processor
 
+from bidirectional import enable_bidirectional_attention
+
 
 def load(model_id: str, device: str = "cuda"):
     model = ColQwen3_5.from_pretrained(
@@ -16,7 +18,7 @@ def load(model_id: str, device: str = "cuda"):
         device_map=device,
         attn_implementation="flash_attention_2",
     ).eval()
-    model.enable_bidirectional_attention()
+    enable_bidirectional_attention(model)
     return model, ColQwen3_5Processor.from_pretrained(model_id)
 
 

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

processor_config.json

@@ -25,7 +25,7 @@
     },
     "temporal_patch_size": 2
   },
-  "processor_class": "ColQwen3_5Processor",
+  "processor_class": "Qwen3VLProcessor",
   "video_processor": {
     "do_convert_rgb": true,
     "do_normalize": true,

reproduce.py

@@ -28,6 +28,8 @@ from PIL import Image
 from colpali_engine.models import ColQwen3_5, ColQwen3_5Processor
 from colpali_engine.utils.maxsim import maxsim_inbatch
 
+from bidirectional import enable_bidirectional_attention
+
 V1 = [
     "arxivqa_test_subsampled",
     "docvqa_test_subsampled",
@@ -241,7 +243,7 @@ def main() -> int:
     model = ColQwen3_5.from_pretrained(
         args.model, torch_dtype=torch.bfloat16, attn_implementation=args.attn
     )
-    model.enable_bidirectional_attention()
+    enable_bidirectional_attention(model)
     model = model.to("cuda").eval()
 
     result, cache = {}, {}

sentence_bert_config.json

@@ -0,0 +1,25 @@
+{
+    "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,
+    "config_kwargs": {
+        "text_config": {
+            "is_causal": false
+        }
+    }
+}

tokenizer_config.json

@@ -21,8 +21,9 @@
     "vision_eos_token": "<|vision_end|>"
   },
   "pad_token": "<|endoftext|>",
+  "padding_side": "left",
   "pretokenize_regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
-  "processor_class": "ColQwen3_5Processor",
+  "processor_class": "Qwen3VLProcessor",
   "split_special_tokens": false,
   "tokenizer_class": "Qwen2Tokenizer",
   "unk_token": null,