---
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`](https://huggingface.co/datasets/qiaojin/PubMedQA) (pqa_labeled)  
**LLM**: [`Jackrong/Qwen3.5-4B-Neo-GGUF`](https://huggingface.co/Jackrong/Qwen3.5-4B-Neo-GGUF) (Q4_K_S, CPU inference)  
**Embedding**: [`BAAI/bge-small-en-v1.5`](https://huggingface.co/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

```bash
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:

```bash
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`.

```bash
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.

```bash
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.

```bash
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}\!\left(\text{name}(v_i)\right) \quad \forall\, v_i \in \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}\!\left(W_1 \cdot \text{MEAN}\!\left(\{x_j\}_{j \in \mathcal{N}(i) \cup \{i\}}\right)\right), \quad h_i^{(1)} \in \mathbb{R}^{256}$$

$$\mu_i = W_\mu \cdot \text{MEAN}\!\left(\{h_j^{(1)}\}_{j \in \mathcal{N}(i) \cup \{i\}}\right), \quad \mu_i \in \mathbb{R}^{128}$$

$$\log \sigma_i = \text{clamp}\!\left(W_\sigma \cdot \text{MEAN}\!\left(\{h_j^{(1)}\}_{j \in \mathcal{N}(i) \cup \{i\}}\right),\; -10,\; 10\right)$$

**Reparameterisation** (training only):

$$z_i = \mu_i + \varepsilon \odot \exp(\log \sigma_i), \quad \varepsilon \sim \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:

$$\hat{A}_{ij} = \sigma\!\left(z_i^\top z_j\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}}$$

$$\mathcal{L}_{\text{recon}} = -\frac{1}{|\mathcal{E}^+| + |\mathcal{E}^-|}\left[\sum_{(i,j)\in\mathcal{E}^+} \log \hat{A}_{ij} + \sum_{(i,j)\in\mathcal{E}^-} \log\!\left(1 - \hat{A}_{ij}\right)\right]$$

$$\mathcal{L}_{\text{KL}} = -\frac{1}{2N}\sum_{i=1}^{N}\left(1 + 2\log\sigma_i - \mu_i^2 - \sigma_i^2\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}$$

Cosine similarity against all node semantic features:

$$s_i^{\text{sem}} = \frac{\mathbf{e}_q^{\text{sem}} \cdot X_i}{\|\mathbf{e}_q^{\text{sem}}\|\,\|X_i\|}$$

#### 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}\!\left(\mathbf{e}_q^{\text{sem}},\; \varnothing\right) \in \mathbb{R}^{128}$$

Cosine similarity against all precomputed structural embeddings $Z$:

$$s_i^{\text{struct}} = \frac{\mathbf{e}_q^{\text{struct}} \cdot Z_i}{\|\mathbf{e}_q^{\text{struct}}\|\,\|Z_i\|}$$

#### 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)} - \min_j s_j^{\,(\cdot)}}{\max_j s_j^{\,(\cdot)} - \min_j s_j^{\,(\cdot)}}$$

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

$$\boxed{f_i = \alpha\,\tilde{s}_i^{\text{sem}} + (1-\alpha)\,\tilde{s}_i^{\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.