Creating a Knowledge Graph Using an LLM
In the realm of artificial intelligence, one of the most interesting applications is the creation of Knowledge Graphs from unstructured data. This article will explore how to construct a Knowledge Graph from a medical log using a Large Language Model (LLM) like GPT-4o-mini. Unlike traditional Natural Language Processing (NLP) methods, which may struggle with messy data, LLMs provide enhanced accuracy and context, making them invaluable for tasks involving complex information.
Getting Started
To embark on creating a Knowledge Graph, we first need to prepare our environment. This involves installing some important dependencies using Python:
!pip install "mirascope[openai]" matplotlib networkx
Obtaining an OpenAI API Key
To use GPT-4o-mini, you’ll need an OpenAI API key. This can be obtained by visiting the OpenAI API Keys page and generating a new key. Keep in mind that new users may need to enter billing details and make an initial payment of $5 to activate their API access.
import os
from getpass import getpass
os.environ["OPENAI_API_KEY"] = getpass('Enter OpenAI API Key: ')
Defining a Graph Schema
Next, we need a way to structure the information we want to extract. A simple schema can be defined using Pydantic, encompassing nodes, edges, and the overall Knowledge Graph. Here’s how we can represent our schema:
from pydantic import BaseModel, Field
class Edge(BaseModel):
source: str
target: str
relationship: str
class Node(BaseModel):
id: str
type: str
properties: dict | None = None
class KnowledgeGraph(BaseModel):
nodes: list[Node]
edges: list[Edge]
Defining the Patient Log
To generate our Knowledge Graph, we need unstructured data. Let’s take a look at a sample patient log for a patient named Mary:
patient_log = """
Mary called for help at 3:45 AM, reporting that she had fallen while going to the bathroom. This marks the second fall incident within a week. She complained of dizziness before the fall.
Earlier in the day, Mary was observed wandering the hallway and appeared confused when asked basic questions. She was unable to recall the names of her medications and asked the same question multiple times.
Mary skipped both lunch and dinner, stating she didn't feel hungry. When the nurse checked her room in the evening, Mary was lying in bed with mild bruising on her left arm and complained of hip pain.
Vital signs taken at 9:00 PM showed slightly elevated blood pressure and a low-grade fever (99.8°F). Nurse also noted increased forgetfulness and possible signs of dehydration.
This behavior is similar to previous episodes reported last month.
"""
Generating the Knowledge Graph
With our schema and patient log defined, we can now harness the power of LLMs to extract structured insights from unstructured text. The following function will analyze the patient log and identify entities and relationships:
from mirascope.core import openai, prompt_template
@openai.call(model="gpt-4o-mini", response_model=KnowledgeGraph)
@prompt_template(
"""
SYSTEM:
Extract a knowledge graph from this patient log.
Use Nodes to represent people, symptoms, events, and observations.
Use Edges to represent relationships like "has symptom", "reported", "noted", etc.
The log:
{log_text}
"""
)
def generate_kg(log_text: str) -> openai.OpenAIDynamicConfig:
return {"log_text": log_text}
kg = generate_kg(patient_log)
print(kg)
Querying the Graph
Once the Knowledge Graph is created, we can query it to find out more about the patient’s health risks or concerns. Below is a function that takes a natural language question and the structured graph, allowing us to retrieve informative responses:
@openai.call(model="gpt-4o-mini")
@prompt_template(
"""
SYSTEM:
Use the knowledge graph to answer the user's question.
Graph:
{knowledge_graph}
USER:
{question}
"""
)
def run(question: str, knowledge_graph: KnowledgeGraph): ...
question = "What health risks or concerns does Mary exhibit based on her recent behavior and vitals?"
print(run(question, kg))
Visualizing the Graph
To better understand the relationships in our Knowledge Graph, we can visualize it. Below is how we can create a clear and interactive representation using Matplotlib and NetworkX:
import matplotlib.pyplot as plt
import networkx as nx
def render_graph(kg: KnowledgeGraph):
G = nx.DiGraph()
for node in kg.nodes:
G.add_node(node.id, label=node.type, **(node.properties or {}))
for edge in kg.edges:
G.add_edge(edge.source, edge.target, label=edge.relationship)
plt.figure(figsize=(15, 10))
pos = nx.spring_layout(G)
nx.draw_networkx_nodes(G, pos, node_size=2000, node_color="lightgreen")
nx.draw_networkx_edges(G, pos, arrowstyle="->", arrowsize=20)
nx.draw_networkx_labels(G, pos, font_size=12, font_weight="bold")
edge_labels = nx.get_edge_attributes(G, "label")
nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_color="blue")
plt.title("Healthcare Knowledge Graph", fontsize=15)
plt.show()
render_graph(kg)
Conclusion
In summary, creating a Knowledge Graph from unstructured medical logs using LLMs like GPT-4o-mini provides a powerful tool for extracting valuable insights. By structuring the data effectively, practitioners can gain a clearer understanding of patient conditions and risks. This approach not only enhances decision-making in healthcare but also opens avenues for further research and application in diverse fields.
FAQs
- What is a Knowledge Graph? A Knowledge Graph is a structured representation of information, where entities are connected by relationships, making it easier to understand complex data.
- How do LLMs improve the extraction of information? LLMs excel in understanding context and semantics, which allows them to extract relevant entities and relationships more accurately from unstructured data.
- What is the importance of defining a schema? Defining a schema is crucial as it provides a framework for organizing information and ensures that the data extracted is meaningful and usable.
- Can this method be applied to other fields? Yes, the principles outlined can be adapted to various fields such as finance, legal research, and customer support to extract and visualize insights from unstructured data.
- How can I visualize the Knowledge Graph? You can use libraries like Matplotlib and NetworkX in Python to create clear graphical representations of your Knowledge Graph.
