import os
import json
import random
import numpy as np
import argparse
import autogen
from toolset_high import *
from medagent import MedAgent
from config import openai_config, llm_config_list
import time
def judge(pred, ans):
old_flag = True
if not ans in pred:
old_flag = False
if "True" in pred:
pred = pred.replace("True", "1")
else:
pred = pred.replace("False", "0")
if ans == "False" or ans == "false":
ans = "0"
if ans == "True" or ans == "true":
ans = "1"
if ans == "No" or ans == "no":
ans = "0"
if ans == "Yes" or ans == "yes":
ans = "1"
if ans == "None" or ans == "none":
ans = "0"
if ", " in ans:
ans = ans.split(', ')
if ans[-2:] == ".0":
ans = ans[:-2]
if not type(ans) == list:
ans = [ans]
new_flag = True
for i in range(len(ans)):
if not ans[i] in pred:
new_flag = False
break
return (old_flag or new_flag)
def set_seed(seed):
random.seed(seed)
np.random.seed(seed)
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--llm", type=str, default="<YOUR_LLM_NAME>")
parser.add_argument("--num_questions", type=int, default=1)
parser.add_argument("--dataset", type=str, default="mimic_iii")
parser.add_argument("--data_path", type=str, default="<YOUR_DATASET_PATH>")
parser.add_argument("--logs_path", type=str, default="<YOUR_LOGS_PATH>")
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--debug", action="store_true")
parser.add_argument("--debug_id", type=str, default="521fd2885f51641a963f8d3e")
parser.add_argument("--start_id", type=int, default=0)
parser.add_argument("--num_shots", type=int, default=4)
args = parser.parse_args()
set_seed(args.seed)
if args.dataset == 'mimic_iii':
from prompts_mimic import EHRAgent_4Shots_Knowledge
else:
from prompts_eicu import EHRAgent_4Shots_Knowledge
config_list = [openai_config(args.llm)]
llm_config = llm_config_list(args.seed, config_list)
chatbot = autogen.agentchat.AssistantAgent(
name="chatbot",
system_message="For coding tasks, only use the functions you have been provided with. Reply TERMINATE when the task is done. Save the answers to the questions in the variable 'answer'. Please only generate the code.",
llm_config=llm_config,
)
user_proxy = MedAgent(
name="user_proxy",
is_termination_msg=lambda x: x.get("content", "") and x.get("content", "").rstrip().endswith("TERMINATE"),
human_input_mode="NEVER",
max_consecutive_auto_reply=10,
code_execution_config={"work_dir": "coding"},
config_list=config_list,
)
# register the functions
user_proxy.register_function(
function_map={
"python": run_code
}
)
user_proxy.register_dataset(args.dataset)
file_path = args.data_path
# read from json file
with open(file_path, 'r') as f:
contents = json.load(f)
# random shuffle
import random
random.shuffle(contents)
file_path = "{}/{}/".format(args.logs_path, args.num_shots) + "{id}.txt"
start_time = time.time()
if args.num_questions == -1:
args.num_questions = len(contents)
long_term_memory = []
init_memory = EHRAgent_4Shots_Knowledge
init_memory = init_memory.split('\n\n')
for i in range(len(init_memory)):
item = init_memory[i]
item = item.split('Question:')[-1]
question = item.split('\nKnowledge:\n')[0]
item = item.split('\nKnowledge:\n')[-1]
knowledge = item.split('\nSolution:')[0]
code = item.split('\nSolution:')[-1]
new_item = {"question": question, "knowledge": knowledge, "code": code}
long_term_memory.append(new_item)
for i in range(args.start_id, args.num_questions):
if args.debug and contents[i]['id'] != args.debug_id:
continue
question = contents[i]['template']
answer = contents[i]['answer']
try:
user_proxy.update_memory(args.num_shots, long_term_memory)
user_proxy.initiate_chat(
chatbot,
message=question,
)
logs = user_proxy._oai_messages
logs_string = []
logs_string.append(str(question))
logs_string.append(str(answer))
for agent in list(logs.keys()):
for j in range(len(logs[agent])):
if logs[agent][j]['content'] != None:
logs_string.append(logs[agent][j]['content'])
else:
argums = logs[agent][j]['function_call']['arguments']
if type(argums) == dict and 'cell' in argums.keys():
logs_string.append(argums['cell'])
else:
logs_string.append(argums)
except Exception as e:
logs_string = [str(e)]
print(logs_string)
file_directory = file_path.format(id=contents[i]['id'])
# f = open(file_directory, 'w')
if type(answer) == list:
answer = ', '.join(answer)
logs_string.append("Ground-Truth Answer ---> "+answer)
with open(file_directory, 'w') as f:
f.write('\n----------------------------------------------------------\n'.join(logs_string))
logs_string = '\n----------------------------------------------------------\n'.join(logs_string)
if '"cell": "' in logs_string:
last_code_start = logs_string.rfind('"cell": "')
last_code_end = logs_string.rfind('"\n}')
last_code = logs_string[last_code_start+9:last_code_end]
else:
last_code_end = logs_string.rfind('Solution:')
prediction_end = logs_string.rfind('TERMINATE')
prediction = logs_string[last_code_end:prediction_end]
result = judge(prediction, answer)
if result:
new_item = {"question": question, "knowledge": user_proxy.knowledge, "code": user_proxy.code}
long_term_memory.append(new_item)
end_time = time.time()
print("Time elapsed: ", end_time - start_time)
if __name__ == "__main__":
main()
