Mahbodez/knee_report_checklist
1
1import json2import numpy as np3import treegraph as tg4import colorama5from colorama import Fore6import networkx as nx7import utils8import re9import logger as lg10 11DEBUG = True12INPUT_COLOR = Fore.LIGHTGREEN_EX13DEBUG_COLOR = Fore.LIGHTBLACK_EX14OUTPUT_COLOR = Fore.LIGHTMAGENTA_EX15INFO_COLOR = Fore.BLUE16HELP_COLOR = Fore.CYAN17 18 19def print_debug(*args, color=DEBUG_COLOR):20 """21 Prints debug messages if DEBUG is set to True.22 """23 if DEBUG:24 for arg in args:25 print(color + str(arg))26 27 28class ReportInterface:29 def __init__(30 self,31 llm: utils.LLM,32 system_prompt: str,33 tree_graph: nx.Graph,34 nodes_dict: dict[str, tg.Node],35 api_key: str = None,36 ):37 self.llm = llm38 self.system_prompt = system_prompt39 self.tree_graph = tree_graph40 self.nodes_dict = nodes_dict41 self.api_key = api_key42 self.build()43 44 def build(self):45 utils.set_api_key(self.api_key)46 self.system_prompt = utils.make_message("system", self.system_prompt)47 self.visitable_nodes = self._get_visitable_nodes()48 self.report_dict = self._get_report_dict()49 50 self.active_node: tg.Node = self.nodes_dict["root"]51 self.unique_visited_nodes = set() # set of nodes visited52 self.node_journey = [] # list of nodes visited53 self.distance_travelled = 0 # number of edges travelled54 self.jumps = 0 # number of jumps55 self.jump_lengths = [] # list of jump lengths56 self.counter = 0 # number of questions asked57 58 colorama.init(autoreset=True) # to reset the color after each print statement59 60 self.help_message = f"""You are presented with a Knee MRI.61 You are asked to fill out a radiology report.62 Please only report the findings in the MRI.63 Please mention your findings with the corresponding anatomical structures.64 There are {len(self.visitable_nodes.keys())} visitable nodes in the tree.65 You must visit as many nodes as possible, while avoiding too many jumps."""66 67 def _get_visitable_nodes(self):68 return dict(69 zip(70 [71 node.name72 for node in self.tree_graph.nodes73 if node.name != "root" and node.has_children() is False74 ],75 [76 node77 for node in self.tree_graph.nodes78 if node.name != "root" and node.has_children() is False79 ],80 )81 )82 83 def _get_report_dict(self):84 return {85 node.name: tg.Node(node.name, "", node.children)86 for node in self.visitable_nodes.values()87 }88 89 @utils.debug(DEBUG, print_debug)90 def _check_question_validity(91 self,92 question: str,93 ):94 # let's ask the question from the model and check if it's valid95 template_json = json.dumps(96 {key: node.value for key, node in self.visitable_nodes.items()},97 indent=4,98 )99 q = f"""the following is a Knee MRI report "template" in a JSON format with keys and values.100 You are given a "finding" phrase from a radiologist.101 Match as best as possible the "finding" with one of keys in the "template".102 <template>103 {template_json}104 </template>105 <finding>106 {question}107 </finding>108 "available": [Is the "finding" relevant to any key in the "template"? say "yes" or "no".109 Make sure the "finding" is relevant to Knee MRI and knee anatomy otherwise say 'no'.110 Do not answer irrelevant phrases.]111 "node": [if the above answer is 'yes', write only the KEY of the most relevant node to the "finding". otherwise, say 'none'.]112 """113 114 keys = ["available", "node"]115 prompt = [self.system_prompt] + [116 utils.make_question(utils.JSON_TEMPLATE, question=q, keys=keys)117 ]118 response = self.llm(prompt)119 print_debug(120 prompt,121 response,122 )123 available = utils.json2dict(response)["available"].strip().lower()124 node = utils.json2dict(response)["node"]125 return available, node126 127 def _update_node(self, node_name, findings):128 self.report_dict[node_name].value += str(findings) + "\n"129 response = f"Updated node '{node_name}' with finding '{findings}'"130 print(OUTPUT_COLOR + response)131 return response132 133 def save_report(self, filename: str):134 # convert performance metrics to json135 metrics = {136 "distance_travelled": self.distance_travelled,137 "jumps": self.jumps,138 "jump_lengths": self.jump_lengths,139 "unique_visited_nodes": [node.name for node in self.unique_visited_nodes],140 "node_journey": [node.name for node in self.node_journey],141 "report": {142 node_name: node.value for node_name, node in self.report_dict.items()143 },144 }145 # save the report146 with open(filename, "w") as file:147 json.dump(metrics, file, indent=4)148 149 def prime_model(self):150 """151 Primes the model with the system prompt.152 """153 q = "Are you ready to begin?\nSay 'yes' or 'no'."154 keys = ["answer"]155 response = self.llm(156 [157 self.system_prompt,158 utils.make_question(utils.JSON_TEMPLATE, question=q, keys=keys),159 ],160 )161 print_debug(q, response)162 if utils.json2dict(response)["answer"].lower() == "yes":163 print(INFO_COLOR + "The model is ready.")164 return True165 else:166 print(INFO_COLOR + "The model is not ready.")167 return False168 169 def performance_summary(self):170 # print out the summary info171 print(INFO_COLOR + "Performance Summary:")172 print(173 INFO_COLOR + f"Total distance travelled: {self.distance_travelled} edge(s)"174 )175 print(INFO_COLOR + f"Jump lengths: {self.jump_lengths}")176 print(INFO_COLOR + f"Jump lengths mean: {np.mean(self.jump_lengths):.1f}")177 print(INFO_COLOR + f"Jump lengths SD: {np.std(self.jump_lengths):.1f}")178 print(INFO_COLOR + f"Nodes visited in order: {self.node_journey}")179 print(INFO_COLOR + f"Unique nodes visited: {self.unique_visited_nodes}")180 print(181 INFO_COLOR182 + f"You have explored {len(self.unique_visited_nodes)/len(self.visitable_nodes):.1%} ({len(self.unique_visited_nodes)}/{len(self.visitable_nodes)}) of the tree."183 )184 print_debug("\n")185 print_debug("Report Summary:".rjust(20))186 for name, node in self.report_dict.items():187 if node.value != "":188 print_debug(f"{name}: {node.value}")189 print(INFO_COLOR + f"total cost: ${self.llm.cost:.4f}")190 print(INFO_COLOR + f"total tokens used: {self.llm.token_counter}")191 192 def get_stats(self):193 report_string = ""194 for name, node in self.report_dict.items():195 if node.value != "":196 report_string += f"{name}: <{node.value}> \n"197 return {198 "Lengths travelled": self.distance_travelled,199 "Number of jumps": self.jumps,200 "Jump lengths": self.jump_lengths,201 "Unique nodes visited": [node.name for node in self.unique_visited_nodes],202 "Visited Nodes": [node.name for node in self.node_journey],203 "Report": report_string,204 }205 206 def visualize_tree(self, **kwargs):207 tg.visualize_graph(tg.from_list(self.node_journey), self.tree_graph, **kwargs)208 209 def get_plot(self, **kwargs):210 return tg.get_graph(tg.from_list(self.node_journey), self.tree_graph, **kwargs)211 212 def process_input(self, input_text: str):213 res = "n/a"214 try:215 finding = input_text216 if finding.strip().lower() == "quit":217 print(INFO_COLOR + "Exiting...")218 return "quit"219 elif finding.strip().lower() == "help":220 return "help"221 222 available, node = self._check_question_validity(finding)223 if available != "yes":224 print(225 OUTPUT_COLOR226 + "Could not find a relevant node.\nWrite more clearly and provide more details."227 )228 return "n/a"229 if node not in self.visitable_nodes.keys():230 print(231 OUTPUT_COLOR232 + "Could not find a relevant node.\nWrite more clearly and provide more details."233 )234 return "n/a"235 else:236 # modify the tree to update the node with findings237 res = self._update_node(node, finding)238 239 print(240 INFO_COLOR241 + f"jumping from node '{self.active_node}' to node '{node}'..."242 )243 distance = tg.num_edges_between_nodes(244 self.tree_graph, self.active_node, self.nodes_dict[node]245 )246 print(INFO_COLOR + f"distance travelled: {distance} edge(s)")247 248 self.active_node = self.nodes_dict[node]249 self.jumps += 1250 self.jump_lengths.append(distance)251 self.distance_travelled += distance252 if self.active_node.name != "root":253 self.unique_visited_nodes.add(self.active_node)254 self.node_journey.append(self.active_node)255 except Exception as ex:256 print_debug(ex, color=Fore.LIGHTRED_EX)257 return "exception"258 259 self.counter += 1260 try:261 self.performance_summary()262 except Exception as ex:263 print_debug(ex, color=Fore.LIGHTRED_EX)264 return res265 266 267class ReportChecklistInterface:268 def __init__(269 self,270 llm: utils.LLM,271 system_prompt: str,272 graph: nx.Graph,273 nodes_dict: dict[str, tg.Node],274 api_key: str = None,275 logger: lg.Logger = None,276 username: str = None,277 ):278 self.llm = llm279 self.system_prompt = system_prompt280 self.tree_graph: nx.Graph = graph281 self.nodes_dict = nodes_dict282 self.api_key = api_key283 self.logger = logger284 self.username = username285 self.build()286 287 def build(self):288 utils.set_api_key(self.api_key)289 self.system_prompt = utils.make_message("system", self.system_prompt)290 self.visitable_nodes = self._get_visitable_nodes()291 292 colorama.init(autoreset=True) # to reset the color after each print statement293 294 self.help_message = f"""You are presented with a Knee MRI.295 You are asked to fill out a radiology report.296 Please only report the findings in the MRI.297 Please mention your findings with the corresponding anatomical structures.298 There are {len(self.visitable_nodes.keys())} visitable nodes in the tree."""299 300 def _get_visitable_nodes(self):301 return dict(302 zip(303 [304 node.name305 for node in self.tree_graph.nodes306 if node.name != "root" and node.has_children() is False307 ],308 [309 node310 for node in self.tree_graph.nodes311 if node.name != "root" and node.has_children() is False312 ],313 )314 )315 316 @utils.debug(DEBUG, print_debug)317 def _check_report(318 self,319 report: str,320 ):321 # let's ask the question from the model and check if it's valid322 checklist_json = json.dumps(323 {key: node.value for key, node in self.visitable_nodes.items()},324 indent=4,325 )326 q = f"""the following is a Knee MRI "checklist" in JSON format with keys as items and values as findings:327 A knee MRI "report" is also provided in raw text format written by a radiologist:328 <checklist>329 {checklist_json}330 </checklist>331 <report>332 {report}333 </report>334 Your task is to find all the corresponding items from the "checklist" in the "report" and fill out a JSON with the same keys as the "checklist" but extract the corresponding values from the "report".335 If a key is not found in the "report", please set the value to "n/a", otherwise set it to the corresponding finding from the "report".336 You must check the "report" phrases one by one and find a corresponding key(s) for EACH phrase in the "report" from the "checklist" and fill out the "report_checked" JSON.337 Try to fill out as many items as possible.338 ALL of the items in the "checklist" must be filled out.339 Don't generate findings that are not present in the "report" (new findings).340 Be comprehensive and don't miss any findings that are present in the "report".341 Watch out for encompassing terms (e.g., "cruciate ligaments" means both "ACL" and "PCL").342 "thought_process": [Think in steps on how you would do this task.]343 "report_ckecked" : [a JSON with the same keys as the "checklist" but take the values from the "report", as described above.]344 """345 346 keys = ["thought_process", "report_checked"]347 prompt = [self.system_prompt] + [348 utils.make_question(utils.JSON_TEMPLATE, question=q, keys=keys)349 ]350 response = self.llm(prompt)351 print_debug(352 prompt,353 response,354 )355 if self.logger:356 # set name to class name357 self.logger(358 name=self.__class__.__name__,359 message=f"prompt: {prompt}\nresponse: {response}",360 )361 report_checked = utils.json2dict(response)362 return report_checked["report_checked"]363 364 def prime_model(self):365 """366 Primes the model with the system prompt.367 """368 q = "Are you ready to begin?\nSay 'yes' or 'no'."369 keys = ["answer"]370 response = self.llm(371 [372 self.system_prompt,373 utils.make_question(utils.JSON_TEMPLATE, question=q, keys=keys),374 ],375 )376 print_debug(q, response)377 if utils.json2dict(response)["answer"].lower() == "yes":378 print(INFO_COLOR + "The model is ready.")379 return True380 else:381 print(INFO_COLOR + "The model is not ready.")382 return False383 384 def process_input(self, input_text: str):385 try:386 report = input_text387 if self.logger:388 self.logger(self.username, f"report: {report}")389 390 if report.strip().lower() == "quit":391 print(INFO_COLOR + "Exiting...")392 if self.logger:393 self.logger(self.username, "Exiting...")394 return "quit"395 elif report.strip().lower() == "help":396 if self.logger:397 self.logger(self.username, "Help")398 return "help"399 400 checked_report: dict = self._check_report(report)401 # make a string of the report402 # replace true with [checkmark emoji] and false with [cross emoji]403 report_string = ""404 CHECKMARK = "\u2705"405 CROSS = "\u274C"406 407 # we need a regex to convert the camelCase keys to a readable format408 def camel2readable(camel: str):409 string = re.sub("([a-z])([A-Z])", r"\1 \2", camel)410 # captialize every word411 string = " ".join([word.capitalize() for word in string.split()])412 return string413 414 for key, value in checked_report.items():415 if str(value).lower() == "n/a":416 report_string += f"{camel2readable(key)}: {CROSS}\n"417 else:418 report_string += f"{camel2readable(key)}: <{value}> {CHECKMARK}\n"419 420 portion_visited: float = report_string.count(CHECKMARK) / len(421 checked_report.keys()422 )423 report_string += f"Portion of the checklist visited: {portion_visited:.1%}"424 if self.logger:425 self.logger(self.__class__.__name__, report_string)426 return report_string427 except Exception as ex:428 print_debug(ex, color=Fore.LIGHTRED_EX)429 if self.logger:430 self.logger(self.__class__.__name__, "Exception: " + ex)431 return "exception"432