Views
No views yet
1import torch
2from qwen_omni_utils import process_mm_info
3from transformers import AutoProcessor, Qwen2_5OmniThinkerForConditionalGeneration
4
5def construct_rerank_message(audio_path, text):
6 return [{
7 "role": "user",
8 "content": [
9 {
10 "type": "text",
11 "text": (
12 "I will provide you with a query and a candidate. "
13 "Please evaluate whether the candidate matches the query. "
14 "If it does, respond with 'Yes'; if it doesn't, respond with 'No'."
15 ),
16 },
17 {"type": "text", "text": "Query:"},
18 {"type": "audio", "audio": audio_path},
19 {
20 "type": "text",
21 "text": "Find a caption describing the sound events in the given audio.",
22 },
23 {"type": "text", "text": "Candidate:"},
24 {"type": "text", "text": text},
25 ],
26 }]
27
28
29def process_input(messages, device):
30 texts = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
31 audios, images, videos = process_mm_info(messages, use_audio_in_video=True)
32 inputs = processor(
33 text=texts,
34 audio=audios,
35 images=images,
36 videos=videos,
37 return_tensors="pt",
38 padding=True,
39 use_audio_in_video=True,
40 )
41 inputs = inputs.to(device)
42 return inputs
43
44def rerank_yes_prob(messages):
45 inputs = process_input(messages, device)
46 yes_id = tokenizer("Yes", add_special_tokens=False).input_ids[0]
47 no_id = tokenizer("No", add_special_tokens=False).input_ids[0]
48
49 with torch.inference_mode():
50 outputs = model.generate(
51 **inputs,
52 max_new_tokens=1,
53 output_scores=True,
54 return_dict_in_generate=True,
55 do_sample=False,
56 )
57 logits = outputs.scores[0]
58 yes_probs = torch.softmax(logits[:, [yes_id, no_id]], dim=-1)[:, 0]
59 return yes_probs
60
61
62# 1) Load model + processor
63model_path = "Jazzcharles/AuroLA-rerank-7B" # or your HF repo id
64
65device = "cuda" if torch.cuda.is_available() else "cpu"
66dtype = torch.bfloat16 if device == "cuda" else torch.float32
67
68model = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained(
69 model_path,
70 torch_dtype=dtype,
71 device_map="auto" if device == "cuda" else None,
72)
73processor = AutoProcessor.from_pretrained(model_path, use_fast=False)
74tokenizer = processor.tokenizer
75
76# 2) Prepare one audio-text pair
77audio_files = [
78 "/mnt/data/AudioCaps/audio/--0w1YA1Hm4_30.wav",
79 "/mnt/data/AudioCaps/audio/--0w1YA1Hm4_30.wav",
80]
81text_queries = [
82 "A vehicle driving as a man and woman are talking and laughing",
83 "Wind is blowing and heavy rain is falling and splashing",
84]
85# 3) Build rerank prompt and score
86messages = [construct_rerank_message(audio_path, text) for (audio_path, text) in zip(audio_files, text_queries)]
87yes_prob = rerank_yes_prob(messages)
88yes_prob = [prob.detach().cpu().item() for prob in yes_prob]
89
90print("The probability of the audio matching the text is:", yes_prob)1@misc{xu2026scalingaudiotextretrievalmultimodal,
2 title={Scaling Audio-Text Retrieval with Multimodal Large Language Models},
3 author={Jilan Xu and Carl Thomé and Danijela Horak and Weidi Xie and Andrew Zisserman},
4 year={2026},
5 eprint={2602.18010},
6 archivePrefix={arXiv},
7 primaryClass={cs.SD},
8 url={https://arxiv.org/abs/2602.18010},
9}