File size: 1,499 Bytes
b8986a1 85d173c b8986a1 85d173c b8986a1 85d173c b8986a1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
import re
import tempfile
import os
from pyvis.network import Network
def extract_key_terms(text: str):
"""
A simple method to extract key terms by matching capitalized words.
In a clinical context, this may highlight important entities.
"""
return re.findall(r"\b[A-Z][a-zA-Z]+\b", text)
def create_medical_graph(query: str, docs: list) -> str:
"""
Creates an interactive knowledge graph using PyVis:
- A central node for the query.
- One node per abstract.
- Sub-nodes for key terms extracted from each abstract.
Returns the HTML content of the graph.
"""
net = Network(height="600px", width="100%", directed=False)
net.add_node("QUERY", label=f"Query: {query}", color="red", shape="star")
for i, doc in enumerate(docs):
doc_id = f"Doc_{i}"
net.add_node(doc_id, label=f"Abstract {i+1}", color="blue")
net.add_edge("QUERY", doc_id)
terms = extract_key_terms(doc)
for term in set(terms):
term_id = f"{doc_id}_{term}"
net.add_node(term_id, label=term, color="green")
net.add_edge(doc_id, term_id)
# Write the graph to a temporary HTML file and return its content.
with tempfile.NamedTemporaryFile(delete=False, suffix=".html") as tmp:
temp_filename = tmp.name
net.show(temp_filename)
with open(temp_filename, "r", encoding="utf-8") as f:
html_content = f.read()
os.remove(temp_filename)
return html_content
|