Delete interface.py

interface.py

@@ -1,406 +0,0 @@
-import json
-import numpy as np
-import treegraph as tg
-import colorama
-from colorama import Fore
-import networkx as nx
-import utils
-import re
-
-DEBUG = True
-INPUT_COLOR = Fore.LIGHTGREEN_EX
-DEBUG_COLOR = Fore.LIGHTBLACK_EX
-OUTPUT_COLOR = Fore.LIGHTMAGENTA_EX
-INFO_COLOR = Fore.BLUE
-HELP_COLOR = Fore.CYAN
-
-
-def print_debug(*args, color=DEBUG_COLOR):
-    """
-    Prints debug messages if DEBUG is set to True.
-    """
-    if DEBUG:
-        for arg in args:
-            print(color + str(arg))
-
-
-class ReportInterface:
-    def __init__(
-        self,
-        llm: utils.LLM,
-        system_prompt: str,
-        tree_graph: nx.Graph,
-        nodes_dict: dict[str, tg.Node],
-        api_key: str = None,
-    ):
-        self.llm = llm
-        self.system_prompt = system_prompt
-        self.tree_graph = tree_graph
-        self.nodes_dict = nodes_dict
-        self.api_key = api_key
-        self.build()
-
-    def build(self):
-        utils.set_api_key(self.api_key)
-        self.system_prompt = utils.make_message("system", self.system_prompt)
-        self.visitable_nodes = self._get_visitable_nodes()
-        self.report_dict = self._get_report_dict()
-
-        self.active_node: tg.Node = self.nodes_dict["root"]
-        self.unique_visited_nodes = set()  # set of nodes visited
-        self.node_journey = []  # list of nodes visited
-        self.distance_travelled = 0  # number of edges travelled
-        self.jumps = 0  # number of jumps
-        self.jump_lengths = []  # list of jump lengths
-        self.counter = 0  # number of questions asked
-
-        colorama.init(autoreset=True)  # to reset the color after each print statement
-
-        self.help_message = f"""You are presented with a Knee MRI.
-        You are asked to fill out a radiology report.
-        Please only report the findings in the MRI.
-        Please mention your findings with the corresponding anatomical structures.
-        There are {len(self.visitable_nodes.keys())} visitable nodes in the tree.
-        You must visit as many nodes as possible, while avoiding too many jumps."""
-
-    def _get_visitable_nodes(self):
-        return dict(
-            zip(
-                [
-                    node.name
-                    for node in self.tree_graph.nodes
-                    if node.name != "root" and node.has_children() is False
-                ],
-                [
-                    node
-                    for node in self.tree_graph.nodes
-                    if node.name != "root" and node.has_children() is False
-                ],
-            )
-        )
-
-    def _get_report_dict(self):
-        return {
-            node.name: tg.Node(node.name, "", node.children)
-            for node in self.visitable_nodes.values()
-        }
-
-    @utils.debug(DEBUG, print_debug)
-    def _check_question_validity(
-        self,
-        question: str,
-    ):
-        # let's ask the question from the model and check if it's valid
-        template_json = json.dumps(
-            {key: node.value for key, node in self.visitable_nodes.items()},
-            indent=4,
-        )
-        q = f"""the following is a Knee MRI report "template" in a JSON format with keys and values.
-        You are given a "finding" phrase from a radiologist.
-        Match as best as possible the "finding" with one of keys in the "template".
-        <template>
-        {template_json}
-        </template>
-        <finding>
-        {question}
-        </finding>
-        "available": [Is the "finding" relevant to any key in the "template"? say "yes" or "no".
-        Make sure the "finding" is relevant to Knee MRI and knee anatomy otherwise say 'no'.
-        Do not answer irrelevant phrases.]
-        "node": [if the above answer is 'yes', write only the KEY of the most relevant node to the "finding". otherwise, say 'none'.]
-        """
-
-        keys = ["available", "node"]
-        prompt = [self.system_prompt] + [
-            utils.make_question(utils.JSON_TEMPLATE, question=q, keys=keys)
-        ]
-        response = self.llm(prompt)
-        print_debug(
-            prompt,
-            response,
-        )
-        available = utils.json2dict(response)["available"].strip().lower()
-        node = utils.json2dict(response)["node"]
-        return available, node
-
-    def _update_node(self, node_name, findings):
-        self.report_dict[node_name].value += str(findings) + "\n"
-        response = f"Updated node '{node_name}' with finding '{findings}'"
-        print(OUTPUT_COLOR + response)
-        return response
-
-    def save_report(self, filename: str):
-        # convert performance metrics to json
-        metrics = {
-            "distance_travelled": self.distance_travelled,
-            "jumps": self.jumps,
-            "jump_lengths": self.jump_lengths,
-            "unique_visited_nodes": [node.name for node in self.unique_visited_nodes],
-            "node_journey": [node.name for node in self.node_journey],
-            "report": {
-                node_name: node.value for node_name, node in self.report_dict.items()
-            },
-        }
-        # save the report
-        with open(filename, "w") as file:
-            json.dump(metrics, file, indent=4)
-
-    def prime_model(self):
-        """
-        Primes the model with the system prompt.
-        """
-        q = "Are you ready to begin?\nSay 'yes' or 'no'."
-        keys = ["answer"]
-        response = self.llm(
-            [
-                self.system_prompt,
-                utils.make_question(utils.JSON_TEMPLATE, question=q, keys=keys),
-            ],
-        )
-        print_debug(q, response)
-        if utils.json2dict(response)["answer"].lower() == "yes":
-            print(INFO_COLOR + "The model is ready.")
-            return True
-        else:
-            print(INFO_COLOR + "The model is not ready.")
-            return False
-
-    def performance_summary(self):
-        # print out the summary info
-        print(INFO_COLOR + "Performance Summary:")
-        print(
-            INFO_COLOR + f"Total distance travelled: {self.distance_travelled} edge(s)"
-        )
-        print(INFO_COLOR + f"Jump lengths: {self.jump_lengths}")
-        print(INFO_COLOR + f"Jump lengths mean: {np.mean(self.jump_lengths):.1f}")
-        print(INFO_COLOR + f"Jump lengths SD: {np.std(self.jump_lengths):.1f}")
-        print(INFO_COLOR + f"Nodes visited in order: {self.node_journey}")
-        print(INFO_COLOR + f"Unique nodes visited: {self.unique_visited_nodes}")
-        print(
-            INFO_COLOR
-            + 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."
-        )
-        print_debug("\n")
-        print_debug("Report Summary:".rjust(20))
-        for name, node in self.report_dict.items():
-            if node.value != "":
-                print_debug(f"{name}: {node.value}")
-        print(INFO_COLOR + f"total cost: ${self.llm.cost:.4f}")
-        print(INFO_COLOR + f"total tokens used: {self.llm.token_counter}")
-
-    def get_stats(self):
-        report_string = ""
-        for name, node in self.report_dict.items():
-            if node.value != "":
-                report_string += f"{name}: <{node.value}> \n"
-        return {
-            "Lengths travelled": self.distance_travelled,
-            "Number of jumps": self.jumps,
-            "Jump lengths": self.jump_lengths,
-            "Unique nodes visited": [node.name for node in self.unique_visited_nodes],
-            "Visited Nodes": [node.name for node in self.node_journey],
-            "Report": report_string,
-        }
-
-    def visualize_tree(self, **kwargs):
-        tg.visualize_graph(tg.from_list(self.node_journey), self.tree_graph, **kwargs)
-
-    def get_plot(self, **kwargs):
-        return tg.get_graph(tg.from_list(self.node_journey), self.tree_graph, **kwargs)
-
-    def process_input(self, input_text: str):
-        res = "n/a"
-        try:
-            finding = input_text
-            if finding.strip().lower() == "quit":
-                print(INFO_COLOR + "Exiting...")
-                return "quit"
-            elif finding.strip().lower() == "help":
-                return "help"
-
-            available, node = self._check_question_validity(finding)
-            if available != "yes":
-                print(
-                    OUTPUT_COLOR
-                    + "Could not find a relevant node.\nWrite more clearly and provide more details."
-                )
-                return "n/a"
-            if node not in self.visitable_nodes.keys():
-                print(
-                    OUTPUT_COLOR
-                    + "Could not find a relevant node.\nWrite more clearly and provide more details."
-                )
-                return "n/a"
-            else:
-                # modify the tree to update the node with findings
-                res = self._update_node(node, finding)
-
-                print(
-                    INFO_COLOR
-                    + f"jumping from node '{self.active_node}' to node '{node}'..."
-                )
-                distance = tg.num_edges_between_nodes(
-                    self.tree_graph, self.active_node, self.nodes_dict[node]
-                )
-                print(INFO_COLOR + f"distance travelled: {distance} edge(s)")
-
-                self.active_node = self.nodes_dict[node]
-                self.jumps += 1
-                self.jump_lengths.append(distance)
-                self.distance_travelled += distance
-                if self.active_node.name != "root":
-                    self.unique_visited_nodes.add(self.active_node)
-                    self.node_journey.append(self.active_node)
-        except Exception as ex:
-            print_debug(ex, color=Fore.LIGHTRED_EX)
-            return "exception"
-
-        self.counter += 1
-        try:
-            self.performance_summary()
-        except Exception as ex:
-            print_debug(ex, color=Fore.LIGHTRED_EX)
-        return res
-
-
-class ReportChecklistInterface:
-    def __init__(
-        self,
-        llm: utils.LLM,
-        system_prompt: str,
-        graph: nx.Graph,
-        nodes_dict: dict[str, tg.Node],
-        api_key: str = None,
-    ):
-        self.llm = llm
-        self.system_prompt = system_prompt
-        self.tree_graph: nx.Graph = graph
-        self.nodes_dict = nodes_dict
-        self.api_key = api_key
-        self.build()
-
-    def build(self):
-        utils.set_api_key(self.api_key)
-        self.system_prompt = utils.make_message("system", self.system_prompt)
-        self.visitable_nodes = self._get_visitable_nodes()
-
-        colorama.init(autoreset=True)  # to reset the color after each print statement
-
-        self.help_message = f"""You are presented with a Knee MRI.
-        You are asked to fill out a radiology report.
-        Please only report the findings in the MRI.
-        Please mention your findings with the corresponding anatomical structures.
-        There are {len(self.visitable_nodes.keys())} visitable nodes in the tree."""
-
-    def _get_visitable_nodes(self):
-        return dict(
-            zip(
-                [
-                    node.name
-                    for node in self.tree_graph.nodes
-                    if node.name != "root" and node.has_children() is False
-                ],
-                [
-                    node
-                    for node in self.tree_graph.nodes
-                    if node.name != "root" and node.has_children() is False
-                ],
-            )
-        )
-
-    @utils.debug(DEBUG, print_debug)
-    def _check_report(
-        self,
-        report: str,
-    ):
-        # let's ask the question from the model and check if it's valid
-        checklist_json = json.dumps(
-            {key: node.value for key, node in self.visitable_nodes.items()},
-            indent=4,
-        )
-        q = f"""the following is a Knee MRI "checklist" in JSON format with keys as items and values as findings:
-        A knee MRI "report" is also provided in raw text format written by a radiologist:
-        <checklist>
-        {checklist_json}
-        </checklist>
-        <report>
-        {report}
-        </report>
-        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".
-        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".
-        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.
-        Try to fill out as many items as possible.
-        ALL of the items in the "checklist" must be filled out.
-        Don't generate findings that are not present in the "report" (new findings).
-        Be comprehensive and don't miss any findings that are present in the "report".
-        Watch out for encompassing terms (e.g., "cruciate ligaments" means both "ACL" and "PCL").
-        "thought_process": [Think in steps on how you would do this task.]
-        "report_ckecked" : [a JSON with the same keys as the "checklist" but take the values from the "report", as described above.]
-        """
-
-        keys = ["thought_process", "report_checked"]
-        prompt = [self.system_prompt] + [
-            utils.make_question(utils.JSON_TEMPLATE, question=q, keys=keys)
-        ]
-        response = self.llm(prompt)
-        print_debug(
-            prompt,
-            response,
-        )
-        report_checked = utils.json2dict(response)
-        return report_checked["report_checked"]
-
-    def prime_model(self):
-        """
-        Primes the model with the system prompt.
-        """
-        q = "Are you ready to begin?\nSay 'yes' or 'no'."
-        keys = ["answer"]
-        response = self.llm(
-            [
-                self.system_prompt,
-                utils.make_question(utils.JSON_TEMPLATE, question=q, keys=keys),
-            ],
-        )
-        print_debug(q, response)
-        if utils.json2dict(response)["answer"].lower() == "yes":
-            print(INFO_COLOR + "The model is ready.")
-            return True
-        else:
-            print(INFO_COLOR + "The model is not ready.")
-            return False
-
-    def process_input(self, input_text: str):
-        try:
-            report = input_text
-            if report.strip().lower() == "quit":
-                print(INFO_COLOR + "Exiting...")
-                return "quit"
-            elif report.strip().lower() == "help":
-                return "help"
-
-            checked_report: dict = self._check_report(report)
-            # make a string of the report
-            # replace true with [checkmark emoji] and false with [cross emoji]
-            report_string = ""
-            CHECKMARK = "\u2705"
-            CROSS = "\u274C"
-
-            # we need a regex to convert the camelCase keys to a readable format
-            def camel2readable(camel: str):
-                string = re.sub("([a-z])([A-Z])", r"\1 \2", camel)
-                # captialize every word
-                string = " ".join([word.capitalize() for word in string.split()])
-                return string
-
-            for key, value in checked_report.items():
-                if str(value).lower() == "true":
-                    report_string += f"{camel2readable(key)}: {CHECKMARK}\n"
-                elif str(value).lower() == "n/a":
-                    report_string += f"{camel2readable(key)}: {CROSS}\n"
-                else:
-                    report_string += f"{camel2readable(key)}: <{value}> {CHECKMARK}\n"
-            return report_string
-        except Exception as ex:
-            print_debug(ex, color=Fore.LIGHTRED_EX)
-            return "exception"