开源协议: mit
数据集:

• Tevatron/bge-ir
• Tevatron/wiki-ss-nq-new
• Tevatron/pixmo-docs
• Tevatron/colpali
• Tevatron/msrvtt
• Tevatron/audiocaps
  基础模型:
• Qwen/Qwen2.5-Omni-7B
• Tevatron/Qwen2.5-Omni-7B-Thinker
  任务类型: 视觉文档检索
  库名称: peft

Tevatron/OmniEmbed-v0.1

OmniEmbed是基于Qwen2.5-Omni-7B构建的多模态嵌入模型，采用我们开发的Tevatron工具包训练——这是一个跨规模、跨语言、跨模态的统一文档检索工具包。
OmniEmbed能生成跨多语言文本、图像、音频和视频的统一嵌入表示，为多样化应用提供高效的跨模态检索能力。

📝 文本   🖼️ 图像   🎧 音频   🎥 视频   🌐 多语言

评估结果:

• 虽然AudioCaps存在多个基线模型，但文献中缺乏统一的查询/语料设置标准，难以直接对比。

OmniEmbed展现出强劲性能，与专为单项任务优化的模型相比毫不逊色。

使用指南

# 导入库、加载模型和处理器
import torch
from transformers import AutoProcessor, Qwen2_5OmniThinkerForConditionalGeneration
from qwen_omni_utils import process_mm_info

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-Omni-7B")
model = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained(
    'ArvinZhuang/OmniEmbed-test',
    attn_implementation="flash_attention_2",
    torch_dtype=torch.bfloat16
).to(device).eval()

processor.tokenizer.padding_side = "left"
model.padding_side = "left"

# 信息编码函数
def encode_message(message):
    texts = processor.apply_chat_template(message, tokenize=False, add_generation_prompt=True)[0] + "<|endoftext|>"
    audio_inputs, image_inputs, video_inputs = process_mm_info(message, use_audio_in_video=True)

    inputs = processor(
        text=texts,
        audio=audio_inputs,
        images=image_inputs,
        videos=video_inputs,
        return_tensors="pt",
        padding="longest",
    )
    for k in inputs:
        inputs[k] = inputs[k].to(device)

    cache_position = torch.arange(0, inputs['input_ids'].shape[1], device=device)
    inputs = model.prepare_inputs_for_generation(**inputs, use_cache=True, cache_position=cache_position)
    model_outputs = model(**inputs, return_dict=True, output_hidden_states=True)

    last_hidden_state = model_outputs.hidden_states[-1]
    reps = last_hidden_state[:, -1]
    reps = torch.nn.functional.normalize(reps, p=2, dim=-1)
    return reps

🎬 视频检索

示例查询 = '查询：如何制作麻婆豆腐？'
示例视频1 = "https://huggingface.co/Tevatron/OmniEmbed-v0/resolve/main/assets/mapo_tofu.mp4"
示例视频2 = "https://huggingface.co/Tevatron/OmniEmbed-v0/resolve/main/assets/zhajiang_noodle.mp4"
查询 = [{'role': 'user', 'content': [{'type': 'text', 'text': 示例查询}]}]
视频1 = [{'role': 'user', 'content': [{'type': 'video', 'video': 示例视频1}]}]
视频2 = [{'role': 'user', 'content': [{'type': 'video', 'video': 示例视频2}]}]

相似度1 = torch.cosine_similarity(encode_message(查询), encode_message(视频1))
相似度2 = torch.cosine_similarity(encode_message(查询), encode_message(视频2))

print("相似度:", 相似度1.item(), 相似度2.item())

🎵 音频检索

示例查询 = '查询：一首轻快的钢琴曲'
示例音频1 = "https://huggingface.co/Tevatron/OmniEmbed-v0/resolve/main/assets/joe_hisaishi_summer.mp3"
示例音频2 = "https://huggingface.co/Tevatron/OmniEmbed-v0/resolve/main/assets/jay_chou_superman_cant_fly.mp3"
查询 = [{'role': 'user', 'content': [{'type': 'text', 'text': 示例查询}]}]
音频1 = [{'role': 'user', 'content': [{'type': 'audio', 'audio': 示例音频1}]}]
音频2 = [{'role': 'user', 'content': [{'type': 'audio', 'audio': 示例音频2}]}]

相似度1 = torch.cosine_similarity(encode_message(查询), encode_message(音频1))
相似度2 = torch.cosine_similarity(encode_message(查询), encode_message(音频2))

print("相似度:", 相似度1.item(), 相似度2.item())

📈 图像文档检索(图像/图表/PDF)

示例查询 = '查询：Qwen2.5-Omni支持多少种输入模态？'
示例图片1 = "https://huggingface.co/Tevatron/OmniEmbed-v0/resolve/main/assets/qwen2.5omni_hgf.png"
示例图片2 = "https://huggingface.co/Tevatron/OmniEmbed-v0/resolve/main/assets/llama4_hgf.png"
查询 = [{'role': 'user', 'content': [{'type': 'text', 'text': 示例查询}]}]
图片1 = [{'role': 'user', 'content': [{'type': 'image', 'image': 示例图片1}]}]
图片2 = [{'role': 'user', 'content': [{'type': 'image', 'image': 示例图片2}]}]

相似度1 = torch.cosine_similarity(encode_message(查询), encode_message(图片1))
相似度2 = torch.cosine_similarity(encode_message(查询), encode_message(图片2))

print("相似度:", 相似度1.item(), 相似度2.item())

🌍 多语言文本检索

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

相似度1 = torch.cosine_similarity(encode_message(查询), encode_message(文本1))
相似度2 = torch.cosine_similarity(encode_message(查询), encode_message(文本2))

print("相似度:", 相似度1.item(), 相似度2.item())

数据与训练

我们已在Tevatron完全开源训练数据和训练代码

联系方式

本模型由以下开发者共同研发：

庄胜尧、马学广、Samantha Zhan、张晶

如有任何疑问或需要进一步讨论，欢迎随时联系我们。