Upload 5 files

interface.py

@@ -0,0 +1,431 @@
+import json
+import numpy as np
+import treegraph as tg
+import colorama
+from colorama import Fore
+import networkx as nx
+import utils
+import re
+import logger as lg
+
+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,
+        logger: lg.Logger = None,
+        username: 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.logger = logger
+        self.username = username
+        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,
+        )
+        if self.logger:
+            # set name to class name
+            self.logger(
+                name=self.__class__.__name__,
+                message=f"prompt: {prompt}\nresponse: {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 self.logger:
+                self.logger(self.username, f"report: {report}")
+
+            if report.strip().lower() == "quit":
+                print(INFO_COLOR + "Exiting...")
+                if self.logger:
+                    self.logger(self.username, "Exiting...")
+                return "quit"
+            elif report.strip().lower() == "help":
+                if self.logger:
+                    self.logger(self.username, "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() == "n/a":
+                    report_string += f"{camel2readable(key)}: {CROSS}\n"
+                else:
+                    report_string += f"{camel2readable(key)}: <{value}> {CHECKMARK}\n"
+
+            portion_visited: float = report_string.count(CHECKMARK) / len(
+                checked_report.keys()
+            )
+            report_string += f"Portion of the checklist visited: {portion_visited:.1%}"
+            if self.logger:
+                self.logger(self.__class__.__name__, report_string)
+            return report_string
+        except Exception as ex:
+            print_debug(ex, color=Fore.LIGHTRED_EX)
+            if self.logger:
+                self.logger(self.__class__.__name__, "Exception: " + ex)
+            return "exception"

knee_template.json

@@ -0,0 +1,353 @@
+{
+    "root": {
+        "value": "root",
+        "parent": null,
+        "children": [
+            "kneeJointEffusion",
+            "kneeMeniscus",
+            "kneeAclPcl",
+            "kneeMcl",
+            "kneePosterolateralCorner",
+            "kneeExtensorMechanism",
+            "kneeCartilage",
+            "kneeBone",
+            "kneeOther"
+        ]
+    },
+    "kneeJointEffusion": {
+        "value": "Presence and/or extent of joint effusion.",
+        "parent": "root",
+        "children": []
+    },
+    "kneeMeniscus": {
+        "value": "",
+        "parent": "root",
+        "children": [
+            "kneeMeniscusMedialTearing",
+            "kneeMeniscusLateralTearing",
+            "kneeMeniscusWrisberg",
+            "kneeMeniscusRootTearing",
+            "kneeMeniscusRampLesion"
+        ]
+    },
+    "kneeMeniscusMedialTearing": {
+        "value": "Presence and/or severity of medial meniscus tearing",
+        "parent": "kneeMeniscus",
+        "children": []
+    },
+    "kneeMeniscusLateralTearing": {
+        "value": "Presence and/or severity of lateral meniscus tearing",
+        "parent": "kneeMeniscus",
+        "children": []
+    },
+    "kneeMeniscusWrisberg": {
+        "value": "Presence and/or severity of Wrisberg variant",
+        "parent": "kneeMeniscus",
+        "children": []
+    },
+    "kneeMeniscusRootTearing": {
+        "value": "Presence and/or severity of meniscus root tearing",
+        "parent": "kneeMeniscus",
+        "children": []
+    },
+    "kneeMeniscusRampLesion": {
+        "value": "Presence and/or severity of ramp lesion",
+        "parent": "kneeMeniscus",
+        "children": []
+    },
+    "kneeAclPcl": {
+        "value": "",
+        "parent": "root",
+        "children": [
+            "kneeAcl",
+            "kneePcl"
+        ]
+    },
+    "kneeAcl": {
+        "value": "",
+        "parent": "kneeAclPcl",
+        "children": [
+            "kneeAclTearing",
+            "kneeAclDegeneration",
+            "kneeAclReconstruction"
+        ]
+    },
+    "kneeAclTearing": {
+        "value": "Presence and/or severity of ACL tearing",
+        "parent": "kneeAcl",
+        "children": []
+    },
+    "kneeAclDegeneration": {
+        "value": "Presence and/or severity of ACL degeneration",
+        "parent": "kneeAcl",
+        "children": []
+    },
+    "kneeAclReconstruction": {
+        "value": "ACL reconstruction status",
+        "parent": "kneeAcl",
+        "children": []
+    },
+    "kneePcl": {
+        "value": "",
+        "parent": "kneeAclPcl",
+        "children": [
+            "kneePclTearing",
+            "kneePclDegeneration",
+            "kneePclReconstruction"
+        ]
+    },
+    "kneePclTearing": {
+        "value": "Presence and/or severity of PCL tearing",
+        "parent": "kneePcl",
+        "children": []
+    },
+    "kneePclDegeneration": {
+        "value": "Presence and/or severity of PCL degeneration",
+        "parent": "kneePcl",
+        "children": []
+    },
+    "kneePclReconstruction": {
+        "value": "PCL reconstruction status",
+        "parent": "kneePcl",
+        "children": []
+    },
+    "kneeMcl": {
+        "value": "",
+        "parent": "root",
+        "children": [
+            "kneeMclTearing",
+            "kneeMclDeepFibers",
+            "kneeMclSuperficialFibers"
+        ]
+    },
+    "kneeMclTearing": {
+        "value": "Presence and/or severity of MCL tearing",
+        "parent": "kneeMcl",
+        "children": []
+    },
+    "kneeMclDeepFibers": {
+        "value": "MCL deep fibers status",
+        "parent": "kneeMcl",
+        "children": []
+    },
+    "kneeMclSuperficialFibers": {
+        "value": "MCL superficial fibers status",
+        "parent": "kneeMcl",
+        "children": []
+    },
+    "kneePosterolateralCorner": {
+        "value": "",
+        "parent": "root",
+        "children": [
+            "kneeIlioTibialBand",
+            "kneeBicepsFemorisTendon",
+            "kneeLateralCollateralLigament"
+        ]
+    },
+    "kneeIlioTibialBand": {
+        "value": "Presence and/or severity of ilio-tibial band findings",
+        "parent": "kneePosterolateralCorner",
+        "children": []
+    },
+    "kneeBicepsFemorisTendon": {
+        "value": "Presence and/or severity of biceps femoris tendon findings",
+        "parent": "kneePosterolateralCorner",
+        "children": []
+    },
+    "kneeLateralCollateralLigament": {
+        "value": "Presence and/or severity of lateral collateral ligament findings",
+        "parent": "kneePosterolateralCorner",
+        "children": []
+    },
+    "kneeExtensorMechanism": {
+        "value": "",
+        "parent": "root",
+        "children": [
+            "kneeQuadricepsTendon",
+            "kneePatellarTendon"
+        ]
+    },
+    "kneeQuadricepsTendon": {
+        "value": "",
+        "parent": "kneeExtensorMechanism",
+        "children": [
+            "kneeQuadricepsTendonTearing",
+            "kneeQuadricepsTendinopathy"
+        ]
+    },
+    "kneeQuadricepsTendonTearing": {
+        "value": "Presence and/or severity of quadriceps tendon tearing",
+        "parent": "kneeQuadricepsTendon",
+        "children": []
+    },
+    "kneeQuadricepsTendinopathy": {
+        "value": "Presence and/or severity of quadriceps tendinopathy",
+        "parent": "kneeQuadricepsTendon",
+        "children": []
+    },
+    "kneePatellarTendon": {
+        "value": "",
+        "parent": "kneeExtensorMechanism",
+        "children": [
+            "kneePatellarTendonTearing",
+            "kneePatellarTendinopathy"
+        ]
+    },
+    "kneePatellarTendonTearing": {
+        "value": "Presence and/or severity of patellar tendon tearing",
+        "parent": "kneePatellarTendon",
+        "children": []
+    },
+    "kneePatellarTendinopathy": {
+        "value": "Presence and/or severity of patellar tendinopathy",
+        "parent": "kneePatellarTendon",
+        "children": []
+    },
+    "kneeCartilage": {
+        "value": "Articular cartilage status",
+        "parent": "root",
+        "children": [
+            "kneeCartilageFemoral",
+            "kneeCartilageTibial",
+            "kneeCartilagePatellar",
+            "kneeOsteochondralLesion"
+        ]
+    },
+    "kneeCartilageFemoral": {
+        "value": "",
+        "parent": "kneeCartilage",
+        "children": [
+            "kneeCartilageFemoralMedial",
+            "kneeCartilageFemoralLateral"
+        ]
+    },
+    "kneeCartilageFemoralMedial": {
+        "value": "Presence and/or severity of knee medial femoral cartilage findings",
+        "parent": "kneeCartilageFemoral",
+        "children": []
+    },
+    "kneeCartilageFemoralLateral": {
+        "value": "Presence and/or severity of knee lateral femoral cartilage findings",
+        "parent": "kneeCartilageFemoral",
+        "children": []
+    },
+    "kneeCartilageTibial": {
+        "value": "",
+        "parent": "kneeCartilage",
+        "children": [
+            "kneeCartilageTibialMedial",
+            "kneeCartilageTibialLateral"
+        ]
+    },
+    "kneeCartilageTibialMedial": {
+        "value": "Presence and/or severity of knee medial tibial cartilage findings",
+        "parent": "kneeCartilageTibial",
+        "children": []
+    },
+    "kneeCartilageTibialLateral": {
+        "value": "Presence and/or severity of knee lateral tibial cartilage findings",
+        "parent": "kneeCartilageTibial",
+        "children": []
+    },
+    "kneeCartilagePatellar": {
+        "value": "",
+        "parent": "kneeCartilage",
+        "children": [
+            "kneeCartilagePatellarMedial",
+            "kneeCartilagePatellarLateral"
+        ]
+    },
+    "kneeOsteochondralLesion": {
+        "value": "Presence and/or severity of knee osteochondral lesions/defects",
+        "parent": "kneeCartilage",
+        "children": []
+    },
+    "kneeCartilagePatellarMedial": {
+        "value": "Presence and/or severity of knee medial patellar cartilage findings",
+        "parent": "kneeCartilagePatellar",
+        "children": []
+    },
+    "kneeCartilagePatellarLateral": {
+        "value": "Presence and/or severity of knee lateral patellar cartilage findings",
+        "parent": "kneeCartilagePatellar",
+        "children": []
+    },
+    "kneeBone": {
+        "value": "",
+        "parent": "root",
+        "children": [
+            "kneeBoneFracture",
+            "kneeBoneMarrowEdema",
+            "kneeSubchondralFracture",
+            "kneeOsteonecrosis",
+            "kneeBoneAvn"
+        ]
+    },
+    "kneeBoneFracture": {
+        "value": "Presence, severity and/or location and/or type of knee bone fracture",
+        "parent": "kneeBone",
+        "children": []
+    },
+    "kneeBoneMarrowEdema": {
+        "value": "Presence and/or severity of knee bone marrow edema/contusion",
+        "parent": "kneeBone",
+        "children": []
+    },
+    "kneeSubchondralFracture": {
+        "value": "Presence and/or severity of knee subchondral fractures",
+        "parent": "kneeBone",
+        "children": []
+    },
+    "kneeOsteonecrosis": {
+        "value": "Presence and/or severity of knee osteonecrosis",
+        "parent": "kneeBone",
+        "children": []
+    },
+    "kneeBoneAvn": {
+        "value": "Presence and/or severity of knee avascular necrosis",
+        "parent": "kneeBone",
+        "children": []
+    },
+    "kneeOther": {
+        "value": "Other knee findings",
+        "parent": "root",
+        "children": [
+            "kneeBursa",
+            "kneePoplitealCyst",
+            "kneeGanglionCyst",
+            "kneeMass",
+            "kneeSynovium",
+            "other"
+        ]
+    },
+    "kneeBursa": {
+        "value": "Presence and/or severity of knee bursa findings, e.g. bursitis",
+        "parent": "kneeOther",
+        "children": []
+    },
+    "kneePoplitealCyst": {
+        "value": "Presence and/or extent of knee popliteal/Baker's cyst",
+        "parent": "kneeOther",
+        "children": []
+    },
+    "kneeGanglionCyst": {
+        "value": "Presence and/or extent of knee ganglion cyst",
+        "parent": "kneeOther",
+        "children": []
+    },
+    "kneeMass": {
+        "value": "Presence and/or extent of knee masses, e.g. lipoma, hemangioma, synovial sarcoma",
+        "parent": "kneeOther",
+        "children": []
+    },
+    "kneeSynovium": {
+        "value": "Presence and/or extent of knee synovial findings, e.g. synovitis, thickening",
+        "parent": "kneeOther",
+        "children": []
+    },
+    "other": {
+        "value": "Any other findings not listed above, such as soft tissue findings",
+        "parent": "kneeOther",
+        "children": []
+    }
+}

logger.py

@@ -0,0 +1,13 @@
+class Logger:
+    def __init__(self, log_file):
+        self.log_file = log_file
+
+    def log(self, name, message):
+        try:
+            with open(self.log_file, "a", encoding="utf-8") as f:
+                f.write(f"[{name}]: {message}\n")
+        except OSError as ex:
+            print(f"Error logging message: {ex}")
+
+    def __call__(self, name, message):
+        self.log(name, message)

report_ckecklist_gradio.py

@@ -0,0 +1,85 @@
+import gradio as gr
+import interface
+import utils
+import treegraph as tg
+import logger as lg
+
+system_prompt = """
+You are a critical AI radiology assistant.
+You are helping a radiologist correctly fill out a radiology report.
+The report is regarding a Knee MRI.
+"""
+
+graph, nodes_dict = tg.build_tree_from_file("knee_template.json")
+report_interface = interface.ReportChecklistInterface(
+    llm=utils.LLM(model="gpt-3.5-turbo"),
+    system_prompt=system_prompt,
+    graph=graph,
+    nodes_dict=nodes_dict,
+)
+logger = None
+
+if report_interface.prime_model() is False:
+    print("Model priming failed. Please try again.")
+    exit()
+else:
+    print("Model priming successful.")
+
+running = True
+
+
+def check_report(report, name):
+    global logger, report_interface
+    if len(name.strip()) < 3:
+        return "Please enter a name."
+    else:
+        logger = lg.Logger(log_file="log.txt")
+        report_interface.logger = logger
+        report_interface.username = name
+    if running:
+        results = report_interface.process_input(report)
+        if results == "quit":
+            quit_fn()
+        elif results == "help":
+            return report_interface.help_message
+        elif results == "exception":
+            return "An exception occurred. Please try again."
+        else:
+            return results
+    else:
+        return "Model has been stopped."
+
+
+def quit_fn():
+    global running
+    running = False
+    return "Model has been stopped."
+
+
+with gr.Blocks(theme="soft") as demo:
+    gr.Markdown("## Radiology Report Assistant")
+    gr.Markdown(report_interface.help_message)
+    name_textbox = gr.Textbox(label="Name")
+
+    report_textbox = gr.TextArea(label="Report", lines=20, max_lines=50)
+    with gr.Row():
+        check_btn = gr.Button(
+            value="Check Report",
+        )
+        clear_btn = gr.ClearButton(
+            value="Clear Messages",
+        )
+        quit_btn = gr.Button(
+            value="Quit",
+        )
+    results_textbox = gr.TextArea(label="Results", lines=20, max_lines=50)
+    clear_btn.add([results_textbox, report_textbox])
+
+    check_btn.click(
+        fn=check_report,
+        inputs=[report_textbox, name_textbox],
+        outputs=[results_textbox],
+    )
+    quit_btn.click(fn=quit_fn, outputs=[results_textbox])
+
+demo.launch(auth=("admin", "radiologyreport"))