Integrate with Sentence Transformers

1_Pooling/config.json

@@ -0,0 +1,5 @@
+{
+    "embedding_dimension": 2048,
+    "pooling_mode": "lasttoken",
+    "include_prompt": true
+}

README.md

@@ -1,8 +1,11 @@
 ---
 license: apache-2.0
+pipeline_tag: feature-extraction
+library_name: sentence-transformers
 tags:
 - multimodal-embedding
 - transformers
+- sentence-transformers
 - feature-extraction
 ---
 # May 2026 update of LCO-Embedding models
@@ -28,6 +31,115 @@ In this version, we make substantial improvements on all 4 modalities (text, ima
 
 # Usage
 
+## Using Sentence Transformers
+
+Install Sentence Transformers with the multimodal extras (for image, audio, and video support):
+
+```bash
+pip install "sentence_transformers[image,audio,video]" "transformers>=5.6.0"
+```
+
+```python
+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
+    },
+)
+```
+
+The same `Summarize the above <modality> in one word:` instruction used in the paper is baked into the chat template, so `encode()` takes plain text, file paths, URLs, or multimodal dicts directly.
+
+### Text Retrieval
+```python
+query = "What is the tallest mountain in the world?"
+documents = [
+    "Mount Everest is Earth's highest mountain above sea level, located in the Mahalangur Himal sub-range of the Himalayas. Its elevation of 8,848.86 metres was established by a joint Chinese-Nepali survey in 2020.",
+    "K2, at 8,611 metres above sea level, is the second-highest mountain on Earth, after Mount Everest. It lies in the Karakoram range on the China-Pakistan border.",
+    "Mount Kilimanjaro is a dormant volcano in Tanzania. It is the highest mountain in Africa, with its summit about 5,895 metres above sea level.",
+]
+
+query_embedding = model.encode(query)
+document_embeddings = model.encode(documents)
+print(model.similarity(query_embedding, document_embeddings))
+# tensor([[0.5368, 0.5053, 0.4989]])
+```
+
+### Image Retrieval
+```python
+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]])
+```
+
+### Audio Retrieval
+```python
+query = "A light piano piece"
+documents = [
+    "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/joe_hisaishi_summer.mp3",
+    "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/jay_chou_superman_cant_fly.mp3",
+]
+
+query_embedding = model.encode(query)
+document_embeddings = model.encode(documents, batch_size=1)
+print(model.similarity(query_embedding, document_embeddings))
+# tensor([[0.3649, 0.0662]])
+```
+
+### Video Retrieval
+```python
+# For video on smaller GPUs, cap the processor up front:
+model[0].processing_kwargs.update({
+    "video": {"max_pixels": 64 * 28 * 28, "do_sample_frames": True, "fps": 1},
+})
+
+query = "How to cook Mapo Tofu?"
+documents = [
+    "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/mapo_tofu.mp4",
+    "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/zhajiang_noodle.mp4",
+]
+
+query_embedding = model.encode(query)
+document_embeddings = model.encode(documents, batch_size=1)
+print(model.similarity(query_embedding, document_embeddings))
+# tensor([[0.6408, 0.4967]])
+```
+
+### Multimodal Inputs
+
+To embed a document that combines multiple modalities, pass a dict with any combination of `"text"`, `"image"`, `"audio"`, and `"video"` keys instead of a single path or string:
+
+```python
+documents = [
+    {
+        "text": "A cooking tutorial for Mapo Tofu",
+        "video": "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/mapo_tofu.mp4",
+    },
+    {
+        "image": "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/qwen2.5omni_hgf.png",
+        "audio": "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/joe_hisaishi_summer.mp3",
+    },
+]
+document_embeddings = model.encode(documents, batch_size=1)
+print(document_embeddings.shape)
+# (2, 2048)
+```
+
+The expected outputs above were produced in `bfloat16` on a CUDA device with the default (`sdpa`) attention. Exact values shift slightly in the fourth decimal with a different dtype or attention implementation.
+
+## Using Transformers
+
 All inference code is the same with our [OG models](https://huggingface.co/LCO-Embedding/LCO-Embedding-Omni-7B) and can seamlessly support the new checkpoint by changing the model name.
 
 # Contributors

additional_chat_templates/sentence_transformers.jinja

@@ -0,0 +1,51 @@
+{%- set audio_count = namespace(value=0) -%}
+{%- set image_count = namespace(value=0) -%}
+{%- set video_count = namespace(value=0) -%}
+<|im_start|>system
+You are a helpful assistant.<|im_end|>
+{% for message in messages -%}
+    {%- if message['role'] == 'system' -%}
+        {#- skip: the fixed system prompt above was already emitted. Any input system
+            message is only present to silence Qwen2.5-Omni's default-prompt warning. -#}
+    {%- else -%}
+        <|im_start|>{{ message['role'] }}
+{% if message['content'] is string -%}
+        {{- message['content'] -}}<|im_end|>
+{% else -%}
+        {%- set seen = namespace(image=false, audio=false, video=false) -%}
+        {%- for content in message['content'] -%}
+            {%- if content['type'] == 'image' or 'image' in content or 'image_url' in content -%}
+                {%- set image_count.value = image_count.value + 1 -%}
+                {%- set seen.image = true -%}
+                {%- if add_vision_id -%}Picture {{ image_count.value }}: {% endif -%}
+                <|vision_bos|><|IMAGE|><|vision_eos|>
+            {%- elif content['type'] == 'audio' or 'audio' in content or 'audio_url' in content -%}
+                {%- set audio_count.value = audio_count.value + 1 -%}
+                {%- set seen.audio = true -%}
+                {%- if add_audio_id -%}Audio {{ audio_count.value }}: {% endif -%}
+                <|audio_bos|><|AUDIO|><|audio_eos|>
+            {%- elif content['type'] == 'video' or 'video' in content -%}
+                {%- set video_count.value = video_count.value + 1 -%}
+                {%- set seen.video = true -%}
+                {%- if add_vision_id -%}Video {{ video_count.value }}: {% endif -%}
+                <|vision_bos|><|VIDEO|><|vision_eos|>
+            {%- elif 'text' in content -%}
+                {{- content['text'] -}}
+            {%- endif -%}
+        {%- endfor -%}
+        {%- if seen.image -%}
+{{ '\n' }}Summarize the above image in one word:
+        {%- elif seen.video -%}
+{{ '\n' }}Summarize the above video in one word:
+        {%- elif seen.audio -%}
+{{ '\n' }}Summarize the above audio in one word:
+        {%- else -%}
+{{ '\n' }}Summarize the above text in one word:
+        {%- endif -%}
+        <|im_end|>
+{% endif -%}
+    {%- endif -%}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+    <|im_start|>assistant
+{% endif -%}

chat_template.jinja

@@ -0,0 +1,7 @@
+{% set audio_count = namespace(value=0) %}{% set image_count = namespace(value=0) %}{% set video_count = namespace(value=0) %}{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|im_start|>system
+You are a helpful assistant.<|im_end|>
+{% endif %}<|im_start|>{{ message['role'] }}
+{% if message['content'] is string %}{{ message['content'] }}<|im_end|>
+{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_bos|><|IMAGE|><|vision_eos|>{% elif content['type'] == 'audio' or 'audio' in content or 'audio_url' in content %}{% set audio_count.value = audio_count.value + 1 %}{% if add_audio_id %}Audio {{ audio_count.value }}: {% endif %}<|audio_bos|><|AUDIO|><|audio_eos|>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_bos|><|VIDEO|><|vision_eos|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>
+{% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant
+{% endif %}

chat_template.json

@@ -1,3 +0,0 @@
-{
-  "chat_template": "{% set audio_count = namespace(value=0) %}{% set image_count = namespace(value=0) %}{% set video_count = namespace(value=0) %}{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n{% endif %}<|im_start|>{{ message['role'] }}\n{% if message['content'] is string %}{{ message['content'] }}<|im_end|>\n{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_bos|><|IMAGE|><|vision_eos|>{% elif content['type'] == 'audio' or 'audio' in content or 'audio_url' in content %}{% set audio_count.value = audio_count.value + 1 %}{% if add_audio_id %}Audio {{ audio_count.value }}: {% endif %}<|audio_bos|><|AUDIO|><|audio_eos|>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_bos|><|VIDEO|><|vision_eos|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>\n{% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}"
-}

config_sentence_transformers.json

@@ -0,0 +1,13 @@
+{
+  "__version__": {
+    "pytorch": "2.10.0+cu128",
+    "sentence_transformers": "5.4.0",
+    "transformers": "5.6.0"
+  },
+  "default_prompt_name": "default",
+  "model_type": "SentenceTransformer",
+  "prompts": {
+    "default": "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech."
+  },
+  "similarity_fn_name": "cosine"
+}

modules.json

@@ -0,0 +1,20 @@
+[
+  {
+    "idx": 0,
+    "name": "0",
+    "path": "",
+    "type": "sentence_transformers.base.modules.transformer.Transformer"
+  },
+  {
+    "idx": 1,
+    "name": "1",
+    "path": "1_Pooling",
+    "type": "sentence_transformers.sentence_transformer.modules.pooling.Pooling"
+  },
+  {
+    "idx": 2,
+    "name": "2",
+    "path": "2_Normalize",
+    "type": "sentence_transformers.sentence_transformer.modules.normalize.Normalize"
+  }
+]

sentence_bert_config.json

@@ -0,0 +1,48 @@
+{
+    "transformer_task": "any-to-any",
+    "modality_config": {
+        "text": {
+            "method": "forward",
+            "method_output_name": [
+                "hidden_states",
+                -1
+            ]
+        },
+        "image": {
+            "method": "forward",
+            "method_output_name": [
+                "hidden_states",
+                -1
+            ]
+        },
+        "audio": {
+            "method": "forward",
+            "method_output_name": [
+                "hidden_states",
+                -1
+            ]
+        },
+        "video": {
+            "method": "forward",
+            "method_output_name": [
+                "hidden_states",
+                -1
+            ]
+        },
+        "message": {
+            "method": "forward",
+            "method_output_name": [
+                "hidden_states",
+                -1
+            ],
+            "format": "structured"
+        }
+    },
+    "module_output_name": "token_embeddings",
+    "processing_kwargs": {
+        "chat_template": {
+            "chat_template": "sentence_transformers",
+            "add_generation_prompt": true
+        }
+    }
+}