Integrate with Sentence Transformers v5.4

1_Pooling/config.json

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

README.md

@@ -11,7 +11,12 @@ base_model:
 - Qwen/Qwen2.5-Omni-7B
 - Tevatron/Qwen2.5-Omni-7B-Thinker
 pipeline_tag: visual-document-retrieval
-library_name: peft
+library_name: sentence-transformers
+tags:
+- sentence-transformers
+- peft
+- multimodal
+- feature-extraction
 ---
 # Tevatron/OmniEmbed-v0.1
 
@@ -37,7 +42,110 @@ OmniEmbed achieves strong performance, comparable to models specifically optimiz
 
 ---
 
-### Usage
+## 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" "peft>=0.19.0"
+```
+
+```python
+import torch
+from sentence_transformers import SentenceTransformer
+
+model = SentenceTransformer(
+    "Tevatron/OmniEmbed-v0.1",
+    model_kwargs={
+        "torch_dtype": torch.bfloat16,
+        "attn_implementation": "flash_attention_2",  # pip install kernels; recommended but not mandatory
+    },
+)
+```
+
+#### 🎬 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},
+})
+
+example_query = "How to cook Mapo Tofu?"
+example_videos = [
+    "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(example_query)
+document_embeddings = model.encode_document(example_videos, batch_size=1)
+print(model.similarity(query_embedding, document_embeddings))
+# tensor([[0.4318, 0.2858]])
+```
+
+#### 🎵 Audio Retrieval
+```python
+example_query = "A light piano piece"
+example_audios = [
+    "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(example_query)
+document_embeddings = model.encode_document(example_audios, batch_size=1)
+print(model.similarity(query_embedding, document_embeddings))
+# tensor([[0.2734, 0.2165]])
+```
+
+#### 📈 Image Document Retrieval (Image, Chart, PDF)
+```python
+example_query = "How many input modality does Qwen2.5-Omni support?"
+example_images = [
+    "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(example_query)
+document_embeddings = model.encode_document(example_images, batch_size=1)
+print(model.similarity(query_embedding, document_embeddings))
+# tensor([[0.4682, 0.2956]])
+```
+
+#### 🌍 Multilingual Text Retrieval
+```python
+example_query = "氧气在空气中占比多少？"
+example_texts = [
+    "空气是指大气层中由不同气体和各类飘浮在其中的固体与液体颗粒（大气颗粒与气溶胶）所组成的气态混合物。地球大气层的空气主要由78.1%的氮气、20.9%氧气、0.9%的氩气和1~4%的水蒸气组成，其成分并不是固定的，随着高度、气压、温度的改变和对流情况不同，局部空气的组成比例也会改变。空气在大气层（特别是对流层）中的流动形成了风和曳流、气旋、龙卷等自然现象，而空气中飘浮的颗粒则形成了云、雾、霾和沙尘暴等短期天气情况。空气在海洋和陆地之间跨区域流动所承载的湿度和热能传导也是水循环和气候变率与变化的关键一环。",
+    "水（化学式：H2O）是一种无机化合物，在常温且无杂质中是无色[1]无味不导电的透明液体，也会通过蒸发产生气态的水蒸气（这种蒸发可以发生在任何温度下，同时取决于与空气接触的表面积和湿度差）。在标准大气压下，水的凝固点是0 °C（32 °F；273 K），沸点是100 °C（212 °F；373 K）。",
+]
+
+query_embedding = model.encode_query(example_query)
+document_embeddings = model.encode_document(example_texts, batch_size=1)
+print(model.similarity(query_embedding, document_embeddings))
+# tensor([[0.3371, 0.2496]])
+```
+
+#### 🧩 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_document(documents, batch_size=1)
+```
+
+### Using Transformers
+
 ```python
 # Import Library, Load Model and Processor
 import torch
@@ -82,7 +190,7 @@ def encode_message(message):
     return reps
 ```
 
-### 🎬 Video Retrieval
+#### 🎬 Video Retrieval
 ```python
 example_query = 'Query: How to cook Mapo Tofu?'
 example_video_1 = "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/mapo_tofu.mp4"
@@ -95,9 +203,12 @@ sim1 = torch.cosine_similarity(encode_message(query), encode_message(video_1))
 sim2 = torch.cosine_similarity(encode_message(query), encode_message(video_2))
 
 print("Similarities:", sim1.item(), sim2.item())
+# Video similarities: 0.408203125 0.283203125
+# Defaults produce ~30 GB of video activations; on smaller GPUs add e.g.
+# `"max_pixels": 64*28*28, "fps": 1, "min_pixels": 16*28*28` to each video dict.
 ```
 
-### 🎵 Audio Retrieval
+#### 🎵 Audio Retrieval
 ```python
 example_query = 'Query: A light piano piece'
 example_audio_1 = "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/joe_hisaishi_summer.mp3"
@@ -110,9 +221,10 @@ sim1 = torch.cosine_similarity(encode_message(query), encode_message(audio_1))
 sim2 = torch.cosine_similarity(encode_message(query), encode_message(audio_2))
 
 print("Similarities:", sim1.item(), sim2.item())
+# Audio similarities: 0.275390625 0.2353515625
 ```
 
-### 📈 Image Document Retrieval (Image, Chart, PDF)
+#### 📈 Image Document Retrieval (Image, Chart, PDF)
 ```python
 example_query = 'Query: How many input modality does Qwen2.5-Omni support?'
 example_image_1 = "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/qwen2.5omni_hgf.png"
@@ -125,9 +237,10 @@ sim1 = torch.cosine_similarity(encode_message(query), encode_message(image_1))
 sim2 = torch.cosine_similarity(encode_message(query), encode_message(image_2))
 
 print("Similarities:", sim1.item(), sim2.item())
+# Image similarities: 0.458984375 0.296875
 ```
 
-### 🌍 Multilingual Text Retrieval
+#### 🌍 Multilingual Text Retrieval
 ```python
 example_query = 'Query: 氧气在空气中占比多少？'
 example_text_1 = "空气是指大气层中由不同气体和各类飘浮在其中的固体与液体颗粒（大气颗粒与气溶胶）所组成的气态混合物。地球大气层的空气主要由78.1%的氮气、20.9%氧气、0.9%的氩气和1~4%的水蒸气组成，其成分并不是固定的，随着高度、气压、温度的改变和对流情况不同，局部空气的组成比例也会改变。空气在大气层（特别是对流层）中的流动形成了风和曳流、气旋、龙卷等自然现象，而空气中飘浮的颗粒则形成了云、雾、霾和沙尘暴等短期天气情况。空气在海洋和陆地之间跨区域流动所承载的湿度和热能传导也是水循环和气候变率与变化的关键一环。"
@@ -140,6 +253,7 @@ sim1 = torch.cosine_similarity(encode_message(query), encode_message(text_1))
 sim2 = torch.cosine_similarity(encode_message(query), encode_message(text_2))
 
 print("Similarities:", sim1.item(), sim2.item())
+# Text similarities: 0.3359375 0.2490234375
 ```
 
 ## Data & Training

additional_chat_templates/sentence_transformers.jinja

@@ -0,0 +1,6 @@
+<|im_start|>system
+You are a helpful assistant.<|im_end|>
+<|im_start|>user
+{% for message in messages %}{% if message['content'] is string %}{{ message['content'] }}{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}<|vision_bos|><|IMAGE|><|vision_eos|>{% elif content['type'] == 'audio' or 'audio' in content or 'audio_url' in content %}<|audio_bos|><|AUDIO|><|audio_eos|>{% elif content['type'] == 'video' or 'video' in content %}<|vision_bos|><|VIDEO|><|vision_eos|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}{% endif %}{% endfor %}<|im_end|>
+<|im_start|>assistant
+<|endoftext|>

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,14 @@
+{
+  "__version__": {
+    "sentence_transformers": "5.4.0",
+    "transformers": "5.6.0",
+    "pytorch": "2.10.0+cu128"
+  },
+  "prompts": {
+    "query": "Query: ",
+    "document": ""
+  },
+  "default_prompt_name": null,
+  "similarity_fn_name": "cosine",
+  "model_type": "SentenceTransformer"
+}

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,47 @@
+{
+    "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"
+        }
+    }
+}