open-set UoT

Author: ChorlingLau

run.py

@@ -20,10 +20,20 @@ def run(args):
         args.task_end_index = min(args.task_end_index, len(task.data))
 
     if args.naive_run:
-        log_file = f'./logs/{args.task}/{args.guesser_model}_as_guesser/{args.dataset}_{args.temperature}_naive_{"" if args.inform else "un"}inform_EXAMINER{args.examiner_model}_{args.task_start_index}-{args.task_end_index}.json'
+        log_file = (f'./logs/{args.task}/{args.guesser_model}_as_guesser/{args.dataset}_{args.temperature}'
+                    f'_naive_{"" if args.inform else "un"}inform_EXAMINER{args.examiner_model}'
+                    f'_{args.task_start_index}-{args.task_end_index}.json')
     else:
-        log_file = f'./logs/{args.task}/{args.guesser_model}_as_guesser/{args.dataset}_{args.temperature}_lambda{args.reward_lambda}_L{args.n_extend_layers}_K{args.n_potential_actions}_PRUN{args.n_pruned_nodes}_EXAMINER{args.examiner_model}_{args.task_start_index}-{args.task_end_index}.json'
-        root_file = f'./roots/{args.task}/{args.guesser_model}_{args.dataset}_{args.temperature}_root.pickle'
+        log_file = (f'./logs/{args.task}/{args.guesser_model}_as_guesser/'
+                    f'{f"OS_init{args.open_set_size}_renew{args.size_to_renew}_" if args.open_set_size > 0 else ""}'
+                    f'{f"pre{args.n_pre_ask}_" if args.n_pre_ask > 0 else ""}'
+                    f'{args.dataset}_{args.temperature}_lambda{args.reward_lambda}_acc{not args.none_acc_reward}'
+                    f'_exp{args.expected_reward_method}_L{args.n_extend_layers}_K{args.n_potential_actions}'
+                    f'_PRUN{args.n_pruned_nodes}_EXAMINER{args.examiner_model}'
+                    f'_{args.task_start_index}-{args.task_end_index}.json')
+        root_file = (f'./roots/{args.task}/{args.guesser_model}'
+                     f'{f"OS_init{args.open_set_size}_" if args.open_set_size > 0 else ""}'
+                     f'_{args.dataset}_{args.temperature}_root.pickle')
         if os.path.exists(root_file):
             r = open(root_file, 'rb')
             root = pickle.load(r)
@@ -55,7 +65,7 @@ def run(args):
 
 def parse_args():
     args = argparse.ArgumentParser()
-    args.add_argument('--guesser_model', type=str, default='gemini-1.0-pro',
+    args.add_argument('--guesser_model', type=str, default='gpt-3.5-turbo',
                       choices=['gpt-4', 'gpt-3.5-turbo',
                                '_claude-2', 'claude-3-opus-20240229', 'claude-3-sonnet-20240229',
                                'palm-2', 'cohere', 'llama-2-70b-chat',
@@ -66,10 +76,13 @@ def parse_args():
 
     args.add_argument('--task', type=str, default='20q',
                       choices=['20q', 'md', 'tb'])
-    args.add_argument('--dataset', type=str, default='bigbench',
-                      choices=['bigbench', 'common', 'DX', 'MedDG', 'FloDial'])
+    args.add_argument('--dataset', type=str, default='common',
+                      choices=['bigbench', 'common', 'thing', 'DX', 'MedDG', 'FloDial'])
     args.add_argument('--task_start_index', type=int, default=-1)
     args.add_argument('--task_end_index', type=int, default=-1)
+    args.add_argument('--open_set_size', type=int, default=-1)
+    args.add_argument('--size_to_renew', type=int, default=-1)  # only used when open_set_size > 0
+    args.add_argument('--n_pre_ask', type=int, default=0)  # only used when open_set_size > 0 and data doesn't contain self-repo
 
     args.add_argument('--naive_run', action='store_true', default=False)
     args.add_argument('--inform', action='store_true', default=False)  # only used when naive_run
@@ -85,6 +98,9 @@ def parse_args():
     args.add_argument('--expected_action_tokens', type=int, default=50)
     args.add_argument('--expected_target_tokens', type=int, default=10)
 
+    args.add_argument('--none_acc_reward', action='store_true', default=False)
+    args.add_argument('--expected_reward_method', type=str, default='avg', choices=['avg', 'max'])
+
     args = args.parse_args()
     return args
 

src/uot/chat_utils.py

@@ -1,3 +1,4 @@
+import copy
 import importlib
 
 from uot.models import get_response_method
@@ -79,7 +80,7 @@ def cls_given_repo(task, items: list, repo, translate=False, self_repo=True):
             message = [{"role": "user", "content": f"Translate to English: {repo}"}]
             gpt3_response = get_response_method("gpt-3.5-turbo")
             repo = gpt3_response(message, model="gpt-3.5-turbo", max_tokens=500)
-        repo = task.prompts.self_report_prompt.format(repo=repo)
+        repo = task.prompts.self_repo_prompt.format(repo=repo)
     else:
         repo = task.prompts.free_answer_prompt.format(repo=repo)
     message = [{"role": "user", "content": task.prompts.classify_prompt.format(item_list_str=', '.join(items), repo=repo)}]
@@ -106,3 +107,35 @@ def extract_items(rsp, keyword):
     except Exception as e:
         print(e)
         return cls_given_repo(task, items, repo, translate, self_repo)
+
+
+def initialize_open_set(task, repo=""):
+    response = get_response_method(task.guesser_model)
+    size = task.open_set_size
+    
+    if isinstance(repo, str):
+        message = [{"role": "user", "content": task.prompts.init_open_set_prompt.format(repo=repo, size=size)}]
+    else:
+        message = repo + [{"role": "user", "content": task.prompts.init_open_set_prompt.format(size=size)}]
+    rsp = response(message, model=task.guesser_model, max_tokens=15*size)
+    print([rsp])
+    try:
+        rsp = set(eval(rsp))
+        return list(rsp)
+    except Exception as e:
+        print(e)
+        return initialize_open_set(task, repo)
+
+
+def renew_open_set(task, history, items):
+    response = get_response_method(task.guesser_model)
+    size = task.open_set_size
+    message = copy.deepcopy(history) + [{"role": "user", "content": task.prompts.renew_open_set_prompt.format(size=size, item_list=str(items))}]
+    rsp = response(message, model=task.guesser_model, max_tokens=15*size)
+    print([rsp])
+    try:
+        rsp = set(eval(rsp))
+        return list(rsp)
+    except Exception as e:
+        print(e)
+        return renew_open_set(task, history, items)

src/uot/data/data_20q.py

@@ -6,3 +6,4 @@
 Objects = ['Hula hoop', 'Calendar', "King Tut's mask", 'CD-ROM', 'Pajamas', 'Treehouse', 'Rocking chair', 'The Mona Lisa', 'T-Rex', 'Light bulb', 'Palm tree', 'Balloon', 'The Crown Jewels', 'Wrapping paper', 'Penny', 'Notebook', 'Fire extinguisher', 'Napkin', 'Beret', 'The Titanic', 'Blender', 'Stamp', 'Yacht', 'Volleyball', 'Tissues', 'Comet', 'Hairbrush', 'Mittens', 'Chopsticks', 'Magazine', 'Piccolo', 'Northern Lights', 'Chessboard', 'Christmas tree', 'Stained glass', 'Hollywood sign', 'Tennis court']
 COMMON = Animals + Food + Objects + Places
 
+THING200 = ['trombone', 'monkey', 'quad', 'speedometer', 'wreck', 'cockroach', 'butterfly', 'cookie', 'hat', 'uniform', 'ferry', 'yarn', 'razor blade', 'cigarette holder', 'rope', 'knife', 'snowboard', 'bone', 'book', 'vest', 'easter egg', 'panda', 'crepe', 'sandal', 'sandpaper', 'brussels sprouts', 'wick', 'wax', 'bullet', 'screw', 'holster', 'train set', 'crayfish', 'needle', 'elephant', 'paint', 'sweater', 'book', 'mussel', 'dandelion', 'seagull', 'float', 'shutter', 'altar', 'bagel', 'coil', 'funnel', 'pie', 'lemon', 'pasta', 'magnifier', 'cornucopia', 'muffin', 'scarecrow', 'whiteboard', 'scraper', 'gargoyle', 'copier', 'rose', 'banner', 'braid', 'dumbwaiter', 'cat', 'gargoyle', 'pepper mill', 'squirrel', 'air conditioner', 'chariot', 'chessboard', 'ice-cream cone', 'bread', 'motherboard', 'bug', 'space shuttle', 'barcode', 'plate', 'box', 'counter', 'breakfast', 'lavender', 'slug', 'coral', 'lipstick', 'soccer ball', 'wick', 'trowel', 'olive', 'gate', 'ship', 'scarecrow', 'cello', 'man', 'barrel', 'lip balm', 'armor', 'flamingo', 'rock', 'sloth', 'buggy', 'cooler', 'coffee', 'basketball', 'bulldozer', 'whoopee cushion', 'breakfast', 'golf cart', 'album', 'milk', 'cash machine', 'potpie', 'potato', 'granite', 'slot machine', 'footprint', 'suit', 'jeep', 'mop', 'garter', 'wine cooler', 'box', 'gas mask', 'spool', 'hookah', 'razor', 'chips', 'pet food', 'canvas', 'polo shirt', 'shield', 'boy', 'plunger', 'treasure', 'cabinet', 'stew', 'dolly', 'frog', 'knitting needle', 'gyroscope', 'satellite', 'clasp', 'anklet', 'lasagna', 'crane', 'pigeon', 'grape', 'carrot', 'mold', 'denture', 'highlighter', 'trigger', 'furnace', 'spur', 'pantsuit', 'spareribs', 'bull', 'shell', 'footrest', 'scuba', 'cooker', 'plum', 'accordion', 'saw', 'gourd', 'hail', 'knitting needle', 'stake', 'bone', 'ruby', 'wax', 'pants', 'chive', 'skin', 'wok', 'toga', 'torpedo', 'straw', 'wire', 'spoon', 'costume', 'trigger', 'rabbit', 'flashlight', 'doormat', 'horseshoe', 'mongoose', 'shortbread', 'sleeping bag', 'pine needle', 'toast', 'swizzle stick', 'snowplow', 'beanie', 'sundae', 'life jacket', 'ivy', 'ivy', 'cotton candy', 'pin', 'zebra', 'boy']

src/uot/method.py

@@ -1,7 +1,8 @@
 import copy
 
+from uot.chat_utils import renew_open_set
 from uot.models import get_response_method
-from uot.uot import select
+from uot.uot import select, renew_node_to_root
 
 
 def get_examiner_response(task, history):
@@ -35,7 +36,7 @@ def simplify_rsp(rsp):
             return n, n.question, True
 
     targeting_prompt_set = task.prompts.targeting_prompt_set_FA if task.free_answer else task.prompts.targeting_prompt_set
-    msg = [{"role": "user", "content": targeting_prompt_set.format(item_list_str=', '.join(node.items))}]
+    msg = copy.deepcopy(history) + [{"role": "user", "content": targeting_prompt_set.format(item_list_str=', '.join(node.items))}]
     return node, simplify_rsp(response(msg, model=task.guesser_model)), False
 
 
@@ -64,23 +65,37 @@ def converse(task, i):
     item = task.data[i]["target"]
     target_decl = task.prompts.target_declaration.format(target=item)
     print(target_decl)
+    print("------ DIALOGUE START ------")
+    count = 0
+
+    if not task.free_answer:
+        history_e = [{'role': 'user', 'content': task.prompts.examiner_prologue.format(item=item)}]
+    else:
+        history_e = [{'role': 'user', 'content': task.prompts.simulator_prologue.format(item=item, conv_hist=task.data[i]["conv_hist"])}]
 
     if "self_repo" in task.data[i]:
         guesser_prologue = task.prompts.guesser_prologue_FA if task.free_answer else task.prompts.guesser_prologue
         history_g = [{'role': 'user', 'content': guesser_prologue.format(repo=task.data[i]["self_repo"])}]
         print("Self-report:", task.data[i]["self_repo"])
+        node = task.root.handle_self_repo(task, task.data[i]["self_repo"])
     else:
         history_g = [{'role': 'user', 'content': task.prompts.guesser_prologue}]
-
-    if not task.free_answer:
-        history_e = [{'role': 'user', 'content': task.prompts.examiner_prologue.format(item=item)}]
-    else:
-        history_e = [{'role': 'user', 'content': task.prompts.simulator_prologue.format(item=item, conv_hist=task.data[i]["conv_hist"])}]
-
-    print("------ DIALOGUE START ------")
-    count = 0
-    node, bot1_response, flag = get_guesser_response(task, history_g, count + 1, task.root)
-    node.print()
+        # !! for openset uot !!
+        if task.open_set_size > 0 and task.n_pre_ask > 0:
+            for _ in range(task.n_pre_ask):
+                bot1_response = get_guesser_naive_response(task, history_g, count+1)
+                print("Bot 2:", bot1_response)
+                history_g.append({'role': 'system', 'content': bot1_response})
+                history_e.append({'role': 'user', 'content': bot1_response})
+                bot2_response = get_examiner_response(task, history_e)
+                print("Bot 1:", bot2_response)
+                history_g.append({'role': 'user', 'content': bot2_response})
+                history_e.append({'role': 'system', 'content': bot2_response})
+                count += 1
+                print('------', count, '-------------')
+        node = task.root.handle_self_repo(task, history_g) if task.open_set_size > 0 else task.root
+
+    node, bot1_response, flag = get_guesser_response(task, history_g, count + 1, node)
     print("Bot 2:", bot1_response)
 
     history_g.append({'role': 'system', 'content': bot1_response})
@@ -110,6 +125,10 @@ def converse(task, i):
             state = -1
             break
 
+        # renew
+        if count <= int(task.max_turn*0.3) + task.n_pre_ask and task.open_set_size > 0 and len(node.items) < task.size_to_renew:
+            node = renew_node_to_root(task, node, history_g)
+
         node, bot1_response, flag = get_guesser_response(task, history_g, count + 1, node)
         print("Bot 2:", bot1_response)
         history_g.append({'role': 'system', 'content': bot1_response})

src/uot/tasks/medical_diagnosis.py

@@ -1,4 +1,5 @@
 import os
+import json
 
 from uot.chat_utils import import_prompts_by_task
 from uot.uot import UoTNode
@@ -11,56 +12,28 @@ def __init__(self, args):
         self.max_turn = 5
         self.prompts = import_prompts_by_task("md")
         self.set = []
-        self.data = self.load_dataset(args.dataset)
+        self.data = json.loads(args.dataset)
         self.root = None
 
     def load_dataset(self, name):
         if name == "DX":
             self.set = ['Allergic rhinitis', 'upper respiratory tract infection (URTI)', 'pneumonia',
-                        'Hand foot and mouth disease in children', 'Infantile diarrhea']
-            return load_dx_dataset(os.path.join(os.path.dirname(__file__), "../data/DX_dialog.txt"))
+                        'Hand foot and mouth disease in children', 'Infantile diarrhea']\
+                if self.open_set_size <= 0 else self.set
         elif name == "MedDG":
             self.free_answer = True
             self.set = ['Enteritis', 'Gastritis', 'Gastroenteritis', 'Esophagitis',
-                        'Cholecystitis', 'Appendicitis', 'Pancreatitis', 'Gastric ulcer']
-            return load_meddg_dataset(os.path.join(os.path.dirname(__file__), "../data/MedDG_dialog.txt"))
+                        'Cholecystitis', 'Appendicitis', 'Pancreatitis', 'Gastric ulcer',
+                        'Constipation', 'Cold', 'Irritable bowel syndrome', 'Diarrhea',
+                        'Allergic rhinitis', 'Upper respiratory tract infection', 'Pneumonia']\
+                if self.open_set_size <= 0 else self.set
         else:
             raise NotImplementedError
+        return json.loads(os.path.join(os.path.dirname(__file__), f"../data/{name}.json").read())
 
     def create_root(self, root=None):
         if not root:
             self.root = UoTNode("ROOT", True, self.set, None, self.guesser_model)
         else:
-            root.n_extend_layers = self.n_extend_layers
+            root.set_config(self.n_extend_layers, not self.none_acc_reward, self.expected_reward_method)
             self.root = root
-
-
-def load_dx_dataset(file_path):
-    dic = {"过敏性鼻炎": 'Allergic rhinitis', "肺炎": 'pneumonia', "小儿腹泻": 'Infantile diarrhea',
-           "上呼吸道感染": 'upper respiratory tract infection (URTI)',
-           "小儿手足口病": 'Hand foot and mouth disease in children'}
-    with open(file_path, 'r', encoding='utf-8') as f:
-        data = eval(f.read())
-        repo_dataset = [{'self_repo': dialog['self_repo_en'], 'target': dic[dialog['disease_tag']]} for dialog in data]
-    return repo_dataset
-
-
-def load_meddg_dataset(file_path):
-    repo_dataset = []
-    flag = 0
-    disease, self_repo, dialog = "", "", ""
-    with open(file_path, 'r', encoding='utf-8') as f:
-        for line in f:
-            if line.startswith("dialog"):
-                flag = 1
-                disease = line.split("|")[1][:-1]
-            elif line.startswith("{"):
-                content = eval(line)
-                dialog += f"{content['id']}: {content['Sentence']}\n"
-                if flag:
-                    self_repo = content['self_repo_en']
-                    flag = 0
-            else:
-                repo_dataset.append({'self_repo': self_repo, 'target': disease, 'conv_hist': dialog})
-                disease, self_repo, dialog = "", "", ""
-    return repo_dataset

src/uot/tasks/prompts/medical_diagnosis.py

@@ -28,20 +28,20 @@
 '''
 
 # conversation
-target_question = "Are you a '{target}'?"
+target_question = "Are you experiencing '{target}'?"
 target_question_FA = "Are you experiencing '{target}'?"
 
 targeting_prompt_free = """Note that you should point out and ask what disease the patient suffer from now. 
-The question must be 'Are you a [disease name]?' You must reply me with 1 question to ask only."""
+The question must be 'Are you experiencing [disease name]?' You must reply me with 1 question to ask only."""
 targeting_prompt_free_FA = """Note that you should point out and ask what disease the patient suffer from now. 
 The question must be 'Are you experiencing [disease name]?' You must reply me with 1 question to ask only."""
 
 targeting_prompt_set = """Note that you should point out and ask what disease the patient suffer from now. The patient may suffer from one of diseases below:
 {item_list_str}, or other.
-The question must be 'Are you a [disease name]?' You must reply me with 1 question."""
+The question must be 'Are you experiencing [disease name]?' You must reply me with 1 question."""
 targeting_prompt_set_FA = """Note that you should point out and ask what disease the patient suffer from now. The patient may suffer from one of diseases below:
 {item_list_str}, or other.
-The question must be 'Are you a [disease name]?' You must reply me with 1 question."""
+The question must be 'Are you experiencing [disease name]?' You must reply me with 1 question."""
 
 guesser_prologue = '''You are a doctor and your patient self-reports that: {repo}.
 You should ask your patient question in English with symptoms which can only be answered by 'Yes' or 'No', in order to find what disease this patient suffers. 
@@ -52,7 +52,7 @@
 Let us begin. Ask me the first question.
 '''
 
-urge_prompt = "Based on the symptons above, if you find out the disease, please ask 'Are you a [disease name]?'"
+urge_prompt = "Based on the symptons above, if you find out the disease, please ask 'Are you experiencing [disease name]?'"
 
 inform_prompt = "The patient may suffer from one of diseases below:\n{item_list_str}"
 
@@ -86,3 +86,11 @@
 Note that never directly tell me what disease is all the time.
 Let us begin. Here is my first question.
 '''
+
+# open set
+init_open_set_prompt = '''You are a doctor and your patient self-reports that: {repo}. Please propose {size} diseases that you think your patient may suffer from.
+Your response should be: ["disease1", "disease2", ...]'''
+
+renew_open_set_prompt = '''Based on the conversation history, please propose {size} diseases that your patient may suffer from.
+The list of {size} diseases should contains {item_list}
+Your response should be: ["disease1", "disease2", ...]'''