README.md

metadata

title: Healthcare GNN GraphRAG
emoji: 🏥
colorFrom: blue
colorTo: green
sdk: docker
app_port: 7860
pinned: false

Healthcare GNN-based GraphRAG Pipeline

Course: Big Data Applications — Lab 03: GNN-based RAG for LLM Inference
Dataset: qiaojin/PubMedQA (pqa_labeled)
LLM: Jackrong/Qwen3.5-4B-Neo-GGUF (Q4_K_S, CPU inference)
Embedding: BAAI/bge-small-en-v1.5 (384-dim)

---

Overview

Standard Retrieval-Augmented Generation (RAG) retrieves context by measuring cosine similarity between a query vector and flat document chunk embeddings. This approach is topology-blind: two entities that co-occur in many medical relationships receive no higher retrieval priority than isolated, semantically similar text fragments.

This project introduces a Graph-augmented RAG pipeline that explicitly models the relational structure of a medical knowledge graph. A Variational Graph AutoEncoder (VGAE) with a GraphSAGE backbone is trained on the knowledge graph via a link-prediction objective, producing structural embeddings that encode each entity's topological neighbourhood. At query time, both semantic similarity and graph-structural proximity are fused through a calibrated linear interpolation to rank candidate context nodes.

---

System Architecture

PubMedQA Dataset
       │
       ▼
┌─────────────────────────┐
│  1. KG Construction     │  (offline, pre-computed)
│  LlamaIndex + Qwen-4B   │
│  → PropertyGraphIndex   │
└────────────┬────────────┘
             │  entities + relations
             ▼
┌─────────────────────────┐
│  2. PyG Conversion      │  (offline, pre-computed)
│  BAAI/bge-small-en-v1.5 │
│  → node features X      │
│  → edge_index E         │
└────────────┬────────────┘
             │  Data(x=X, edge_index=E)
             ▼
┌─────────────────────────┐
│  3. VGAE Training       │  (offline, pre-computed)
│  GraphSAGE encoder      │
│  Link-prediction loss   │
│  → structural emb. Z    │
└────────────┬────────────┘
             │  gnn_model.pth, pyg_data.pt
             ▼
┌─────────────────────────┐
│  4. Hybrid Retrieval    │  (online, per-query)
│  GNNHybridRetriever     │
│  α·sem + (1-α)·struct   │
└────────────┬────────────┘
             │  Top-K context nodes
             ▼
┌─────────────────────────┐
│  5. LLM Generation      │  (online, per-query)
│  Qwen3.5-4B Q4_K_S      │
│  llama-cpp CPU          │
└─────────────────────────┘

---

Quick Start

Requirements

• Python ≥ 3.10
• RAM ≥ 8 GB (16 GB recommended)
• No GPU required — all inference runs on CPU

Installation

git clone <repo-url>
cd HealthcareGraphRAG
python -m venv .venv
source .venv/bin/activate        # Windows: .venv\Scripts\activate
pip install -r requirements.txt

Run the app

The pre-computed artifacts (storage_graph/) are committed to the repository, so you can launch the app directly:

python app.py

Open http://localhost:7860 in your browser. The system initialises in the background (loading the LLM takes ~30 s on first run); the status indicator in the top-right corner turns green when ready.

---

Reproduce Artifacts

Run the three offline steps in order from the repository root. Each step reads the output of the previous one.

Step 1 — Build the Knowledge Graph

Downloads PubMedQA, extracts SPO triples with Qwen-4B, and persists a PropertyGraphIndex.

python -m src.indexing.kg_builder
# Output: ./storage_graph/  (LlamaIndex graph store)

MAX_DOCS = 10 by default. Edit src/indexing/kg_builder.py to process more documents.

Step 2 — Convert to PyG

Encodes graph nodes with BGE-small and builds a PyTorch Geometric Data object.

python -m src.graph.pyg_converter
# Output: ./storage_graph/pyg_data.pt

Step 3 — Train the VGAE

Trains the GraphSAGE-VGAE on a link-prediction objective and saves structural embeddings.

python -m src.gnn.trainer
# Output: ./storage_graph/pyg_data.pt  (updated with structural_embeddings)
#         ./storage_graph/gnn_model.pth

---

Step-by-Step Pipeline

Step 1 — Knowledge Graph Construction

PubMedQA abstracts are segmented into 150-word chunks and fed to LlamaIndex's PropertyGraphIndex with a SimpleLLMPathExtractor. The extractor prompts Qwen3.5-4B (via llama-cpp) to parse each chunk into subject–predicate–object triples, which are accumulated into a labelled property graph $\mathcal{G} = (\mathcal{V}, \mathcal{E})$.

Input: raw PubMed abstracts with MeSH annotations
Output: a persisted LlamaIndex PropertyGraphIndex

---

Step 2 — PyG Graph Conversion

The property graph is converted into a PyTorch Geometric Data object suitable for GNN training.

Node feature matrix $X \in \mathbb{R}^{N \times 384}$, where $N$ is the number of entities:

$$X_i=\text{BGE-small}\text{\hspace{0.17em}⁣}(\text{name}(v_i))\mathrm{\forall }{v}_i\in \mathcal{V}X\_i\; =\; \backslash text\{BGE-small\}\backslash !\backslash left(\backslash text\{name\}(v\_i)\backslash right)\; \backslash quad\; \backslash forall\backslash ,\; v\_i\; \backslash in\; \backslash mathcal\{V\}$$

Edge index $E \in \mathbb{Z}^{2 \times |\mathcal{E}|}$: source and target node indices for each directed relation.

Output: a checkpoint containing the PyG Data object with node features and edge indices, plus bidirectional node ID mappings.

---

Step 3 — VGAE Training

3.1 Model: GraphSAGE-VGAE

The encoder is a two-layer GraphSAGE network that outputs the parameters of a Gaussian posterior over each node's latent representation:

$$h_i^{(1)}=\text{ReLU}\text{\hspace{0.17em}⁣}(W_1\cdot \text{MEAN}\text{\hspace{0.17em}⁣}(\{x_j{\}}_{j\in \mathcal{N}(i)\cup \{i\}})),h_i^{(1)}\in {\mathbb{R}}^{256}h\_i^\{(1)\}\; =\; \backslash text\{ReLU\}\backslash !\backslash left(W\_1\; \backslash cdot\; \backslash text\{MEAN\}\backslash !\backslash left(\backslash \{x\_j\backslash \}\_\{j\; \backslash in\; \backslash mathcal\{N\}(i)\; \backslash cup\; \backslash \{i\backslash \}\}\backslash right)\backslash right),\; \backslash quad\; h\_i^\{(1)\}\; \backslash in\; \backslash mathbb\{R\}^\{256\}$$

$${\mu }_i=W_{\mu }\cdot \text{MEAN}\text{\hspace{0.17em}⁣}(\{h_j^{(1)}{\}}_{j\in \mathcal{N}(i)\cup \{i\}}),{\mu }_i\in {\mathbb{R}}^{128}\backslash mu\_i\; =\; W\_\backslash mu\; \backslash cdot\; \backslash text\{MEAN\}\backslash !\backslash left(\backslash \{h\_j^\{(1)\}\backslash \}\_\{j\; \backslash in\; \backslash mathcal\{N\}(i)\; \backslash cup\; \backslash \{i\backslash \}\}\backslash right),\; \backslash quad\; \backslash mu\_i\; \backslash in\; \backslash mathbb\{R\}^\{128\}$$

$$\log {\sigma }_i=\text{clamp}\text{\hspace{0.17em}⁣}(W_{\sigma }\cdot \text{MEAN}\text{\hspace{0.17em}⁣}(\{h_j^{(1)}{\}}_{j\in \mathcal{N}(i)\cup \{i\}}),-10,10)\backslash log\; \backslash sigma\_i\; =\; \backslash text\{clamp\}\backslash !\backslash left(W\_\backslash sigma\; \backslash cdot\; \backslash text\{MEAN\}\backslash !\backslash left(\backslash \{h\_j^\{(1)\}\backslash \}\_\{j\; \backslash in\; \backslash mathcal\{N\}(i)\; \backslash cup\; \backslash \{i\backslash \}\}\backslash right),\backslash ;\; -10,\backslash ;\; 10\backslash right)$$

Reparameterisation (training only):

$$z_i={\mu }_i+\epsilon \odot \exp (\log {\sigma }_i),\epsilon \sim \mathcal{N}(0,I)z\_i\; =\; \backslash mu\_i\; +\; \backslash varepsilon\; \backslash odot\; \backslash exp(\backslash log\; \backslash sigma\_i),\; \backslash quad\; \backslash varepsilon\; \backslash sim\; \backslash mathcal\{N\}(0,\; I)$$

At inference the deterministic mean $z_i = \mu_i$ is used, eliminating stochastic variance.

3.2 Decoder

The decoder computes the probability of an edge between nodes $i$ and $j$ as the inner product of their latent vectors:

$${\widehat{A}}_{ij}=\sigma \text{\hspace{0.17em}⁣}\left(z_i^{\mathrm{\top }}{z}_j\right)\backslash hat\{A\}\_\{ij\}\; =\; \backslash sigma\backslash !\backslash left(z\_i^\backslash top\; z\_j\backslash right)$$

3.3 Training Objective

The model is trained with a link-prediction binary cross-entropy loss plus a KL regularisation term:

$$\mathcal{L}={\mathcal{L}}_{\text{recon}}+{\mathcal{L}}_{\text{KL}}\backslash mathcal\{L\}\; =\; \backslash mathcal\{L\}\_\{\backslash text\{recon\}\}\; +\; \backslash mathcal\{L\}\_\{\backslash text\{KL\}\}$$

$${\mathcal{L}}_{\text{recon}}=-\frac{1}{\mathrm{\mid }{\mathcal{E}}^{+}\mathrm{\mid }+\mathrm{\mid }{\mathcal{E}}^{-}\mathrm{\mid }}[\sum _{(i,j)\in {\mathcal{E}}^{+}}\log {\widehat{A}}_{ij}+\sum _{(i,j)\in {\mathcal{E}}^{-}}\log \text{\hspace{0.17em}⁣}(1-{\widehat{A}}_{ij})]\backslash mathcal\{L\}\_\{\backslash text\{recon\}\}\; =\; -\backslash frac\{1\}\{|\backslash mathcal\{E\}^+|\; +\; |\backslash mathcal\{E\}^-|\}\backslash left[\backslash sum\_\{(i,j)\backslash in\backslash mathcal\{E\}^+\}\; \backslash log\; \backslash hat\{A\}\_\{ij\}\; +\; \backslash sum\_\{(i,j)\backslash in\backslash mathcal\{E\}^-\}\; \backslash log\backslash !\backslash left(1\; -\; \backslash hat\{A\}\_\{ij\}\backslash right)\backslash right]$$

$${\mathcal{L}}_{\text{KL}}=-\frac{1}{2N}\sum _{i=1}^N(1+2\log {\sigma }_i-{\mu }_i^2-{\sigma }_i^2)\backslash mathcal\{L\}\_\{\backslash text\{KL\}\}\; =\; -\backslash frac\{1\}\{2N\}\backslash sum\_\{i=1\}^\{N\}\backslash left(1\; +\; 2\backslash log\backslash sigma\_i\; -\; \backslash mu\_i^2\; -\; \backslash sigma\_i^2\backslash right)$$

Negative edges $\mathcal{E}^-$ are sampled uniformly at random with $|\mathcal{E}^-| = |\mathcal{E}^+|$ per epoch.

Optimiser: Adam, $\text{lr} = 0.01$, 100 epochs.

Output: trained VGAE weights and structural embeddings $Z \in \mathbb{R}^{N \times 128}$ persisted alongside the graph checkpoint.

---

Step 4 — Hybrid Retrieval

At query time, GNNHybridRetriever fuses two complementary similarity signals.

4.1 Semantic Score

The query $q$ is encoded by the same BGE-small model used during graph construction:

$${\mathbf{e}}_q^{\text{sem}}=\text{BGE-small}(q)\in {\mathbb{R}}^{384}\backslash mathbf\{e\}\_q^\{\backslash text\{sem\}\}\; =\; \backslash text\{BGE-small\}(q)\; \backslash in\; \backslash mathbb\{R\}^\{384\}$$

Cosine similarity against all node semantic features:

$$s_i^{\text{sem}}=\frac{{\mathbf{e}}_q^{\text{sem}}\cdot X_i}{\mathrm{\parallel }{\mathbf{e}}_q^{\text{sem}}\mathrm{\parallel }\mathrm{\parallel }{X}_i\mathrm{\parallel }}s\_i^\{\backslash text\{sem\}\}\; =\; \backslash frac\{\backslash mathbf\{e\}\_q^\{\backslash text\{sem\}\}\; \backslash cdot\; X\_i\}\{\backslash |\backslash mathbf\{e\}\_q^\{\backslash text\{sem\}\}\backslash |\backslash ,\backslash |X\_i\backslash |\}$$

4.2 Structural Score

The query is treated as an isolated node (zero in-degree / out-degree) and its structural embedding is computed by forwarding $\mathbf{e}_q^{\text{sem}}$ through the frozen VGAE encoder with an empty edge index:

$${\mathbf{e}}_q^{\text{struct}}=\text{VGAE\_encoder}\text{\hspace{0.17em}⁣}({\mathbf{e}}_q^{\text{sem}},\mathrm{\varnothing })\in {\mathbb{R}}^{128}\backslash mathbf\{e\}\_q^\{\backslash text\{struct\}\}\; =\; \backslash text\{VGAE\backslash \_encoder\}\backslash !\backslash left(\backslash mathbf\{e\}\_q^\{\backslash text\{sem\}\},\backslash ;\; \backslash varnothing\backslash right)\; \backslash in\; \backslash mathbb\{R\}^\{128\}$$

Cosine similarity against all precomputed structural embeddings $Z$:

$$s_i^{\text{struct}}=\frac{{\mathbf{e}}_q^{\text{struct}}\cdot Z_i}{\mathrm{\parallel }{\mathbf{e}}_q^{\text{struct}}\mathrm{\parallel }\mathrm{\parallel }{Z}_i\mathrm{\parallel }}s\_i^\{\backslash text\{struct\}\}\; =\; \backslash frac\{\backslash mathbf\{e\}\_q^\{\backslash text\{struct\}\}\; \backslash cdot\; Z\_i\}\{\backslash |\backslash mathbf\{e\}\_q^\{\backslash text\{struct\}\}\backslash |\backslash ,\backslash |Z\_i\backslash |\}$$

4.3 Score Fusion

Because $\mathbf{e}_q^{\text{sem}}$ and $\mathbf{e}_q^{\text{struct}}$ live in different metric spaces (384-dim vs. 128-dim), their cosine scores have different numerical ranges. Min-Max normalisation maps each score vector independently to $[0, 1]$:

$${\tilde{s}}_i^{(\cdot )}=\frac{s_i^{(\cdot )}-\underset{j}{\min }{s}_j^{(\cdot )}}{\underset{j}{\max }{s}_j^{(\cdot )}-\underset{j}{\min }{s}_j^{(\cdot )}}\backslash tilde\{s\}\_i^\{\backslash ,(\backslash cdot)\}\; =\; \backslash frac\{s\_i^\{\backslash ,(\backslash cdot)\}\; -\; \backslash min\_j\; s\_j^\{\backslash ,(\backslash cdot)\}\}\{\backslash max\_j\; s\_j^\{\backslash ,(\backslash cdot)\}\; -\; \backslash min\_j\; s\_j^\{\backslash ,(\backslash cdot)\}\}$$

The final ranking score is a convex combination controlled by $\alpha \in [0, 1]$:

$$\boxed{{\displaystyle f_i=\alpha {\tilde{s}}_i^{\text{sem}}+(1-\alpha ){\tilde{s}}_i^{\text{struct}}}}\backslash boxed\{f\_i\; =\; \backslash alpha\backslash ,\backslash tilde\{s\}\_i^\{\backslash text\{sem\}\}\; +\; (1-\backslash alpha)\backslash ,\backslash tilde\{s\}\_i^\{\backslash text\{struct\}\}\}$$

The top-$K$ nodes ranked by $f_i$ are retrieved and their text representations are concatenated as the context passage prepended to the LLM prompt ($\alpha = 0.5$, $K = 3$ by default).

---

Step 5 — LLM Generation

The retrieved context and user query are composed into a prompt and fed to Qwen3.5-4B quantised to Q4_K_S (≈ 2.4 GB), run on CPU via llama-cpp-python:

Parameter	Value
Context window	8 192 tokens
Max new tokens	2 048
Temperature	0.0 (deterministic)
Threads	4 (HF Free CPU)

The model may emit <think>…</think> chain-of-thought tokens. The UI collapses these into a collapsible Reasoning block and surfaces only the final answer.

Continue mode: if the model is interrupted mid-stream (user clicks Stop), typing continue / cont resumes generation from the exact token position where streaming stopped, without re-running retrieval.

---

Deployment Constraints

Constraint	Mitigation
No GPU	GGUF Q4_K_S quantisation; llama-cpp CPU backend
16 GB RAM cap	KG construction is fully offline; only inference runs live
Cold-start latency	Model path resolution checks local cache before downloading
OOM risk	Chunk size capped at 150 words; context window at 8 192 tokens

---

Repository Structure

.
├── app.py                          # Gradio UI + LLM inference entry point
├── requirements.txt
├── src/
│   ├── utils.py                    # Shared utilities (UTF8FS, model path resolver)
│   ├── indexing/
│   │   └── kg_builder.py           # Step 1: KG extraction (offline)
│   ├── graph/
│   │   └── pyg_converter.py        # Step 2: PyG conversion (offline)
│   ├── gnn/
│   │   ├── model.py                # VGAE / GraphSAGE encoder
│   │   └── trainer.py              # Step 3: link-prediction training (offline)
│   └── retrieval/
│       └── hybrid_retriever.py     # Step 4: dual-score retrieval (online)
└── storage_graph/                  # Pre-computed artefacts (committed)
    ├── pyg_data.pt                 # PyG graph + structural embeddings
    └── gnn_model.pth               # Trained VGAE weights

---

Key Design Choices

GraphSAGE over GCN: GCN applies symmetric degree-normalised aggregation $\hat{D}^{-1/2}\hat{A}\hat{D}^{-1/2}$, which penalises high-degree nodes disproportionately on sparse graphs. GraphSAGE's mean aggregation is degree-agnostic and empirically more stable on knowledge graphs with heterogeneous degree distributions.

VGAE over deterministic GAE: The variational posterior $q(Z \mid X, E) = \prod_i \mathcal{N}(z_i \mid \mu_i, \sigma_i^2 I)$ regularises the latent space via the KL term, preventing degenerate embeddings when the graph is sparse. At inference, using the mean $\mu_i$ (rather than sampling) provides deterministic, reproducible retrieval scores.

Isolated-node query projection: Rather than training a separate query encoder or approximating the graph neighbourhood of the query, we exploit the VGAE encoder's ability to process a degree-zero node. This avoids data leakage (the query has no ground-truth edges) and requires no additional parameters.

Min-Max normalisation over softmax: Softmax introduces a temperature-sensitive denominator that interacts poorly when score distributions differ in sharpness across the two spaces. Min-Max normalisation is a simple, parameter-free linear rescaling that preserves the ordinal ranking within each space before fusion.