tencent

/

WeMM-Embedding-2B

like 17

Follow

Tencent 11.7k

Hello @JUNJIE99 @kekekeke @xiaoxiaoshadiao and team!

Congratulations on this release! As far as I can, this is beating just about any model at any benchmark at the same size, very nice to see! And thanks for shipping wemm_sentence_transformers.py alongside the model. This PR builds on it further so Sentence Transformers can load WeMM-Embedding straight from the repository id, and it picks up two upstream transformers changes that silently alter the inputs this model sees.

As a heads up, this PR is AI-generated but human-reviewed.

Pull Request overview

• Integrate WeMM-Embedding-2B with Sentence Transformers v5.7+ via modules.json and trust_remote_code
• Restore the training-time tokenization on transformers newer than the pinned 5.2.0

Details

The pipeline is one module, WeMMTransformer(feature-extraction) -> sentence_embedding, with no
Pooling or Normalize behind it: WeMMEmbedding.embedding already reads the <embedding> position
and L2-normalizes, so its return value goes straight into sentence_embedding. modality_config maps
text, image, video, image+text, text+video and message to that method, so
model.encode("..."), model.encode({"image": ..., "text": ...}) and raw chat messages all work
through one entry point. config_sentence_transformers.json sets cosine similarity and declares no
prompts, since the instruction belongs in the input text rather than in a fixed prefix.

modeling_st_wemm.py is a Transformer subclass, so batching, prompts, truncation, encode_query /
encode_document, similarity and saving stay stock. It overrides only _apply_chat_template,
building the processor inputs through qwen_vl_utils.process_vision_info (with image_patch_size
read from vision_config.patch_size) exactly as the model card's snippet does. That matters mostly
for video, where the processor's own frame sampling picks 114 frames against process_vision_info's
112. Its constructor also seeds the method-signature cache from the inner model's parameters, since
embedding takes multimodal inputs as **kwargs and would otherwise have pixel_values filtered out.

Two things had to change for the model to see its trained input format on current transformers.
Both were measured to be no-ops on the pinned 5.2.0.

1. tokenizer_class is now PreTrainedTokenizerFast. From transformers 5.3-ish, TokenizersBackend
   rebuilds Qwen2 tokenizers from vocab / merges and Qwen2Tokenizer.__init__ hardcodes
   normalizers.NFC(), dropping this repository's Replace(Regex("\n$"), "") (the normalizer the
   embedding_chat_template.jinja comment describes). Without it every input gains a \n before
   <embedding>, and image inputs also gain the role newline the model was trained without.
   PreTrainedTokenizerFast takes the loader branch that reads tokenizer.json verbatim; on 5.2.0 the
   two classes tokenize identically (10/10 sample strings, with AutoProcessor still resolving
   Qwen3VLProcessor with the same special-token ids).
2. additional_chat_templates/sentence_transformers.jinja renders video as a bare <|video_pad|>
   rather than <|vision_start|><|video_pad|><|vision_end|>. 5.2.0 replaced the whole wrapped block
   with the per-frame <seconds><|vision_start|>...<|vision_end|> expansion; newer versions replace
   only <|video_pad|> and leave the outer pair in place, adding a second, unwanted pair of vision
   boundaries. The bare form gives byte-identical ids on 5.2.0 and correct ids on newer versions, which
   is also why embedding_chat_template.jinja is bare. chat_template.jinja is untouched, since it
   matches the upstream Qwen3-VL convention and is what SGLang reads.

modeling_wemm_embedding.py gains one line: embedding resets self.model.rope_deltas, the same
guard wemm_sentence_transformers.py already carried. Below transformers 5.15 a text-only forward
reuses the deltas cached by the previous multimodal forward, shifting its position ids (max |diff|
0.0015 bf16, on text encoded right after an image); newer versions guard against it. Doing it in
embedding rather than in the caller covers the model card's transformers path too.

Verification, in fp32 over text, image, image+text and video inputs: on transformers 5.2.0 the
Sentence Transformers path reproduces AutoModel + AutoProcessor + process_vision_info
bit-exactly, max |diff| 0.0 on every embedding and on the full similarity matrix. Against that same
5.2.0 reference, Sentence Transformers on 5.15.0 produces bit-identical input_ids, attention_mask,
mm_token_type_ids, pixel_values, pixel_values_videos and grids, with embeddings within 1.4e-06
and similarities within 8.9e-07 (GPU kernel differences between the two builds). Checked on
sentence-transformers 5.7.0 and 6.0.0. Dict inputs, bare URLs, PIL.Image objects, mixed-modality
batches, encode_query / encode_document and truncate_dim were each compared against the
equivalent chat message form.

wemm_sentence_transformers.py is removed, since modules.json supersedes it and also works from a
repository id. Happy to keep it if you would rather not drop it, though it renders video through
chat_template.jinja and so hits the doubled-boundary issue above on current transformers.

The model card's Sentence Transformers section is rewritten around the snippet below, plus a
sentence-transformers tag, a truncate_dim note under Matryoshka Embeddings, and
sentence-transformers>=5.7.0 in the install line.

from sentence_transformers import SentenceTransformer

model_id = "tencent/WeMM-Embedding-2B"
model = SentenceTransformer(model_id, trust_remote_code=True)

queries = [
    "Which Llama 4 model variants are available?",
    "How is mapo tofu prepared?",
]
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]])

To try it before merging, load the PR revision:
SentenceTransformer(model_id, revision="refs/pr/1", trust_remote_code=True).

Nothing else changes: the transformers, vLLM and SGLang paths are unaffected, and the rest is
purely additive, apart from the wemm_sentence_transformers.py removal.

Happy to tweak anything you'd like changed. Please let me know if you have any questions or feedback!

• Tom Aarsen