trying huggingface spaces...

Makefile

@@ -0,0 +1,47 @@
+.PHONY: install run
+
+install:
+	poetry install --no-root
+
+run:
+	poetry run streamlit run app.py
+
+build-tree:
+	poetry run python -c 'from main import build_tree_cache; build_tree_cache()'
+
+test-search:
+	poetry run python -c 'from main import test_search; test_search()'
+
+test-semantic-search:
+	poetry run python -c 'from main import test_semantic_search; test_semantic_search()'
+
+fetch-data:
+	mkdir -p data
+	mkdir -p data/addendum
+	mkdir -p data/code_descriptions
+	mkdir -p data/poa_exempt_codes
+	mkdir -p data/table_index
+
+	@if [ ! -f data/code_descriptions/code_descriptions.zip ]; then \
+	    wget https://ftp.cdc.gov/pub/Health_Statistics/NCHS/Publications/ICD10CM/2025/ICD10-CM%20Code%20Descriptions%202025.zip -O data/code_descriptions/code_descriptions.zip; \
+	    unzip data/code_descriptions/code_descriptions.zip -d data/code_descriptions; \
+	fi
+
+	@if [ ! -f data/addendum/addendum.zip ]; then \
+	    wget https://ftp.cdc.gov/pub/Health_Statistics/NCHS/Publications/ICD10CM/2025/ICD10cm-addendum-2025.zip -O data/addendum/addendum.zip; \
+	    unzip data/addendum/addendum.zip -d data/addendum; \
+	fi
+
+	@if [ ! -f data/table_index/table_index.zip ]; then \
+	    wget https://ftp.cdc.gov/pub/Health_Statistics/NCHS/Publications/ICD10CM/2025/icd10cm-table-index-2025.zip -O data/table_index/table_index.zip; \
+	    unzip data/table_index/table_index.zip -d data/table_index; \
+	fi
+
+	@if [ ! -f data/poa_exempt_codes/poa_exempt_codes.zip ]; then \
+	    wget https://ftp.cdc.gov/pub/Health_Statistics/NCHS/Publications/ICD10CM/2025/POAexemptCodesFY25.zip -O data/poa_exempt_codes/poa_exempt_codes.zip; \
+	    unzip data/poa_exempt_codes/poa_exempt_codes.zip -d data/poa_exempt_codes; \
+	fi
+
+	@if [ ! -f data/guidelines.pdf ]; then \
+	    wget https://ftp.cdc.gov/pub/Health_Statistics/NCHS/Publications/ICD10CM/2025/icd-10-cm-FY25-guidelines-october%20-2024.pdf -O data/guidelines.pdf; \
+	fi

README.md

@@ -1,12 +1,4 @@
----
-title: Icdbot
-emoji: 📊
-colorFrom: indigo
-colorTo: pink
-sdk: streamlit
-sdk_version: 1.44.1
-app_file: app.py
-pinned: false
----
+# ICDBot
+
+ICDBot is a web-based agent with access to the ICD-10-CM vocabulary via search, hierarchically, and other ways. It is currently considered experimental.
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

agents.py

@@ -0,0 +1,148 @@
+from kani_utils.base_kanis import StreamlitKani
+from kani import AIParam, ai_function, ChatMessage, ToolCall
+import streamlit as st
+
+from typing import Annotated
+import pandas as pd
+import re
+import yaml
+
+from icdtree import ICDTree, get_tree
+
+
+class ICDBot(StreamlitKani):
+    # Be sure to override the __init__ method to pass any parameters to the superclass
+    def __init__(self, *args, **kwargs):
+        #kwargs['chat_history'] = self._build_fewshot()
+
+        kwargs['system_prompt'] = """
+        You are ICDBot, a chatbot with programmatic access to the ICD-10-CM vocabulary. You can search for ICD-10 codes via and query for terms' parent and children nodes in the hierarchy.
+
+        * Use the `get_context` function to find the most specific code relevant to the user's query by navigating up and down the term hierarchy. For example, if the search for "fracture humerus" returns S42.2, you might use `get_context` to find the parent code S42.2 and its children codes S42.20, S42.21, ..., to identify the most specific code that matches the user's query.
+        """.strip()
+
+        super().__init__(*args, **kwargs)
+
+        # Define avatars for the agent and user
+        # Can be URLs or emojis
+        self.avatar = "🤖"
+        self.user_avatar = "👤"
+
+        # The name and greeting are shown at the start of the chat
+        # The greeting is not known to the LLM, it serves as a prompt for the user
+        self.name = "ICDBot"
+        self.greeting = "Hello, I'm ICDBot. Nice to meet you!"
+
+        # The description is shown in the sidebar and provides more information about the agent
+        self.description = "An agent with access to the ICD-10-CM vocabulary."
+
+        self.icd10_tree = get_tree()
+
+
+    @ai_function()
+    def semantic_search(self,
+                        query: Annotated[str, AIParam(desc="Query string to search for.")],
+                        limit: Annotated[int, AIParam(desc="Maximum number of results to return. Default to 10 unless necessary.")]=10,
+                        offset: Annotated[int, AIParam(desc="Number of results to skip.")]=0):
+        
+        """Search for ICD-10 codes similar to the query using an embedding-search strategy."""
+
+        search_results = self.icd10_tree.semantic_search(query, limit=limit, offset=offset)
+
+        results_strings = [f"{node.code}: {node.short_desc}; Valid: {node.valid}" for node, _ in search_results]
+        result = yaml.dump({"Matches": results_strings})
+        return result
+
+    @ai_function()
+    def think(self, thought: Annotated[str, AIParam(desc="The thought to think.")]):
+        """Use this function to record internal thoughts or notes, and for planning complex actions. You may call this function more than once in succession for complex plans."""
+        return None
+
+    @ai_function()
+    def get_context(self, 
+                    code: Annotated[str, AIParam(desc="ICD-10 code to retrieve context for")]):
+        """Given an ICD-10 code, return information about its parent, children, and siblings."""
+        node = self.icd10_tree.get_node_by_code(code)
+        if node:
+            parent = node.parent
+            children = node.children
+            siblings = parent.children if parent else []
+            parent_string = f"{parent.code}: {parent.short_desc}; Valid: {parent.valid}" if parent else "No parent"
+            children_strings = [f"{child.code}: {child.short_desc}; Valid: {child.valid}" for child in children]
+            self_string = f"{node.code}: {node.short_desc}; Valid: {node.valid}"
+            siblings_strings = [f"{sibling.code}: {sibling.short_desc}; Valid: {sibling.valid}" for sibling in siblings]
+            # format the result as yaml
+            result = yaml.dump({"Node": self_string, "Parent": parent_string, "Siblings": siblings_strings, "Children": children_strings})
+
+            return result
+        else:
+            return f"Code {code} not found."
+
+
+
+    def _build_fewshot(self):
+
+        history = [ChatMessage.user("Is there a condition that corresponds to rheumatoid arthritis in a patient's wrist, when they also have neuropathy?")]
+        history.append(
+                ChatMessage.assistant(
+                    content=None,
+                    # use a walrus operator to save a reference to the tool call here...
+                    tool_calls=[
+                        tc1 := ToolCall.from_function("search_icd", pattern="[Rr]heumatoid arthritis", max_depth=1)
+                    ],
+                ))
+
+        res = self.search_icd(pattern="[Rr]heumatoid arthritis", max_depth=1)
+        # assert that the string contains "M05.5"
+        assert "M05.5" in res
+
+        history.append(ChatMessage.function("search_icd", res, tc1.id))
+
+        history.append(
+                ChatMessage.assistant(
+                    content=None,
+                    # use a walrus operator to save a reference to the tool call here...
+                    tool_calls=[
+                        tc2 := ToolCall.from_function("get_children", code="M05.5")
+                    ],
+                ))
+
+        res = self.get_children(code="M05.5")
+        # assert that the string contains "M05.53"
+        assert "M05.53" in res
+
+        history.append(ChatMessage.function("get_children", res, tc2.id))
+
+        history.append(
+                ChatMessage.assistant(
+                    content=None,
+                    # use a walrus operator to save a reference to the tool call here...
+                    tool_calls=[
+                        tc3 := ToolCall.from_function("get_children", code="M05.53")
+                    ],
+                ))
+
+        res = self.get_children(code="M05.53")
+        # assert that the string contains "M05.539"
+        assert "M05.539" in res
+
+        history.append(ChatMessage.function("get_children", res, tc3.id))
+
+        # we'll also look for children of M05.539, but there will be none. Still, we'll add it to the history
+        history.append(
+                ChatMessage.assistant(
+                    content=None,
+                    # use a walrus operator to save a reference to the tool call here...
+                    tool_calls=[
+                        tc4 := ToolCall.from_function("get_children", code="M05.539")
+                    ],
+                ))
+
+        res = self.get_children(code="M05.539")
+        # assert that the string is ""
+        assert res == "No children."
+
+        history.append(ChatMessage.function("get_children", res, tc4.id))
+
+        history.append(ChatMessage.assistant("Yes, there are multiple codes for rheumatoid arthritis in the wrist with polyneuropathy. Since you didn't mention left or right, M05.539 'Rheumatoid polyneurop w rheumatoid arthritis of unsp wrist' appears to be the most appropriate."))
+        return history

app.py

@@ -0,0 +1,68 @@
+# Example usage of StreamlitKani
+
+########################
+##### 0 - load libs
+########################
+
+# kani_streamlit imports
+import kani_utils.kani_streamlit_server as ks
+
+# for reading API keys from .env file
+import os
+import dotenv # pip install python-dotenv
+
+# kani imports
+from kani.engines.openai import OpenAIEngine
+
+# load app-defined agents
+from agents import ICDBot
+
+# read API keys .env file (e.g. set OPENAI_API_KEY=.... in .env and gitignore .env)
+dotenv.load_dotenv() 
+
+########################
+##### 1 - Configuration
+########################
+
+# initialize the application and set some page settings
+# parameters here are passed to streamlit.set_page_config, 
+# see more at https://docs.streamlit.io/library/api-reference/utilities/st.set_page_config
+# this function MUST be run first
+ks.initialize_app_config(
+    show_function_calls = True,
+    page_title = "ICDBot",
+    page_icon = "🤖", # can also be a URL
+    initial_sidebar_state = "expanded", # or "expanded"
+    menu_items = {
+            "Get Help": "https://github.com/.../issues",
+            "Report a Bug": "https://github.com/.../issues",
+            "About": "An AI agent with access to the ICD-10-CM vocabulary.",
+        }
+)
+
+
+########################
+##### 2 - Define Agents
+########################
+
+# define an engine to use (see Kani documentation for more info)
+engine = OpenAIEngine(os.environ["OPENAI_API_KEY"], model="gpt-4o")
+# engine = OpenAIEngine(os.environ["OPENAI_API_KEY"], model="o1")
+
+# We also have to define a function that returns a dictionary of agents to serve
+# Agents are keyed by their name, which is what the user will see in the UI
+def get_agents():
+    return {
+            "ICDBot": ICDBot(engine, prompt_tokens_cost = 0.005, completion_tokens_cost = 0.015),
+           }
+
+
+# tell the app to use that function to create agents when needed
+ks.set_app_agents(get_agents)
+
+
+########################
+##### 3 - Serve App
+########################
+
+ks.serve_app()

globals.py

@@ -0,0 +1,3 @@
+# for global stuff if needed
+total_nodes = None
+num_embedded = 0

icdnode.py

@@ -0,0 +1,98 @@
+import re
+from sentence_transformers import SentenceTransformer
+import globals
+import sys
+
+class ICDNode:
+    def __init__(self, order, code, valid, short_desc, long_desc):
+        self.order = order
+
+        code = code.replace('.', '')
+        if len(code) > 3:
+            self.code = code[:3] + '.' + code[3:]
+        else:
+            self.code = code
+
+        self.valid = valid
+        self.short_desc = short_desc
+        self.long_desc = long_desc
+        self.children = []
+        self.parent = None
+        self.embedding = None
+        self.level = 0
+
+    def _set_levels(self):
+        """The level of this node is the level of it's parent (if it's not none). The level of children is one more than the level of the parent."""
+        if self.parent:
+            self.level = self.parent.level + 1
+        for child in self.children:
+            child._set_levels()
+
+    def get_parent(self):
+        """Return the parent node of this node."""
+        return self.parent
+
+    def get_children(self):
+        """Return the immediate children nodes of this node."""
+        return self.children
+
+    def get_descendants(self):
+        """Return the children nodes of this node and all subnodes as a list, in a depth-first order."""
+        children = []
+        for child in self.children:
+            children.append(child)
+            children.extend(child.get_descendants())
+        return children
+    
+    def get_ancestors(self):
+        """Return the parent nodes of this node and all super nodes as a list."""
+        ancestors = []
+        parent = self.parent
+        while parent:
+            ancestors.append(parent)
+            parent = parent.parent
+        return ancestors
+    
+    def _add_embeddings(self, embedding_model, prefix = ""):
+        """Add embeddings to this node and all its descendants."""
+        self.embedding = embedding_model.encode([prefix + self.short_desc])[0]
+        globals.num_embedded += 1
+        percent_complete = round(100 * (globals.num_embedded / globals.total_nodes), 2)
+        sys.stderr.write(f"({percent_complete}%) " + str(self) + "\n")
+        for child in self.children:
+            child._add_embeddings(embedding_model, prefix)
+
+    def _sort_children(self):
+        """Sort the children nodes by their code."""
+        self.children.sort(key=lambda x: x.code)
+
+    def _add_child(self, child_node):
+        """Add a child node to this node."""
+        self.children.append(child_node)
+        child_node.parent = self
+
+    def regex_search(self, pattern, max_depth=None, current_depth=0, valid = None):
+        """Returns a list of node objects whose short_desc, long_desc, 
+        or code matches the regex pattern. If valid is set to True, only
+        return valid nodes; if False, only return invalid nodes. If None,
+        return all matching nodes. max_depth is the maximum depth to search."""
+
+        if max_depth is None:
+            max_depth = float('inf')
+
+        # only return the node if it matches the pattern and is within the max_depth in the hierarchy
+        results = []
+        if re.search(pattern, self.short_desc, re.IGNORECASE) or re.search(pattern, self.long_desc, re.IGNORECASE) or re.search(pattern, self.code, re.IGNORECASE):
+            if valid is None or self.valid == valid:
+                results.append(self)
+        if current_depth < max_depth:
+            for child in self.children:
+                if valid is None or child.valid == valid:
+                    results.extend(child.regex_search(pattern, max_depth = max_depth, current_depth = current_depth + 1, valid = valid))
+
+        return results
+
+
+    def __repr__(self):
+        embed_str_rep = "[" + ", ".join([str(round(x, 5)) for x in self.embedding[:4]]) + ", ...]" if hasattr(self, 'embedding') else None
+        return f"Code: {self.code}\tLevel: {self.level}\tValid: {self.valid}\tDesc: {self.short_desc}\tVec: {embed_str_rep}"

icdtree.py

@@ -0,0 +1,206 @@
+from sentence_transformers import SentenceTransformer
+from icdnode import ICDNode
+import globals
+import sys
+import faiss
+import numpy as np
+import os
+import zipfile
+import requests
+import shutil
+
+def get_tree(file_path = "data/code_descriptions/icd10cm-order-2025.txt", st_embedding_model = 'intfloat/e5-small-v2'):
+    # ok, so we're going to cache the result of the tree building in a pickle file, with a name determined based on the file_path and st_embedding_model
+    # if the pickle file exists, we'll load it and return it
+    # if it doesn't exist, we'll build the tree, save it to the pickle file, and return it
+
+    # first, let's determine the pickle file name, using a hash of the file_path and st_embedding_model
+    import hashlib
+    m = hashlib.md5()
+    m.update(file_path.encode('utf-8'))
+    m.update(st_embedding_model.encode('utf-8'))
+    pickle_file_name = "data/" + m.hexdigest() + ".pkl"
+
+    # now, let's try to load the pickle file
+    import pickle
+    icd_tree = None
+    try:
+        with open(pickle_file_name, 'rb') as f:
+            icd_tree = pickle.load(f)
+            sys.stderr.write("Loaded tree from pickle file\n")
+
+    except FileNotFoundError:
+        try:
+            url = "https://ftp.cdc.gov/pub/Health_Statistics/NCHS/Publications/ICD10CM/2025/ICD10-CM%20Code%20Descriptions%202025.zip"
+            zip_file_path = "data/code_descriptions/code_descriptions.zip"
+
+            # Ensure the directory exists
+            os.makedirs("data/code_descriptions", exist_ok=True)
+
+            # Download the file if it doesn't exist
+            if not os.path.exists(zip_file_path):
+                try:
+                    sys.stderr.write("Downloading ICD data...\n")
+                    response = requests.get(url)
+                    if response.status_code != 200:
+                        sys.stderr.write(f"Failed to download file: {response.status_code}\n")
+                        sys.exit(1)
+                    with open(zip_file_path, 'wb') as f:
+                        f.write(response.content)
+                except Exception as e:
+                    sys.stderr.write(f"Error downloading file: {e}\n")
+                    sys.exit(1)
+
+            # Unzip the file
+            try:
+                sys.stderr.write("Unzipping ICD data...\n")
+                with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
+                    zip_ref.extractall("data/code_descriptions")
+            except zipfile.BadZipFile:
+                sys.stderr.write("Error: The downloaded file is not a valid ZIP file.\n")
+                sys.exit(1)
+
+            # Copy the file to the desired location
+            file_path = "data/code_descriptions/icd10cm-order-2025.txt"
+            if not os.path.exists(file_path):
+                try:
+                    sys.stderr.write("Copying ICD data...\n")
+                    shutil.copy("data/code_descriptions/icd10cm-order-2025.txt", file_path)
+                except Exception as e:
+                    sys.stderr.write(f"Error copying file: {e}\n")
+                    sys.exit(1)
+        except Exception as e:
+            sys.stderr.write("Error downloading or unzipping the file: " + str(e))
+            return None
+
+        # if the file doesn't exist, we'll build the tree
+        sys.stderr.write("Building tree\n")
+        icd_tree = ICDTree(file_path, st_embedding_model)
+
+        # and save it to the pickle file
+        with open(pickle_file_name, 'wb') as f:
+            pickle.dump(icd_tree, f)
+    
+    icd_tree.build_faiss_index()
+    return icd_tree
+
+class ICDTree:
+    def __init__(self, path, st_embedding_model):
+        self.nodes_dict = {}  # Dictionary to hold all nodes by code without dots
+        self.roots = []       # List to hold root nodes
+        self.build_tree(path)
+        globals.total_nodes = len(self.nodes_dict)
+
+        # e5-small-v2 seems to work well on synonym mapping (and not gated unlike the jina model): https://arxiv.org/html/2401.01943v2
+        # intfloat's e5 models all require a "query: " prefix for semantic similarity tasks: see e.g. https://huggingface.co/intfloat/e5-small-v2
+
+        self.embedding_model = SentenceTransformer(st_embedding_model)
+        prefix = ""
+        if "intfloat" in st_embedding_model and "e5" in st_embedding_model:
+            prefix = "query: "
+
+        for root in self.roots:
+            root._add_embeddings(self.embedding_model, prefix)
+
+
+    def build_faiss_index(self):
+        # we need to keep a faiss index for the embeddings
+        # IndexFlatL2 seems to work slightly better than IndexFlatIP, at least for the one query I tried ('hurt tummy')
+        self.index = faiss.IndexFlatL2(self.embedding_model.get_sentence_embedding_dimension())
+        #self.index = faiss.IndexFlatIP(self.embedding_model.get_sentence_embedding_dimension())
+        embeddings = [node.embedding for node in self.nodes_dict.values()]
+       
+        self.index.add(np.array(embeddings))
+
+
+    def semantic_search(self, query, limit = 10, offset = 0):
+        """Search for nodes similar to the query using the sentence transformer model."""
+        query_embedding = self.embedding_model.encode([query])[0]
+
+
+        # search the faiss index
+        D, I = self.index.search(np.array([query_embedding]), limit + offset)
+        results = []
+        for i in range(limit + offset):
+            node = list(self.nodes_dict.values())[I[0][i]]
+            results.append((node, float(D[0][i])))
+
+        # sort by the node's level
+        # results.sort(key=lambda x: x[0].level)
+        # sort results by distance
+        results.sort(key=lambda x: x[1])
+
+        return results[offset:]
+
+
+    # def get_descendants_yaml(self, code):
+    #     """Returns a YAML string representing the node and its descendants, with indentation."""
+    #     root = self.get_node_by_code(code)
+    #     if root:
+    #         return root.get_descendants_yaml()
+    #     else:
+    #         return None
+        
+
+    def regex_search(self, pattern, max_depth=None, valid = None):
+        """Returns a list of node objects whose short_desc, long_desc, or 
+        code matches the regex pattern. max_depth is the maximum depth to 
+        search in the hierarchy; if None, search all levels.
+        
+        If valid is set to 1, only return valid nodes; if 0, only return invalid nodes. If None, return all matching nodes.
+        """
+
+        results = []
+        for root in self.roots:
+            results.extend(root.regex_search(pattern, max_depth = max_depth, valid = valid))
+        return results
+
+    def build_tree(self, file_path):
+        # First pass: Read the file and create all nodes
+        with open(file_path, 'r') as f:
+            for line in f:
+                order = line[0:5].strip()
+                code = line[6:13].strip()
+                valid = line[14:15].strip()
+                short_desc = line[16:76].strip()
+                long_desc = line[77:].strip()
+
+                node = ICDNode(order, code, valid, short_desc, long_desc)
+                code_key = node.code.replace('.', '')  # Key without dots for consistent lookup
+                self.nodes_dict[code_key] = node
+
+        # Second pass: Set parent-child relationships
+        for code_key, node in self.nodes_dict.items():
+            parent_code_key = self.find_parent_code(code_key)
+            if parent_code_key:
+                parent_node = self.nodes_dict[parent_code_key]
+                parent_node._add_child(node)
+            else:
+                # No parent found, this node is a root
+                self.roots.append(node)
+
+        # Sort children of all nodes
+        for code_key, node in self.nodes_dict.items():
+            node._sort_children()
+
+        # Set levels for all nodes
+        for root in self.roots:
+            root._set_levels()
+
+    def find_parent_code(self, code_key):
+        """Find the parent code by progressively stripping characters from the end."""
+        for i in range(len(code_key) - 1, 0, -1):
+            parent_code = code_key[:i]
+            if parent_code in self.nodes_dict:
+                return parent_code
+        return None
+
+    def print_tree(self):
+        """Print all root nodes and their subtrees."""
+        for root in self.roots:
+            root.print_node()
+
+    def get_node_by_code(self, code):
+        """Retrieve a node by its code (with or without dots)."""
+        code_key = code.replace('.', '')
+        return self.nodes_dict.get(code_key)

main.py

@@ -0,0 +1,82 @@
+from icdtree import ICDTree, get_tree
+
+# run as e.g.
+# poetry run python -c 'import main; main.test_semantic_search()'  
+
+def build_tree_cache():
+    icd_tree = get_tree()
+    print("Tree built and saved.")
+
+
+def test_semantic_search():
+    icd10_tree = get_tree()
+    #icd10_tree.build_faiss_index()
+
+    # search for the term "Torus"
+    search_results = icd10_tree.semantic_search("pneumonia", limit=10)
+    for result in search_results:
+        print(result)
+
+    print("getting more")
+    search_results = icd10_tree.semantic_search("pneumonia", limit=10, offset=10)
+    for result in search_results:
+        print(result)
+
+
+def test_search():
+    icd10_tree = get_tree()
+
+    search_nodes = icd10_tree.regex_search(r'Torus', max_depth=5)
+    if len(search_nodes) > 0:
+        # get the highest-level nodes (closest to 0)
+        search_nodes.sort(key=lambda x: x.level)
+        stack = [node for node in search_nodes if node.level == search_nodes[0].level]
+
+        # do a quick depth first search to print all nodes in order
+        while len(stack) > 0:
+            node = stack.pop()
+            print(node)
+            stack.extend(node.get_children())
+
+
+if __name__ == '__main__':
+    print("Try running this file like `python -c 'import main; main.build_tree_cache()'` to build the tree cache.")
+
+    # s42 = icd10_tree.get_node_by_code('S42.001B')
+    # # index 0 is the yaml string, index 1 is the level
+    # print(s42.get_ancestors_yaml()[0])
+
+    # s42 = icd10_tree.get_node_by_code('S42')
+    # print(s42)
+    # for child in s42.get_children():
+    #     print(child)
+    #     for grandchild in child.get_children():
+    #         print(grandchild)
+    # print(s42.get_descendants_yaml())
+
+    # res = icd10_tree.regex_search('.*', max_depth=0)
+    # for node in res:
+    #     print(node)
+
+    # print("#######")
+
+    # focal = icd10_tree.get_node_by_code('M05.5')
+    # children = focal.get_children()
+    # for node in children:
+    #     print(node)
+
+    # print("#######")
+
+    # focal = icd10_tree.get_node_by_code('M05.53')
+    # children = focal.get_children()
+    # for node in children:
+    #     print(node)
+
+    # print("#######")
+
+    # focal = icd10_tree.get_node_by_code('M05.539')
+    # children = focal.get_children()
+    # for node in children:
+    #     print(node)
+
+

oldcode.py

@@ -0,0 +1,31 @@
+# Not really useful, but don't want to delete it yet
+def get_ancestors_yaml(self):
+    """Returns a YAML string representing the node and its ancestors, with indentation. For example:
+    - A00 Cholera:
+      - A00.0 Cholera due to Vibrio cholerae 01, biovar cholerae
+        - A00.00 Cholera due to Vibrio cholerae 01, biovar cholerae, cholera gravis  
+      ..."""
+    my_yaml = f"- {self.code} {self.short_desc}:\n"
+    if self.parent:
+        parent_yaml, parent_level = self.parent.get_ancestors_yaml()
+        my_level = parent_level + 1
+        my_yaml = '  ' * (my_level) + my_yaml
+        yaml = parent_yaml + my_yaml
+    else:
+        yaml = my_yaml
+        my_level = 0
+    return yaml, my_level
+
+def get_descendants_yaml(self):
+    """Returns a YAML string representing the node and its descendants, with indentation. For example:
+    - A00 Cholera:
+        - A00.0 Cholera due to Vibrio cholerae 01, biovar cholerae
+        - A00.1 Cholera due to Vibrio cholerae 01, biovar eltor
+        - A00.9 Cholera, unspecified
+        ..."""
+    yaml = f"- {self.code} {self.short_desc}:\n"
+    for child in self.children:
+        child_yaml = child.get_descendants_yaml()
+        child_yaml = child_yaml.replace('\n', '\n  ')
+        yaml += f"  {child_yaml}"
+    return yaml

pyproject.toml

@@ -0,0 +1,27 @@
+[tool.poetry]
+name = "icdbot"
+version = "0.1.0"
+description = "An AI agent for intelligent querying of the ICD10CM vocabulary."
+authors = ["Shawn T O'Neil <[email protected]>"]
+license = "MIT"
+readme = "README.md"
+package-mode = false
+
+[tool.poetry.dependencies]
+python = "^3.10"
+kani-utils = {git = "https://github.com/oneilsh/kani-utils.git"}
+pandas = "^2.2.3"
+kani = "^1.2.2"
+transformers = "^4.48.0"
+sentence-transformers = "^3.3.1"
+peft = "^0.14.0"
+faiss-cpu = "^1.9.0.post1"
+numpy = "^2.2.2"
+pyyaml = "^6.0.2"
+streamlit = "^1.41.1"
+requests = "^2.32.3"
+
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"