| 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') |
|
|
| |
| 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}" |
|
|
