1from transformers import AutoProcessor, AutoTokenizer
2from vllm import LLM, SamplingParams
3from qwen_vl_utils import process_vision_info
4
5########################
6# === Configuration ===
7########################
8IMAGE_PATH = "./assets/images/demo_example.jpg"
9QUESTION = "When the canister is momentarily stopped by the spring, by what distance $d$ is the spring compressed?"
10
11MLLM_MODEL_PATH = "KaiChen1998/RACRO-7B-CRO"
12LLM_MODEL_PATH = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B" # feel free to use more advanced reasoners!
13
14########################
15# === Prompts ===
16########################
17SYSTEM_PROMPT_CAP = "You are given an image and a relevant question. Based on the query, please describe the image in details. Do not try to answer the question."
18SYSTEM_PROMPT_LLM = "You are a helpful assistant."
19
20CAPTION_PROMPT = "Question: {}\nPlease describe the image. DO NOT try to answer the question!"
21LLM_PROMPT = """In the following text, you will receive a detailed caption of an image and a relevant question. In addition, you will be provided with a tentative model response. You goal is to answer the question using these information.
22
23### The detailed caption of the provided image: {}
24
25### Note that the caption might contain incorrect solutions, do not be misguided by them.
26
27### A problem to be solved: {}
28
29### A tentative model response: {}
30
31### Note that the above tentative response might be inaccurate (due to calculation errors, incorrect logic/reasoning and so on), under such a case, please ignore it and give your own solutions. However, if you do not have enough evidence to show it is wrong, please output the tentative response."""
32
33########################
34# === Initialize Models ===
35########################
36processor = AutoProcessor.from_pretrained(MLLM_MODEL_PATH)
37tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_PATH)
38
39mllm = LLM(model=MLLM_MODEL_PATH, tensor_parallel_size=1, gpu_memory_utilization=0.8,
40 device='cuda:0', dtype="bfloat16", limit_mm_per_prompt={"image": 1})
41
42llm = LLM(model=LLM_MODEL_PATH, tensor_parallel_size=1, gpu_memory_utilization=0.8,
43 device='cuda:1', dtype="bfloat16")
44
45mllm_sampling = SamplingParams(temperature=0, max_tokens=8192)
46llm_sampling = SamplingParams(temperature=0.6, top_p=0.95, max_tokens=8192)
47
48########################
49# === Build Prompts ===
50########################
51def build_messages(image_path, question):
52 cap_msgs = [
53 {"role": "system", "content": SYSTEM_PROMPT_CAP},
54 {"role": "user", "content": [{"type": "image", "image": image_path}, {"type": "text", "text": CAPTION_PROMPT.format(question)}]}
55 ]
56 qa_msgs = [
57 {"role": "user", "content": [{"type": "image", "image": image_path}, {"type": "text", "text": question + " Please think step by step. The final answer MUST BE put in \\boxed{}."}]}
58 ]
59 return cap_msgs, qa_msgs
60
61# === Run Captioning and QA ===
62def run_mllm(image_tensor, cap_prompt, qa_prompt):
63 cap_output = mllm.generate([{"multi_modal_data": {"image": image_tensor}, "prompt": cap_prompt[0]}], sampling_params=mllm_sampling)
64 qa_output = mllm.generate([{"multi_modal_data": {"image": image_tensor}, "prompt": qa_prompt[0]}], sampling_params=mllm_sampling)
65 return cap_output[0].outputs[0].text, qa_output[0].outputs[0].text
66
67# === Final Reasoning Step ===
68def run_llm_reasoning(caption, question, answer):
69 messages = [
70 {"role": "system", "content": SYSTEM_PROMPT_LLM},
71 {"role": "user", "content": LLM_PROMPT.format(caption, question, answer)}
72 ]
73 prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
74 output = llm.generate([{"prompt": prompt}], sampling_params=llm_sampling)
75 return output[0].outputs[0].text
76
77########################
78# === Pipeline ===
79########################
80cap_msgs, qa_msgs = build_messages(IMAGE_PATH, QUESTION)
81cap_prompt = processor.apply_chat_template([cap_msgs], tokenize=False, add_generation_prompt=True)
82qa_prompt = processor.apply_chat_template([qa_msgs], tokenize=False, add_generation_prompt=True)
83
84image_tensor, _ = process_vision_info(cap_msgs)
85caption_text, tentative_answer = run_mllm(image_tensor, cap_prompt, qa_prompt)
86final_answer = run_llm_reasoning(caption_text, QUESTION, tentative_answer)
87
88print("Final Answer:\n", final_answer)