Scaling Audio-Text Retrieval with Multimodal Large Language Models
Paper • 2602.18010 • Published • 1
How to use Jazzcharles/AuroLA-Omni-3B with sentence-transformers:
from sentence_transformers import SentenceTransformer
model = SentenceTransformer("Jazzcharles/AuroLA-Omni-3B")
sentences = [
"The weather is lovely today.",
"It's so sunny outside!",
"He drove to the stadium."
]
embeddings = model.encode(sentences)
similarities = model.similarity(embeddings, embeddings)
print(similarities.shape)
# [3, 3]AuroLA-Omni-3B is a BF16 omni-modal embedding model built from both Qwen2.5-Omni-3B and AuroLA-3B. It supports text, image, audio, and video inputs through a SentenceTransformers-compatible checkpoint layout.
| Modality / group | Metric | Score |
|---|---|---|
| Image | hit@1 | 65.1 |
| Video | hit@1 | 48.9 |
| VisDoc | ndcg_linear@5 | 69.2 |
| Text | hit@1 | 32.4 |
| Audio | hit@1 | 46.8 |
| Agentic | hit@1 | 36.5 |
| Average | - | 46.7 |
pip install "sentence_transformers[image,audio,video]" "transformers>=5.6.0" qwen-omni-utils
import torch
from sentence_transformers import SentenceTransformer
model_id = "Jazzcharles/AuroLA-Omni-3B"
model = SentenceTransformer(
model_id,
model_kwargs={
"torch_dtype": torch.bfloat16,
"attn_implementation": "sdpa", # use "flash_attention_2" if flash-attn is installed
},
)
model[0].processing_kwargs.update({
"video": {"max_pixels": 64 * 28 * 28, "do_sample_frames": True, "fps": 1},
})
query_embedding = model.encode_query("A person cooking in a kitchen")
image_embeddings = model.encode_document(["path/to/image.jpg"], batch_size=1)
video_embeddings = model.encode_document(["path/to/video.mp4"], batch_size=1)
audio_embeddings = model.encode_document(["path/to/audio.wav"], batch_size=1)
print("image:", model.similarity(query_embedding, image_embeddings))
print("video:", model.similarity(query_embedding, video_embeddings))
print("audio:", model.similarity(query_embedding, audio_embeddings))
When using raw Transformers, use the sentence_transformers chat template explicitly and take the last non-padding hidden state.
import torch
import torch.nn.functional as F
from transformers import AutoProcessor, Qwen2_5OmniThinkerForConditionalGeneration
from qwen_omni_utils import process_mm_info
model_id = "Jazzcharles/AuroLA-Omni-3B"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
model = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained(
model_id,
torch_dtype=torch.bfloat16,
attn_implementation="sdpa", # use "flash_attention_2" if flash-attn is installed
).to(device).eval()
processor.tokenizer.padding_side = "left"
def encode_message(message):
text = processor.apply_chat_template(
message,
chat_template="sentence_transformers",
tokenize=False,
add_generation_prompt=True,
)
if isinstance(text, list):
text = text[0]
text = text + "<|endoftext|>"
audio_inputs, image_inputs, video_inputs = process_mm_info(
message,
use_audio_in_video=False,
)
inputs = processor(
text=text,
audio=audio_inputs,
images=image_inputs,
videos=video_inputs,
videos_kwargs={"do_sample_frames": False, "fps": 1},
return_tensors="pt",
padding="longest",
).to(device)
with torch.inference_mode():
outputs = model(
**inputs,
return_dict=True,
output_hidden_states=True,
use_cache=False,
)
hidden_states = outputs.hidden_states[-1]
attention_mask = inputs["attention_mask"]
last_token_indices = attention_mask.size(1) - 1 - attention_mask.flip(dims=[1]).argmax(dim=1)
reps = hidden_states[torch.arange(hidden_states.size(0), device=device), last_token_indices]
return F.normalize(reps, p=2, dim=-1)
query = [{"role": "user", "content": [{"type": "text", "text": "A person cooking in a kitchen"}]}]
image = [{"role": "user", "content": [{"type": "image", "image": "path/to/image.jpg"}]}]
video = [{
"role": "user",
"content": [{
"type": "video",
"video": "path/to/video.mp4",
"max_pixels": 128 * 28 * 28,
"fps": 1,
}],
}]
audio = [{"role": "user", "content": [{"type": "audio", "audio": "path/to/audio.wav"}]}]
query_embedding = encode_message(query)
for name, document in {"image": image, "video": video, "audio": audio}.items():
similarity = torch.cosine_similarity(query_embedding, encode_message(document))
print(name, similarity.item())
If you use this model, please cite:
@misc{xu2026scalingaudiotextretrievalmultimodal,
title={Scaling Audio-Text Retrieval with Multimodal Large Language Models},
author={Jilan Xu and Carl Thom{\'e} and Danijela Horak and Weidi Xie and Andrew Zisserman},
year={2026},
eprint={2602.18010},
archivePrefix={arXiv},
primaryClass={cs.SD},
url={https://arxiv.org/abs/2602.18010}
}
Base model
Qwen/Qwen2.5-Omni-3B