Integrate with Sentence Transformers

#1
by tomaarsen HF Staff - opened

Hello @gowitheflow ,

Congratulations on the May 2026 update, the gains across all four modalities look great! This PR is very similar to what I did for https://huggingface.co/LCO-Embedding/LCO-Embedding-Omni-3B, and should make the model a bit easier to use.

Details

I copied over these files from the 3B model:

  • modules.json
  • sentence_bert_config.json
  • config_sentence_transformers.json
  • 1_Pooling/config.json
  • chat_template.jinja (NOTE: this is a byte-identical replacement for the legacy chat_template.json, which is removed)
  • additional_chat_templates/sentence_transformers.jinja

I also modified the README to add a Using Sentence Transformers section with text, image, audio, video, and multimodal examples, plus the sentence-transformers tag, library_name, and pipeline_tag in the frontmatter. There shouldn't be any changes to the model itself, and the transformers usage documented in the README continues to work unchanged.

import torch
from sentence_transformers import SentenceTransformer

model = SentenceTransformer(
    "LCO-Embedding/LCO-Embedding-Omni-3B-2605",
    model_kwargs={
        "dtype": torch.bfloat16,
        # "attn_implementation": "flash_attention_2",  # recommended, if a flash-attn build exists for your platform
    },
)

query = "How many input modalities does Qwen2.5-Omni support?"
documents = [
    "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/qwen2.5omni_hgf.png",
    "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/llama4_hgf.png",
]

query_embedding = model.encode(query)
document_embeddings = model.encode(documents, batch_size=1)
print(model.similarity(query_embedding, document_embeddings))
# tensor([[0.6544, 0.3852]])

The expected outputs in the README were produced in bfloat16 on CUDA with the default sdpa attention, so they are reproducible from the snippet exactly as written.

Please let me know if you have any questions or feedback!

  • Tom Aarsen
tomaarsen changed pull request status to open
LCO-Embedding org

amazing, thanks again!

gowitheflow changed pull request status to merged

Sign up or log in to comment