main.py

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)