Spaces:
Sleeping
Sleeping
| 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. | |