1import re
2from io import BytesIO
3
4import requests
5import torch
6from PIL import Image
7from transformers import (
8 AutoModelForImageTextToText,
9 AutoProcessor,
10 Qwen2_5_VLForConditionalGeneration,
11)
12
13REAG_MODEL_NAME = "aimagelab/ReAG-7B"
14CRITIC_MODEL_NAME = "aimagelab/ReAG-Critic"
15
16SYSTEM_PROMPT_REASONING = (
17 "A conversation between User and Assistant. The user asks a question, "
18 "and the Assistant solves it. The assistant first thinks about the "
19 "reasoning process in the mind and then provides the user with the answer. "
20 "The reasoning process and answer are enclosed within <think> </think> and "
21 "<answer> </answer> tags, respectively, i.e., "
22 "<think>reasoning process here</think><answer>short answer here</answer>"
23)
24
25RELEVANCY_EVAL_SYSTEM_PROMPT = """You are a multimodal reasoning assistant specialized in Knowledge-Based Visual Question Answering (KB-VQA).
26Your task is to evaluate whether a given text passage provides useful and relevant information for answering a question about an image.
27
28You will be given:
29- Image: a visual scene containing entities, actions, and context.
30- Question: a natural-language question that refers to the image.
31- Text Passage: an external knowledge snippet retrieved from a database or the web.
32
33You must analyze the semantic alignment between the text, the image, and the question.
34Follow these steps carefully before giving your final decision:
351. Understand the visual scene: Identify the key objects, people, actions, and context visible in the image.
362. Interpret the question: Determine what information the question seeks.
373. Analyze the text passage: Extract the main claims, facts, and entities mentioned in the text.
38
39Compare for relevance: Assess whether the information in the text:
40- Contains at least one sentence that supports answering the question about the image, OR
41- Provides background knowledge needed to interpret or reason about the image-question pair.
42
43Important:
44- If even a single sentence in the passage is relevant or useful, consider the entire passage as relevant and answer "Yes".
45- If no part of the passage contributes meaningfully to answering the question, answer "No".
46
47Output only one word:
48"Yes" -> if the text provides relevant or useful information for answering the question.
49"No" -> if the text is irrelevant or unhelpful."""
50
51SECTION_EVAL_USER_TEMPLATE = """Here is the question on the image above:
52{question}
53
54Here is the text passage to analyze:
55{passage}
56
57Does the text passage contain at least one sentence that may have some information useful to answer the user question?
58"Yes"/"No" answer:"""
59
60CONTEXT_VQA_PROMPT = """\
61{question}
62
63The following paragraphs may contain useful information to help answer the question correctly:
64{context}
65"""
66
67
68def load_image(image_url: str) -> Image.Image:
69 response = requests.get(image_url, timeout=30, headers={"User-Agent": "Mozilla/5.0"})
70 response.raise_for_status()
71 return Image.open(BytesIO(response.content)).convert("RGB")
72
73
74def get_model_kwargs():
75 if torch.cuda.is_available():
76 return {
77 "device": "cuda",
78 "device_map": "balanced",
79 "torch_dtype": torch.bfloat16,
80 "attn_implementation": "flash_attention_2",
81 }
82 return {
83 "device": "cpu",
84 "device_map": "auto",
85 "torch_dtype": torch.float32,
86 }
87
88
89def parse_reag_output(text: str):
90 answer_match = re.search(r"<answer>(.*?)</answer>", text, re.DOTALL)
91 think_match = re.search(r"<think>(.*?)</think>", text, re.DOTALL)
92 return {
93 "raw_output": text.strip(),
94 "reasoning": think_match.group(1).strip() if think_match else "",
95 "answer": answer_match.group(1).strip() if answer_match else text.strip(),
96 }
97
98
99def run_reag_generator(model, processor, image: Image.Image, question: str):
100 messages = [
101 {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT_REASONING}]},
102 {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": question}]},
103 ]
104 prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
105 inputs = processor(text=[prompt + "<think>"], images=[image], return_tensors="pt", padding=True)
106 inputs = {k: v.to(model.device) for k, v in inputs.items()}
107
108 generated_ids = model.generate(
109 **inputs, max_new_tokens=512, stop_strings=["</answer>"], tokenizer=processor.tokenizer
110 )
111 input_length = inputs["input_ids"].shape[1]
112 generated_text = "<think>" + processor.batch_decode(
113 generated_ids[:, input_length:], skip_special_tokens=True, clean_up_tokenization_spaces=False
114 )[0]
115 return parse_reag_output(generated_text)
116
117
118def run_reag_critic(critic, processor, image: Image.Image, question: str, passage: str, yes_prob_threshold: float = 0.1):
119 messages = [
120 {"role": "system", "content": [{"type": "text", "text": RELEVANCY_EVAL_SYSTEM_PROMPT}]},
121 {
122 "role": "user",
123 "content": [
124 {"type": "image"},
125 {"type": "text", "text": SECTION_EVAL_USER_TEMPLATE.format(question=question, passage=passage)},
126 ],
127 },
128 ]
129 prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
130 inputs = processor(images=[image], text=[prompt], return_tensors="pt", padding=True)
131 inputs = {k: v.to(critic.device) for k, v in inputs.items()}
132
133 with torch.inference_mode():
134 outputs = critic(**inputs)
135 logits = outputs.logits[:, -1, :].float()
136 probs = torch.softmax(logits, dim=-1)
137
138 yes_token_id = processor.tokenizer.convert_tokens_to_ids("Yes")
139 no_token_id = processor.tokenizer.convert_tokens_to_ids("No")
140 return {
141 "relevant": probs[0, yes_token_id].item() > yes_prob_threshold,
142 "yes_probability": probs[0, yes_token_id].item(),
143 "no_probability": probs[0, no_token_id].item(),
144 }
145
146
147# ── Example ──────────────────────────────────────────────────────────────────
148
149image = load_image(
150 "https://upload.wikimedia.org/wikipedia/commons/thumb/5/54/Clinopodium_vulgare_inflorescence.jpg/250px-Clinopodium_vulgare_inflorescence.jpg"
151)
152question = "What kind of properties does this plant have?"
153passages = [
154 "# Description:\nWild basil is a perennial rhizomatous herb ...",
155 "# Distribution:\nWild basil occurs in suitable locations in most of Europe ...",
156 "# Uses:\nThe leaves of wild basil are used as an aromatic herb ... It has been shown to have anti-bacterial properties.",
157]
158
159model_kwargs = get_model_kwargs()
160device = model_kwargs.pop("device")
161
162# 1. Load and run the critic
163critic_processor = AutoProcessor.from_pretrained(CRITIC_MODEL_NAME, padding_side="left", use_fast=True)
164critic = Qwen2_5_VLForConditionalGeneration.from_pretrained(CRITIC_MODEL_NAME, **model_kwargs)
165critic.eval()
166
167relevant_passages = []
168for passage in passages:
169 result = run_reag_critic(critic, critic_processor, image, question, passage)
170 if result["relevant"]:
171 relevant_passages.append(passage)
172
173# 2. Load the generator and answer with filtered context
174context = "\n\n\n".join(relevant_passages) if relevant_passages else ""
175processor = AutoProcessor.from_pretrained(REAG_MODEL_NAME, padding_side="left", use_fast=True)
176generator = AutoModelForImageTextToText.from_pretrained(REAG_MODEL_NAME, **model_kwargs)
177generator.eval()
178
179question_with_context = CONTEXT_VQA_PROMPT.format(question=question, context=context)
180output = run_reag_generator(generator, processor, image, question_with_context)
181print("Answer:", output["answer"])
182print("Reasoning:", output["reasoning"])