ViDoRAG/vidorag_agents.py at main · Alibaba-NLP/ViDoRAG · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
from PIL import Image
import os
import json
import sys

from agent.agent_prompt import seeker_prompt,inspector_prompt,answer_prompt
from agent.map_dict import arrangement_map_dict,page_map_dict_normal,page_map_dict

from utils.parse_tool import extract_json
from utils.image_preprosser import concat_images_with_bbox

class Seeker:
    def __init__(self, vlm):
        self.vlm = vlm
        self.seeker_multi_image = False

        if self.seeker_multi_image:
            self.page_map = page_map_dict_normal
        else:
            self.page_map = page_map_dict
    def run(self, query=None, images_path=None, feedback=None):

        if query is not None and images_path is not None:
            self.buffer_images = images_path
            self.query = query
            prompt = seeker_prompt.replace('{question}', self.query).replace('{page_map}', self.page_map[len(self.buffer_images)])

        elif feedback is not None:
            additional_information = self.query + '\n\n## Additional Information\n' + feedback
            prompt = seeker_prompt.replace('{question}', additional_information).replace('{page_map}', self.page_map[len(self.buffer_images)])

        if self.seeker_multi_image:
            input_images = self.buffer_images
        else:
            input_images = [concat_images_with_bbox(self.buffer_images, arrangement=arrangement_map_dict[len(self.buffer_images)], scale=1, line_width=40)]

        times = 0
        while True:
            if times > 2:
                # return None, None, None
                raise Exception('seeker time out')
            times += 1
            select_response = self.vlm.generate(query=prompt, image=input_images)
            print(select_response)
            try:
                select_response_json = extract_json(select_response)
                reason = select_response_json.get('reason', None)
                summary = select_response_json.get('summary', None)
                select_page_num = select_response_json.get('choice', None)
                if reason is None or summary is None or any([page >= len(self.buffer_images) for page in select_page_num]):
                    raise Exception(f'select json format error: length: {len(self.buffer_images)}')

                selected_images = [self.buffer_images[page] for page in select_page_num]
                self.buffer_images = [image for image in self.buffer_images if image not in selected_images]

            except Exception as e:
                print(e)
                print(select_response)
                print('seeker')
                continue
            break

        return selected_images, summary, reason

class Inspector:
    def __init__(self, vlm):
        self.vlm = vlm
        self.inspector_multi_image = False

        if self.inspector_multi_image:
            self.page_map = page_map_dict_normal
        else:
            self.page_map = page_map_dict

        self.buffer_images = []

    def run(self, query, images_path):
        # answer or not, (images, candidate)/feedback
        if len(self.buffer_images) == 0 and len(images_path) == 0:
            return None, None, None
        elif len(images_path) == 0:
            return 'synthesizer', None, self.buffer_images
        elif len(images_path) != 0:
            self.buffer_images.extend(images_path)

        if self.inspector_multi_image:
            input_images = self.buffer_images
        else:
            input_images = [concat_images_with_bbox(self.buffer_images, arrangement=arrangement_map_dict[len(self.buffer_images)], scale=1, line_width=40)]

        prompt = inspector_prompt.replace('{question}',query).replace('{page_map}',self.page_map[len(self.buffer_images)])

        times = 0
        while True:
            if times >2:
                raise Exception('Inspector time out')
                # return None, None, None
            times +=1

            response = self.vlm.generate(query=prompt,image=input_images)
            print(response)
            try:
                response_json = extract_json(response)
                # thought
                reason = response_json.get('reason',None)
                # if feedback
                info = response_json.get('information',None)
                choice = response_json.get('choice',None)
                # can answer
                answer = response_json.get('answer',None)
                ref = response_json.get('reference',None)

                if reason is None:
                    raise Exception('answer no reason')
                elif answer is not None and ref is not None:
                    if any([page >= len(self.buffer_images) for page in ref]) or len(ref)==0:
                        raise Exception('ref error')
                    if len(ref) == len(self.buffer_images):
                        return 'answer', answer, self.buffer_images
                    else:
                        ref_images = [self.buffer_images[page] for page in ref]
                        return 'synthesizer', answer, ref_images
                elif info is not None and choice is not None:
                    if any([page >= len(self.buffer_images) for page in choice]):
                        raise Exception('choice error')
                    self.buffer_images = [self.buffer_images[page] for page in choice]
                    return 'seeker', info, self.buffer_images

            except Exception as e:
                print(e)
                print("inspector")
                continue

class Synthesizer:
    def __init__(self, vlm):
        self.vlm = vlm
        self.synthesizer_multi_image = False
        if self.synthesizer_multi_image:
            self.page_map = page_map_dict_normal
        else:
            self.page_map = page_map_dict

    def run(self, query, candidate_answer, ref_images):
        if candidate_answer is not None:
            query = query + '\n\n## Related Information\n' + candidate_answer
        prompt = answer_prompt.replace('{question}',query).replace('{page_map}',self.page_map[len(ref_images)])

        if self.synthesizer_multi_image:
            input_images = ref_images
        else:
            input_images = [concat_images_with_bbox(ref_images, arrangement=arrangement_map_dict[len(ref_images)], scale=1, line_width=40)]

        while True:
            final_answer_response = self.vlm.generate(query=prompt,image=input_images)
            print(final_answer_response)
            try:
                final_answer_response_json = extract_json(final_answer_response)
                reason = final_answer_response_json.get('reason',None)
                answer = final_answer_response_json.get('answer',None)
                if reason is None or answer is None :
                    raise Exception('Synthesizer time out')
                return reason, answer
            except Exception as e:
                print(e)
                print(final_answer_response)
                print("answer")
                continue

class ViDoRAG_Agents:
    def __init__(self, vlm):
        self.seeker = Seeker(vlm)
        self.inspector = Inspector(vlm)
        self.synthesizer = Synthesizer(vlm)

    def run_agent(self, query, images_path):
        # initial
        self.seeker.buffer_images = None
        self.inspector.buffer_images = []

        selected_images, summary, reason = self.seeker.run(query=query, images_path=images_path)
        # iter
        while True:
            status, information, images = self.inspector.run(query, selected_images)
            if status == 'answer':
                return information
            elif status == 'synthesizer':
                reason, answer = self.synthesizer.run(query, information, images)
                return answer
            elif status == 'seeker':
                selected_images, summary, reason = self.seeker.run(feedback=information)
                continue
            else:
                print('No related information')
                return None

if __name__ == '__main__':
    from llms.llm import LLM
    vlm = LLM('qwen-vl-max')
    agent = ViDoRAG_Agents(vlm)
    re=agent.run_agent(query='Who is Tim?', images_path=['./data/ExampleDataset/img/00a76e3a9a36255616e2dc14a6eb5dde598b321f_1.jpg','./data/ExampleDataset/img/00a76e3a9a36255616e2dc14a6eb5dde598b321f_2.jpg','./data/ExampleDataset/img/00a76e3a9a36255616e2dc14a6eb5dde598b321f_3.jpg','./data/ExampleDataset/img/00a76e3a9a36255616e2dc14a6eb5dde598b321f_4.jpg'])
    print(re)