Integrate with Sentence Transformers, restore training-time tokenization on newer transformers

README.md

@@ -9,6 +9,7 @@ base_model:
 pipeline_tag: feature-extraction
 
 tags:
+- sentence-transformers
 - multimodal-embedding
 - text-embedding
 - image-embedding
@@ -32,7 +33,7 @@ WeMM-Embedding-2B is a universal multimodal embedding model built on Qwen3.5. It
 
 ```bash
 pip install torch transformers==5.2.0 "qwen-vl-utils[decord]==0.0.14" \
-  sentence-transformers==5.7.0 "accelerate>=1.1.0"
+  "sentence-transformers>=5.7.0" "accelerate>=1.1.0"
 ```
 
 ## Transformers
@@ -85,23 +86,43 @@ Use any subset of the content items to encode text, image, or video independentl
 ## Sentence Transformers
 
 ```python
-from wemm_sentence_transformers import load_wemm_sentence_transformer
+from sentence_transformers import SentenceTransformer
 
 model_id = "tencent/WeMM-Embedding-2B"
-model = load_wemm_sentence_transformer(model_id, device="cuda:0")
-inputs = [
-    "A dog is running on a beach.",
-    {"image": "/path/to/image.jpg", "text": "Represent this image."},
-    {"video": "/path/to/video.mp4", "text": "Represent this video."},
+model = SentenceTransformer(model_id, trust_remote_code=True)
+
+queries = [
+    "Which Llama 4 model variants are available?",
+    "How is mapo tofu prepared?",
 ]
-embeddings = model.encode(
-    inputs,
-    batch_size=1,
-    truncate_dim=2048,
-    normalize_embeddings=True,
-)
+documents = [
+    "Mapo tofu is a Sichuan dish of soft tofu simmered in a spicy, numbing sauce of chili bean paste and Sichuan peppercorn.",
+    {
+        "image": "https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/llama4_hgf.png",
+        "text": "Represent this image.",
+    },
+    {
+        "video": "https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/mapo_tofu.mp4",
+        "text": "Represent this video.",
+    },
+]
+
+query_embeddings = model.encode_query(queries)
+document_embeddings = model.encode_document(documents)
+print(query_embeddings.shape, document_embeddings.shape)
+# (2, 2048) (3, 2048)
+
+similarities = model.similarity(query_embeddings, document_embeddings)
+print(similarities)
+# tensor([[0.0683, 0.4972, 0.0309],
+#         [0.7829, 0.1428, 0.5492]])
 ```
 
+Each input is a string, a URL or path, a `PIL.Image`, or a dict combining `image`,
+`video`, and `text` keys. Put `image` or `video` before `text` so the prompt matches
+the ordering used above. Chat messages such as
+`{"role": "user", "content": [{"type": "image", "image": ...}, {"type": "text", "text": ...}]}`
+are also accepted, which is the way to interleave several images or videos in one input.
 
 ## Matryoshka Embeddings
 
@@ -109,6 +130,12 @@ embeddings = model.encode(
 embedding_256 = torch.nn.functional.normalize(embedding[..., :256], dim=-1)
 ```
 
+With Sentence Transformers, pass `truncate_dim` and let it renormalize:
+
+```python
+embeddings_256 = model.encode_document(documents, truncate_dim=256, normalize_embeddings=True)
+```
+
 Use a dimension listed in `model.config.matryoshka_dimensions`. On MMEB-v2, 256-dimensional embeddings retain 98.7% of the full-dimensional image and video performance.
 
 ## Serving

additional_chat_templates/sentence_transformers.jinja

@@ -0,0 +1,32 @@
+{#- Embedding template used by Sentence Transformers.
+
+    Renders the same surface form as chat_template.jinja for text and images. Video
+    is the exception: it renders as a bare <|video_pad|>, because the processor
+    expands that token into one <seconds><|vision_start|>...<|vision_end|> block per
+    frame and, from transformers 5.3, no longer consumes a surrounding pair. The
+    wrapped form therefore gains a second, unwanted pair of vision boundaries there.
+    This matches embedding_chat_template.jinja, which is bare for the same reason.
+
+    The tokenizer normalizer drops the role newline before a leading image and the
+    newline after the final <|im_end|>, then the post-processor appends <embedding>. #}
+{%- for message in messages %}
+{{- '<|im_start|>' + message.role + '\n' }}
+{%- if message.content is string %}
+{{- message.content }}
+{%- elif message.content is iterable and message.content is not mapping %}
+{%- for item in message.content %}
+{%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
+{{- '<|vision_start|><|image_pad|><|vision_end|>' }}
+{%- elif 'video' in item or item.type == 'video' %}
+{{- '<|video_pad|>' }}
+{%- elif 'text' in item %}
+{{- item.text }}
+{%- else %}
+{{- raise_exception('Unexpected item type in content.') }}
+{%- endif %}
+{%- endfor %}
+{%- elif message.content is not none %}
+{{- raise_exception('Unexpected content type.') }}
+{%- endif %}
+{{- '<|im_end|>\n' }}
+{%- endfor %}

config_sentence_transformers.json

@@ -0,0 +1,11 @@
+{
+  "__version__": {
+    "pytorch": "2.11.0+cu128",
+    "sentence_transformers": "5.7.0",
+    "transformers": "5.2.0"
+  },
+  "default_prompt_name": null,
+  "model_type": "SentenceTransformer",
+  "prompts": {},
+  "similarity_fn_name": "cosine"
+}

modeling_st_wemm.py

@@ -0,0 +1,70 @@
+"""Sentence Transformers module for WeMM-Embedding.
+
+Reproduces the `transformers` usage from the model card inside a Sentence Transformers
+pipeline: vision inputs are prepared with `qwen_vl_utils.process_vision_info` and the
+embedding is read from `WeMMEmbedding.embedding`, which pools the `<embedding>` position and
+L2-normalizes. Everything else (batching, prompts, truncation, `encode_query` /
+`encode_document`, similarity) comes from the stock `Transformer` module.
+"""
+
+from __future__ import annotations
+
+import inspect
+from typing import Any
+
+from sentence_transformers.models import Transformer
+
+
+class WeMMTransformer(Transformer):
+    """`Transformer` that prepares images and videos the way the model card's snippet does."""
+
+    def __init__(self, model_name_or_path: str, **kwargs: Any) -> None:
+        super().__init__(model_name_or_path, **kwargs)
+        vision_config = getattr(self.config, "vision_config", None)
+        self.image_patch_size = int(getattr(vision_config, "patch_size", 16))
+
+        # `embedding` hands its **kwargs to the inner model, so filtering on its own signature
+        # would drop `pixel_values`. Filter on the inner model's parameters instead, plus the
+        # processor's input names for anything the model only accepts as **kwargs.
+        inner_model = getattr(self.model, "model", self.model)
+        signature = set(inspect.signature(inner_model.forward).parameters)
+        signature |= set(getattr(self.processor, "model_input_names", ()))
+        for modality_params in self.modality_config.values():
+            method_name = modality_params["method"]
+            if method_name != "forward":
+                self._method_signature_cache.setdefault(method_name, signature)
+
+    def _apply_chat_template(
+        self,
+        messages: list[list[dict[str, Any]]],
+        modality_kwargs: dict[str, dict[str, Any]],
+        common_kwargs: dict[str, Any],
+        chat_template_kwargs: dict[str, Any],
+    ) -> dict[str, Any]:
+        """Render the chat template and prepare images / videos exactly as the model card does."""
+        from qwen_vl_utils import process_vision_info
+
+        chat_template_kwargs = {"add_generation_prompt": False, **chat_template_kwargs}
+        texts = [
+            self.processor.apply_chat_template(conversation, tokenize=False, **chat_template_kwargs)
+            for conversation in messages
+        ]
+        images, videos, video_kwargs = process_vision_info(
+            [list(conversation) for conversation in messages],
+            image_patch_size=self.image_patch_size,
+            return_video_kwargs=True,
+            return_video_metadata=True,
+        )
+        if videos is not None:
+            videos, video_metadata = (list(part) for part in zip(*videos))
+            video_kwargs = {**video_kwargs, "video_metadata": video_metadata}
+
+        return self.processor(
+            text=texts,
+            images=images,
+            videos=videos,
+            text_kwargs=modality_kwargs["text"],
+            images_kwargs=modality_kwargs["image"],
+            videos_kwargs={**modality_kwargs["video"], **video_kwargs},
+            common_kwargs=common_kwargs,
+        )

modeling_wemm_embedding.py

@@ -5,6 +5,10 @@ from transformers import Qwen3_5ForConditionalGeneration
 
 class WeMMEmbedding(Qwen3_5ForConditionalGeneration):
     def embedding(self, input_ids=None, attention_mask=None, **kwargs):
+        # transformers < 5.15 reuses the rope_deltas cached by the previous multimodal
+        # forward for a text-only one, which shifts its position ids.
+        self.model.rope_deltas = None
+
         outputs = self.model(
             input_ids=input_ids,
             attention_mask=attention_mask,

modules.json

@@ -0,0 +1,8 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "modeling_st_wemm.WeMMTransformer"
+  }
+]

sentence_bert_config.json

@@ -0,0 +1,37 @@
+{
+    "transformer_task": "feature-extraction",
+    "modality_config": {
+        "text": {
+            "method": "embedding",
+            "method_output_name": null
+        },
+        "image": {
+            "method": "embedding",
+            "method_output_name": null
+        },
+        "video": {
+            "method": "embedding",
+            "method_output_name": null
+        },
+        "image+text": {
+            "method": "embedding",
+            "method_output_name": null
+        },
+        "text+video": {
+            "method": "embedding",
+            "method_output_name": null
+        },
+        "message": {
+            "method": "embedding",
+            "method_output_name": null,
+            "format": "structured"
+        }
+    },
+    "module_output_name": "sentence_embedding",
+    "processing_kwargs": {
+        "chat_template": {
+            "chat_template": "sentence_transformers"
+        }
+    },
+    "unpad_inputs": false
+}

tokenizer_config.json

@@ -24,7 +24,7 @@
   "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": "Qwen3VLProcessor",
   "split_special_tokens": false,
-  "tokenizer_class": "Qwen2TokenizerFast",
+  "tokenizer_class": "PreTrainedTokenizerFast",
   "unk_token": null,
   "video_token": "<|video_pad|>",
   "vision_bos_token": "<|vision_start|>",

wemm_sentence_transformers.py

@@ -1,198 +0,0 @@
-#!/usr/bin/env python3
-"""Sentence-Transformers adapter for WeMM-Embedding."""
-
-from __future__ import annotations
-
-from pathlib import Path
-from typing import Any
-
-import torch
-from qwen_vl_utils import process_vision_info
-from sentence_transformers import SentenceTransformer
-from sentence_transformers.sentence_transformer.modules import InputModule
-from transformers import AutoModel, AutoProcessor
-
-
-def _as_messages(
-    value: Any,
-    prompt: str | None = None,
-) -> list[dict[str, Any]]:
-    if isinstance(value, str):
-        text = f"{prompt or ''}{value}"
-        return [{"role": "user", "content": [{"type": "text", "text": text}]}]
-
-    if isinstance(value, list):
-        if not value or not all(
-            isinstance(message, dict) and "role" in message for message in value
-        ):
-            raise TypeError("A list input must be one chat-style conversation.")
-        messages = [dict(message) for message in value]
-        if prompt:
-            messages.insert(0, {"role": "user", "content": prompt})
-        return messages
-
-    if not isinstance(value, dict):
-        raise TypeError(
-            "Input must be text, a chat conversation, or a multimodal dict."
-        )
-    if "messages" in value:
-        return _as_messages(value["messages"], prompt=prompt)
-
-    content: list[dict[str, Any]] = []
-    for modality in ("image", "video"):
-        if value.get(modality) is not None:
-            content.append({"type": modality, modality: value[modality]})
-    if value.get("text") is not None:
-        content.append(
-            {"type": "text", "text": f"{prompt or ''}{value['text']}"}
-        )
-    elif prompt:
-        content.append({"type": "text", "text": prompt})
-    if not content:
-        raise ValueError("Input must contain text, image, or video.")
-    return [{"role": "user", "content": content}]
-
-
-class WeMMInputModule(InputModule):
-    config_keys = ["model_name_or_path"]
-    save_in_root = False
-
-    def __init__(
-        self,
-        model_name_or_path: str,
-        *,
-        device: str | torch.device | None = None,
-        **_: Any,
-    ) -> None:
-        super().__init__()
-        self.model_name_or_path = str(model_name_or_path)
-        self.processor = AutoProcessor.from_pretrained(
-            self.model_name_or_path,
-            trust_remote_code=True,
-        )
-        self.tokenizer = self.processor.tokenizer
-        self.tokenizer.padding_side = "right"
-        self.auto_model = AutoModel.from_pretrained(
-            self.model_name_or_path,
-            trust_remote_code=True,
-            dtype=torch.bfloat16,
-            device_map=device,
-            attn_implementation="sdpa",
-            experts_implementation="eager",
-        ).eval()
-        self.embedding_token_id = self.tokenizer.convert_tokens_to_ids(
-            "<embedding>"
-        )
-        text_config = getattr(
-            self.auto_model.config,
-            "text_config",
-            self.auto_model.config,
-        )
-        self.embedding_dimension = int(text_config.hidden_size)
-
-    @property
-    def modalities(self):
-        return [
-            "text",
-            "image",
-            "video",
-            "message",
-            ("image", "text"),
-            ("text", "video"),
-        ]
-
-    @property
-    def max_seq_length(self) -> int:
-        return int(self.tokenizer.model_max_length)
-
-    def preprocess(
-        self,
-        inputs: list[Any],
-        prompt: str | None = None,
-        **_: Any,
-    ) -> dict[str, torch.Tensor | Any]:
-        conversations = [
-            _as_messages(value, prompt=prompt)
-            for value in inputs
-        ]
-        prompts = [
-            self.processor.apply_chat_template(
-                messages,
-                tokenize=False,
-                add_generation_prompt=False,
-            )
-            for messages in conversations
-        ]
-        images, videos, video_kwargs = process_vision_info(
-            conversations,
-            image_patch_size=16,
-            return_video_kwargs=True,
-            return_video_metadata=True,
-        )
-        if videos is not None:
-            videos, video_metadata = zip(*videos)
-            videos, video_metadata = list(videos), list(video_metadata)
-        else:
-            video_metadata = None
-
-        features = self.processor(
-            text=prompts,
-            images=images,
-            videos=videos,
-            video_metadata=video_metadata,
-            padding=True,
-            return_tensors="pt",
-            **video_kwargs,
-        )
-        input_ids = features["input_ids"]
-        positions = features["attention_mask"].sum(dim=1) - 1
-        batch = torch.arange(input_ids.shape[0])
-        terminal_ids = input_ids[batch, positions]
-        if not torch.all(terminal_ids == self.embedding_token_id):
-            raise RuntimeError("Each input must end with <embedding>.")
-        features["modality"] = "message"
-        return dict(features)
-
-    def forward(
-        self,
-        features: dict[str, torch.Tensor | Any],
-        **_: Any,
-    ) -> dict[str, torch.Tensor | Any]:
-        model_inputs = {
-            key: value
-            for key, value in features.items()
-            if key != "modality"
-        }
-        inner = getattr(self.auto_model, "model", None)
-        if inner is not None and hasattr(inner, "rope_deltas"):
-            inner.rope_deltas = None
-        features["sentence_embedding"] = self.auto_model.embedding(
-            **model_inputs
-        ).float()
-        return features
-
-    def get_embedding_dimension(self) -> int:
-        return self.embedding_dimension
-
-    def save(
-        self,
-        output_path: str,
-        *args: Any,
-        safe_serialization: bool = True,
-        **kwargs: Any,
-    ) -> None:
-        del args, safe_serialization, kwargs
-        self.save_config(output_path)
-
-
-def load_wemm_sentence_transformer(
-    model_name_or_path: str | Path,
-    *,
-    device: str | torch.device = "cuda:0",
-) -> SentenceTransformer:
-    return SentenceTransformer(
-        modules=[WeMMInputModule(str(model_name_or_path), device=device)],
-        device=str(device),
-        similarity_fn_name="cosine",
-    )
-