File size: 12,165 Bytes
cd6775d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
from socratic_node import SocraticNode

import os
import json
import re
import pdb
import torch
import time
from openai import OpenAI
import requests

client = OpenAI(
        api_key=os.environ.get("CHATANYWHERE_API_KEY", ""),
        base_url='https://api.chatanywhere.tech/v1', 
    )

# client = OpenAI(api_key=os.environ.get("DEEPSEEK_API_KEY", ""), base_url="https://api.deepseek.com")

# [ragqa] 可插拔 LLM:设为 callable(system_define:str, prompt:str)->str 即用它替代默认的 gpt-4o-mini。
_LLM_OVERRIDE = None


def set_llm(fn):
    """注入自定义 LLM(如本地模型)。传 None 恢复默认 gpt-4o-mini。"""
    global _LLM_OVERRIDE
    _LLM_OVERRIDE = fn


# [ragqa] 可插拔检索器:设为 callable(query:str)->str(top-1 passage 文本)即用它替代默认 HTTP(8895)。
# 进程内函数版默认注入空检索桩,从而完全不发 HTTP;纯查询分解不需要真实检索。
_RETRIEVER_OVERRIDE = None


def set_retriever(fn):
    """注入自定义检索器。传 None 恢复默认 HTTP ColBERT。"""
    global _RETRIEVER_OVERRIDE
    _RETRIEVER_OVERRIDE = fn


def get_retrieval(query_item):
        if _RETRIEVER_OVERRIDE is not None:
            return _RETRIEVER_OVERRIDE(query_item)
        url = 'http://localhost:8895/api/search?query='+query_item+'&k=1'
        response = requests.get(url=url)
        res_dic = response.json()
        corpus_list_topk = res_dic['topk']
        top1_passage = corpus_list_topk[0]['text']
        
        return top1_passage

class SocraticTree:
    def __init__(self, backbone, api, prompt_map, dataset, num_question, max_turn, max_depth, save_dir):
        self.backbone = backbone
        if backbone == 'gpt':
            pass
           
        elif backbone == 'falcon':
            # host model by text-generation: 
            from text_generation import Client
            self.pipeline = Client("http://127.0.0.1:8080")
            
        self.prompt_map = prompt_map
        self.dataset = dataset
        self.num_question = num_question
        self.max_turn = max_turn
        self.max_depth = max_depth
        
        self.root_node = None
        
        self.save_dir = save_dir    
        self.log_path = None
    
        
    def select_prompt(self, isMultipleChoice, type):
        try:
            # question_type = 'multipleChoice' if isMultipleChoice else 'normalQA'
            question_type = 'multipleChoice_cot' if isMultipleChoice else 'normalQA'
            prompt = self.prompt_map[type][question_type][self.dataset]['prompt']
            tip = self.prompt_map[type][question_type][self.dataset]['tip']
            system_define = self.prompt_map[type][question_type][self.dataset]['system_define']
        except Exception as e:
            print(e)
            pdb.set_trace()
        return prompt, tip, system_define
    
# while True:
#             try:
#                 global client
#                 response = client.chat.completions.create(
#                     model='deepseek-chat',

#                     messages=[
#                         {"role": "system", "content": system_define},
#                         {"role": "user", "content": prompt},
#                     ]
#                 )
#                 break
#             except Exception as e:
#                 # current time in human-readable format
#                 obj = time.localtime()
#                 t = time.asctime(obj)
#                 print(t, 'Request failed. Retrying...')
#                 pdb.set_trace()
#                 # time.sleep(10)
    def request_gpt(self, system_define, prompt):

        # [ragqa] 若注入了自定义 LLM,则用它(进程内本地模型等),否则走默认 gpt-4o-mini。
        if _LLM_OVERRIDE is not None:
            return _LLM_OVERRIDE(system_define, prompt).strip()

        response = client.chat.completions.create(
                    model='gpt-4o-mini',
                    messages=[
                        {"role": "system", "content": system_define},
                        {"role": "user", "content": prompt},
                    ]
                )
                
        return response.choices[0].message.content.strip()
    
    
    
    def request_falcon(self, system_define, prompt):
        input_str = system_define + '\n\n\n' + prompt
        answer = self.pipeline.generate(input_str, max_new_tokens=300).generated_text 
        return answer.strip()
    
    def request(self, system_define, prompt):
        if self.backbone == 'gpt':
            return self.request_gpt(system_define, prompt)
        elif self.backbone == 'falcon':
            return self.request_falcon(system_define, prompt)
    
    def qa2hint(self, question, answer):
        # request chatGPT
        system_define = """Imagine you are an editor. You are given a question-and-answer pair. You need to merge the question and answer into a statement sentence."""
        prompt = "Question: " + question + "\nAnswer: " + answer
        return self.request(system_define, prompt)
    
    def get_clean_answer(self, raw_answer):
        answer = raw_answer
        # remove " in the raw_answer string
        while '"' in answer:
            answer = answer.replace('"', '')
        if ': ' in answer:
            answer = answer.split(': ')[1]
        if 'the answer is ' in answer:
            idx = answer.index('the answer is ')
            answer = answer[idx+14]
        if 'Option ' in answer:
            answer = answer.split('Option ')[1][0]
        if len(answer) > 0 and answer[-1] == '.':
            answer = answer[:-1]
        if 'A. ' in answer:
            answer = 'A'
        elif 'B. ' in answer:
            answer = 'B'
        elif 'C. ' in answer:
            answer = 'C'
        elif 'D. ' in answer:
            answer = 'D'
        elif '. ' in raw_answer:
            answer = raw_answer.split('. ')[1]
        if len(answer) == 2 and answer[1] == '.':
            answer = answer[0]
        return answer
    
        
    def answer_question(self, node):
        passage = get_retrieval(node.question)
        
        system_define = "You are an expert in open domain Q&A and your task is to assess how relevant the retrieved passages are to the question. According to the degree of help of the retrieved corpus to answer the question, it can be divided into three correlation degrees: high, middle and low."
        tip = "Only allowed to answer high or middle or low. If you are not sure, choose the one you think is more likely."
        
        prompt = ''
        prompt += 'Question: '
        prompt += node.question + '\n'
        prompt += 'Passage: '
        prompt += passage + '\n'
        
        prompt += tip + '\n'
        prompt += 'Answer: '
        relevant = self.request(system_define, prompt)
        
        if 'high' in relevant or 'High' in relevant:
            confidence = 'high'
        elif 'middle' in relevant or 'Middle' in relevant:
            confidence = 'middle'
        elif 'low' in relevant or 'Low' in relevant:
            confidence = 'low'
        else:
            confidence = 'middle'
        
        return confidence
  
            
    
    def raise_question(self, node):
        system_define = "Imagine that you are a thoughtful and logical problem solver. You will get a question. However, this question is too complex or lacking information to answer. You need to break down the original question into several simpler subquestions to help the information retrieval system retrieve the relevant information. Important: Do not use pronouns or indefinite pronoun phrases in generated questions. The questions asked must be self-contained questions. Each question can contain only one parameter. Don't just ask yes/no questions."
        
        prompt = ''
        prompt += "For example, Question: Could the Great Wall of China connect the Dodgers to the White Sox? Note: The raised question has to be a self-contained question. Do not use pronouns or indefinite pronoun phrases in the generated questions. Copy context from the original question if needed. Deep Questions: 1. What is the most commonly cited figure for the total length of the Great Wall? 2. What is the straight-line distance between Chicago, Illinois, and Los Angeles, California?" + '\n'
        prompt += 'Question: '
        prompt += node.question + '\n'
        prompt += 'Deep Questions: '
        
        raw_answer = self.request(system_define, prompt)
        
        # post-process answer
        ls_questions = raw_answer.split('\n')
        ls_questions = [re.sub(r'^\d+\.\s+', '', q) for q in ls_questions]
        ls_questions = ls_questions[:self.num_question] if len(ls_questions) > self.num_question else ls_questions
        return ls_questions
    
    
        
    def self_questioning(self, node):
        confidence = self.answer_question(node)      
        if (confidence == 'high' or node.isLeaf) and (node.hasHint() or node.depth != 1):
            # self.log(node, confidence)
            return False, confidence
        else:
            ls_deepQuestion = self.raise_question(node)
            # self.log(node, confidence, ls_deepQuestion=ls_deepQuestion)
            return True, ls_deepQuestion
        
    
    def run(self, node, collected_questions=None):
        if collected_questions is None:
            collected_questions = []

        # self questioning
        continue_deeper, output = self.self_questioning(node)
        if continue_deeper:  # output is subquestion list
            for deep_question in output:
                # Collect the deep question
                collected_questions.append(deep_question)

                deep_node = SocraticNode(
                    deep_question,
                    node.turn,
                    node.depth + 1,
                    self.max_turn,
                    self.max_depth,
                    context=node.context
                )
                deep_answer = self.run(deep_node, collected_questions)
                node.add_hint(1)
            node.update_turn_num()
            return self.run(node, collected_questions)

        else:  # output is answer
            return collected_questions

        
        
    def start(self, id, question, options, context=None):
        self.log_path = self.save_dir + '/log/' + str(id) + '.txt'
        if os.path.exists(self.log_path):
            os.remove(self.log_path)
        self.root_node = SocraticNode(question, 1, 1, self.max_turn, self.max_depth, context=context, options=options)
        collected_questions = self.run(self.root_node)
        return collected_questions
    
        
    def log(self, node, raw_answer, ls_deepQuestion=None):
        turn = node.turn
        depth = node.depth
        isLeaf = node.isLeaf
        context = node.context
        question = node.question
        options = None
        if node.isMultipleChoice():
            options = node.add_optionID_toText()
        hints = node.get_textHints()
        raw_answer = raw_answer
        answer = str(node.answer)
        ls_deepQuestion = ls_deepQuestion       
        
        with open(self.log_path, 'a') as f:
            f.write('=========================================================================================================\n')
            f.write('Turn: ' + str(turn) + ', Depth: ' + str(depth) + ', isLeaf: ' + str(isLeaf) + '\n')
            if context is not None:
                f.write('Context:\t' + context + '\n')
            f.write('Question:\t' + question + '\n')
            if options is not None:
                f.write('Options:\t' + options + '\n')
            if len(hints) > 0:
                f.write('Hints:\t' + hints + '\n')
            f.write('Raw Answer:\t' + raw_answer + '\n')
            f.write('Answer:\t' + answer + '\n')
            if ls_deepQuestion is not None:
                f.write('\nDeep Questions:\n')
                for i, deepQuestion in enumerate(ls_deepQuestion):
                    f.write(str(i+1) + '. ' + deepQuestion + '\n')
            f.write('\n')