Tutorial

Building a Clinical Knowledge Graph with BioBERT and Neo4j

Difficulty: Advanced Time: 20 min read

Introduction

In this cookbook, we will explore how to construct a robust Clinical Knowledge Graph (CKG) by extracting medical entities and relations from unstructured clinical notes using BioBERT, and loading them into a Neo4j graph database for advanced querying and inference.

Architecture Overview


graph LR
    EHR(Unstructured Clinical Notes) --> BioBERT(BioBERT NER Pipeline)
    BioBERT --> Entities(Entities: Disease, Drug)
    BioBERT --> Rel(Relation Extraction Model)
    Entities --> Neo(Neo4j Graph Database)
    Rel --> Neo
    Neo --> Query(Cypher Queries / GDS Algorithms)

Prerequisites

  • Python 3.9+
  • Transformers library (Huggingface)
  • Neo4j desktop or AuraDB instance
  • neo4j Python driver

Step 1: Entity Extraction with BioBERT

We'll use a pre-trained BioBERT model fine-tuned for Named Entity Recognition (NER) to identify diseases, symptoms, and medications from clinical text. Using the aggregation_strategy="simple" is crucial for fusing sub-word tokens back into full clinical terms.

from transformers import AutoTokenizer, AutoModelForTokenClassification
from transformers import pipeline

tokenizer = AutoTokenizer.from_pretrained("d4data/biomedical-ner-all")
model = AutoModelForTokenClassification.from_pretrained("d4data/biomedical-ner-all")

nlp = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple")

clinical_note = "Patient prescribed 20mg Lisinopril for severe hypertension and heart failure."
entities = nlp(clinical_note)

print([(ent['word'], ent['entity_group']) for ent in entities])
# Output: [('Lisinopril', 'Medication'), ('severe hypertension', 'MedicalCondition'), ('heart failure', 'MedicalCondition')]

Step 2: Connecting to Neo4j and Merging Entities

Ensure your Neo4j instance is running and connect using the official Python driver. We use the MERGE command to ensure we don't create duplicate nodes for the same medical concept.

from neo4j import GraphDatabase

URI = "bolt://localhost:7687"
AUTH = ("neo4j", "password")

driver = GraphDatabase.driver(URI, auth=AUTH)

def create_entities(tx, entities):
    for ent in entities:
        # Sanitize entity group for label use
        label = ent['entity_group'].replace(" ", "")
        query = f"MERGE (e:{label} )"
        tx.run(query, name=ent['word'].lower())

with driver.session() as session:
    session.execute_write(create_entities, entities)

Step 3: Creating Relationships

We can establish relationships between extracted entities using co-occurrence within the same sentence, or by piping the entities into a secondary Relation Extraction (RE) model. Here we assume a co-occurrence heuristic.

def create_relationship(tx, med, condition):
    query = (
        "MATCH (m:Medication {name: $med}) "
        "MATCH (c:MedicalCondition {name: $condition}) "
        "MERGE (m)-[:PRESCRIBED_FOR]->(c)"
    )
    tx.run(query, med=med.lower(), condition=condition.lower())

with driver.session() as session:
    session.execute_write(create_relationship, "Lisinopril", "severe hypertension")
    session.execute_write(create_relationship, "Lisinopril", "heart failure")

Step 4: Querying the Graph and Graph Data Science (GDS)

Once populated, you can perform powerful Cypher queries to discover new insights.

MATCH (m:Medication)-[:PRESCRIBED_FOR]->(c:MedicalCondition {name: 'severe hypertension'})
RETURN m.name;

You can also project the graph into memory using the Neo4j Graph Data Science library to run PageRank, identifying the most centrally prescribed medications across a massive patient cohort.

CALL gds.pageRank.stream({
  nodeProjection: ['Medication', 'MedicalCondition'],
  relationshipProjection: 'PRESCRIBED_FOR'
})
YIELD nodeId, score
RETURN gds.util.asNode(nodeId).name AS name, score
ORDER BY score DESC LIMIT 5

Conclusion

By marrying state-of-the-art NLP models like BioBERT with powerful graph databases like Neo4j, healthcare organizations can unlock massive value from their unstructured clinical data, enabling advanced analytics, clinical decision support, and cohort discovery.