import base64

from typing import Literal, Union

from openai import OpenAI
from openai._types import NOT_GIVEN, NotGiven
from openai.types.chat import ChatCompletionMessageParam
from openai.types.create_embedding_response import CreateEmbeddingResponse

import torch

def file_base64(input_file: str) -> str:
    
    with open(input_file, "rb") as f:
        encoded_string = base64.b64encode(f.read()).decode('utf-8')
    return encoded_string

def create_chat_embeddings(
    client: OpenAI,
    *,
    messages: list[ChatCompletionMessageParam],
    model: str,
    encoding_format: Union[Literal["base64", "float"], NotGiven] = NOT_GIVEN,
) -> CreateEmbeddingResponse:
    return client.post(
        "/embeddings",
        cast_to=CreateEmbeddingResponse,
        body={"messages": messages, "model": model, "encoding_format": encoding_format},
    )

def message2embedding(
    client: OpenAI,
    *,
    message: ChatCompletionMessageParam,
    model: str,
    encoding_format: Union[Literal["base64", "float"], NotGiven] = NOT_GIVEN,
) -> list[float]:
    response = create_chat_embeddings(
        client,
        messages=message,
        model=model,
        encoding_format=encoding_format,
    )
    return torch.tensor(response.data[0].embedding)

client = OpenAI(base_url="http://localhost:8000/v1", api_key="dummy_api_key")

🌍 Multilingual Text Retrieval

python

example_query = 'Query: 氧气在空气中占比多少？'
example_text_1 = "空气是指大气层中由不同气体和各类飘浮在其中的固体与液体颗粒（大气颗粒与气溶胶）所组成的气态混合物。地球大气层的空气主要由78.1%的氮气、20.9%氧气、0.9%的氩气和1~4%的水蒸气组成，其成分并不是固定的，随着高度、气压、温度的改变和对流情况不同，局部空气的组成比例也会改变。空气在大气层（特别是对流层）中的流动形成了风和曳流、气旋、龙卷等自然现象，而空气中飘浮的颗粒则形成了云、雾、霾和沙尘暴等短期天气情况。空气在海洋和陆地之间跨区域流动所承载的湿度和热能传导也是水循环和气候变率与变化的关键一环。"
example_text_2 = "水（化学式：H2O）是一种无机化合物，在常温且无杂质中是无色[1]无味不导电的透明液体，也会通过蒸发产生气态的水蒸气（这种蒸发可以发生在任何温度下，同时取决于与空气接触的表面积和湿度差）。在标准大气压下，水的凝固点是0 °C（32 °F；273 K），沸点是100 °C（212 °F；373 K）。"
query = [{'role': 'user', 'content': [{'type': 'text', 'text': example_query}]}]
text_1 = [{'role': 'user', 'content': [{'type': 'text', 'text': example_text_1}]}]
text_2 = [{'role': 'user', 'content': [{'type': 'text', 'text': example_text_2}]}]

query_embedding = message2embedding(client, message=query, model="nkwbtb/OmniEmbed-v0.1")
text_1_embedding = message2embedding(client, message=text_1, model="nkwbtb/OmniEmbed-v0.1")
text_2_embedding = message2embedding(client, message=text_2, model="nkwbtb/OmniEmbed-v0.1")

sim1 = torch.nn.functional.cosine_similarity(query_embedding.unsqueeze(0), text_1_embedding.unsqueeze(0))
sim2 = torch.nn.functional.cosine_similarity(query_embedding.unsqueeze(0), text_2_embedding.unsqueeze(0))

print("Similarities:", sim1.item(), sim2.item())

📈 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"
example_image_2 = "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/llama4_hgf.png"
query = [{'role': 'user', 'content': [{'type': 'text', 'text': example_query}]}]
image_1 = [{'role': 'user', 'content': [{'type': 'image_url', 'image_url': {'url': example_image_1}}]}]
image_2 = [{'role': 'user', 'content': [{'type': 'image_url', 'image_url': {'url': example_image_2}}]}]

query_embedding = message2embedding(client, message=query, model="nkwbtb/OmniEmbed-v0.1")
image_1_embedding = message2embedding(client, message=image_1, model="nkwbtb/OmniEmbed-v0.1")
image_2_embedding = message2embedding(client, message=image_2, model="nkwbtb/OmniEmbed-v0.1")

sim1 = torch.cosine_similarity(query_embedding.unsqueeze(0), image_1_embedding.unsqueeze(0))
sim2 = torch.cosine_similarity(query_embedding.unsqueeze(0), image_2_embedding.unsqueeze(0))

print("Similarities:", sim1.item(), sim2.item())

python

# Or you can use base64 encoded image
example_image_1 = file_base64("qwen2.5omni_hgf.jpg")
example_image_2 = file_base64("llama4_hgf.jpg")
image_1 = [{'role': 'user', 'content': [{'type': 'image_url', 'image_url': {'url': f"data:image/jpeg;base64,{example_image_1}"}}]}]
image_2 = [{'role': 'user', 'content': [{'type': 'image_url', 'image_url': {'url': f"data:image/jpeg;base64,{example_image_2}"}}]}]

🎵 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"
example_audio_2 = "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/jay_chou_superman_cant_fly.mp3"
query = [{'role': 'user', 'content': [{'type': 'text', 'text': example_query}]}]
audio_1 = [{'role': 'user', 'content': [{'type': 'audio_url', 'audio_url': {'url': example_audio_1}}]}]
audio_2 = [{'role': 'user', 'content': [{'type': 'audio_url', 'audio_url': {'url': example_audio_2}}]}]

query_embedding = message2embedding(client, message=query, model="nkwbtb/OmniEmbed-v0.1")
audio_1_embedding = message2embedding(client, message=audio_1, model="nkwbtb/OmniEmbed-v0.1")
audio_2_embedding = message2embedding(client, message=audio_2, model="nkwbtb/OmniEmbed-v0.1")

sim1 = torch.cosine_similarity(query_embedding.unsqueeze(0), audio_1_embedding.unsqueeze(0))
sim2 = torch.cosine_similarity(query_embedding.unsqueeze(0), audio_2_embedding.unsqueeze(0))

print("Similarities:", sim1.item(), sim2.item())

python

# Or you can use base64 encoded audio
audio_1_base64 = file_base64("joe_hisaishi_summer.mp3")
audio_2_base64 = file_base64("jay_chou_superman_cant_fly.mp3")
audio_1 = [{'role': 'user', 'content': [{'type': 'audio_url', 'audio_url': {'url': f"data:audio/mp3;base64,{audio_1_base64}"}}]}]
audio_2 = [{'role': 'user', 'content': [{'type': 'audio_url', 'audio_url': {'url': f"data:audio/mp3;base64,{audio_2_base64}"}}]}]

🎬 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"
example_video_2 = "https://huggingface.co/Tevatron/OmniEmbed-v0.1/resolve/main/assets/zhajiang_noodle.mp4"
query = [{'role': 'user', 'content': [{'type': 'text', 'text': example_query}]}]
video_1 = [{'role': 'user', 'content': [{'type': 'video_url', 'video_url': {'url': example_video_1}}]}]
video_2 = [{'role': 'user', 'content': [{'type': 'video_url', 'video_url': {'url': example_video_2}}]}]

query_embedding = message2embedding(client, message=query, model="nkwbtb/OmniEmbed-v0.1")
video_1_embedding = message2embedding(client, message=video_1, model="nkwbtb/OmniEmbed-v0.1")
video_2_embedding = message2embedding(client, message=video_2, model="nkwbtb/OmniEmbed-v0.1")

sim1 = torch.cosine_similarity(query_embedding.unsqueeze(0), video_1_embedding.unsqueeze(0))
sim2 = torch.cosine_similarity(query_embedding.unsqueeze(0), video_2_embedding.unsqueeze(0))

print("Similarities:", sim1.item(), sim2.item())

python

# Or you can use base64 encoded video
video_1_base64 = file_base64("mapo_tofu.mp4")
video_2_base64 = file_base64("zhajiang_noodle.mp4")
video_1 = [{'role': 'user', 'content': [{'type': 'video_url', 'video_url': {'url': f"data:video/mp4;base64,{video_1_base64}"}}]}]
video_2 = [{'role': 'user', 'content': [{'type': 'video_url', 'video_url': {'url': f"data:video/mp4;base64,{video_2_base64}"}}]}]

nkwbtb/OmniEmbed-v0.1

Merged model for Tevatron/OmniEmbed-v0.1

vLLM hosting

Usage

🌍 Multilingual Text Retrieval

📈 Image Document Retrieval (Image, Chart, PDF)

🎵 Audio Retrieval

🎬 Video Retrieval

References

nkwbtb/
OmniEmbed-v0.1