---
title: Lab3
emoji: 📊
colorFrom: indigo
colorTo: yellow
sdk: gradio
sdk_version: 6.13.0
app_file: app.py
pinned: false
---

# GNN-based Medical RAG System

> Healthcare Question Answering using Graph Neural Network-enhanced Retrieval Augmented Generation.

---

## Table of Contents

1. [Algorithm Overview](#algorithm-overview)
2. [Pipeline Steps](#pipeline-steps)
   - [Dataset](#1-dataset)
   - [Graph Index Construction](#2-llamaindex-property-graph-construction)
   - [PyTorch Geometric Graph](#3-pytorch-geometric-graph-construction)
   - [Graph Neural Network](#4-graph-neural-network-gcn)
   - [Hybrid Retrieval](#5-hybrid-retrieval-gnnhybridretriever)
   - [Generation](#6-generation-qwen35-4b-neo-gguf)
3. [System Architecture](#system-architecture)
4. [Repository Structure](#repository-structure)
5. [References](#references)

---

## Algorithm Overview

This system implements a GNN-based RAG pipeline for medical QA, combining **semantic retrieval** with **structural graph reasoning**. Each query is answered by retrieving relevant documents through a hybrid scoring mechanism that blends raw embedding similarity with GNN-enriched structural embeddings.

---

## Pipeline Steps

### 1. Dataset

| Property | Value |
|---|---|
| **Source** | `keivalya/MedQuad-MedicalQnADataset` (HuggingFace) |
| **Size** | 250 answer documents |
| **Truncation** | 800 characters per document |

MedQuad contains medical question-answer pairs from trusted sources including **NIH**, **CDC**, and **NLM**. Each document is converted to a LlamaIndex `Document` object, with the answer as the main text and the original question stored as metadata.

---

### 2. LlamaIndex Property Graph Construction

We use **LlamaIndex `PropertyGraphIndex`** with `ImplicitPathExtractor` to parse documents into a structured graph index:

```python
index = PropertyGraphIndex.from_documents(
    documents,
    embed_model=embed_model,
    kg_extractors=[ImplicitPathExtractor()],
)
```

`ImplicitPathExtractor` requires no LLM — it creates implicit node relationships based on document structure. Text nodes are extracted from `index.docstore.docs` for downstream graph construction.

---

### 3. PyTorch Geometric Graph Construction

LlamaIndex text nodes are converted to a **PyTorch Geometric** `Data` object:

- **Nodes** — Each text chunk is one graph node.  
  Node features `x ∈ ℝ^(N×384)` are the BGE-Small embeddings.

- **Edges** — Two nodes are connected if their cosine similarity exceeds the threshold:

$$\text{edge}(i, j) \text{ exists} \iff \text{cosine\_sim}(\text{embed}_i,\ \text{embed}_j) > 0.70$$

Self-loops are added for GCN stability: $\tilde{A} = A + I$

---

### 4. Graph Neural Network (GCN)

A **2-layer Graph Convolutional Network** following [Kipf & Welling (2017)](#references) enriches node embeddings with structural context.

**Architecture:**

```
x [N, 384] → GCNConv(384→256) → ReLU → Dropout(0.3)
           → GCNConv(256→384) → structural_embs [N, 384]
```

**GCN Update Rule:**

$$H^{(l+1)} = \sigma\!\left(\tilde{D}^{-1/2}\,\tilde{A}\,\tilde{D}^{-1/2}\,H^{(l)}\,W^{(l)}\right)$$

where $\tilde{A} = A + I$ (adjacency + self-loops) and $\tilde{D}$ is the corresponding degree matrix.

> **Key insight:** The GNN acts as a *structural feature transformer*, not a classifier. No training is required — the graph topology itself provides the structural signal. After 2 message-passing layers, each node's embedding captures information from its **2-hop neighborhood**.

---

### 5. Hybrid Retrieval (GNNHybridRetriever)

We subclass LlamaIndex's `BaseRetriever` to implement dual-score hybrid retrieval.

| Score | Formula | Description |
|---|---|---|
| **Semantic** | $\text{sem}(q, i) = \cos(\text{embed}(q),\ \text{raw\_embed}[i])$ | Pure text similarity |
| **Structural** | $\text{struct}(q, i) = \cos(\text{embed}(q),\ \text{gnn\_embed}[i])$ | Neighborhood-aware similarity |
| **Hybrid** | $\text{score}(q, i) = \alpha \cdot \text{sem}(q, i) + (1-\alpha) \cdot \text{struct}(q, i)$ | Combined score |

Default $\alpha = 0.6$ (60% semantic, 40% structural).

**Why hybrid is better:**

| Method | Weakness |
|---|---|
| Semantic only | Misses documents with different terminology |
| Structural only | May retrieve irrelevant neighbors of relevant nodes |
| **Hybrid** | Robust to vocabulary mismatch while leveraging graph topology ✓ |

The top-K nodes by hybrid score are returned as context for generation.

---

### 6. Generation (Qwen3.5-4B-Neo GGUF)

We use **llama-cpp-python** to run `Qwen3.5-4B-Neo-GGUF` (Q4_K_M quantization) on CPU for efficient inference on free-tier hardware:

```python
llm = Llama(
    model_path=model_path,
    n_ctx=2048,
    n_threads=4,
    n_gpu_layers=0,   # CPU only
)
```

The prompt template explicitly constrains the model to answer only from retrieved context, reducing hallucination risk.

---

## System Architecture

```
MedQuad Dataset (250 docs)
│
▼
LlamaIndex PropertyGraphIndex
(ImplicitPathExtractor)
│
▼
Text Nodes → BGE-Small Embeddings [N × 384]
│
├─── Cosine Similarity → PyG Graph ──► 2-layer GCN
│    (edges where sim > 0.70)               │
│                                Structural Embeddings [N × 384]
│                                           │
└────────────── GNNHybridRetriever ─────────┘
       α × semantic + (1-α) × structural
                     │
               Top-K Context Nodes
                     │
          Qwen3.5-4B-Neo (GGUF, CPU)
                     │
                   Answer
```

---

## Repository Structure

```
├── app.py                  # Main Gradio application
├── requirements.txt        # Python dependencies
├── README.md               # This file (also serves as the lab report)
└── utils/
    ├── __init__.py
    ├── gnn_model.py        # MedicalGNN (2-layer GCN) + get_structural_embeddings()
    ├── graph_builder.py    # Dataset loading, PropertyGraphIndex, PyG conversion
    └── retriever.py        # GNNHybridRetriever (BaseRetriever subclass)
```

---

## References

1. Lewis, P. et al. (2020). *Retrieval Augmented Generation for Knowledge-Intensive NLP Tasks.* NeurIPS 2020, 33, 9459–9474.

2. Edge, D. et al. (2024). *From Local to Global: A Graph RAG Approach to Query-Focused Summarization.* [arXiv:2404.16130](https://arxiv.org/abs/2404.16130).

3. Kipf, T. N. & Welling, M. (2017). *Semi-Supervised Classification with Graph Convolutional Networks.* ICLR 2017.