""" Graph Builder cho Medical RAG Pipeline. Pipeline: 1. Load MedQuad dataset từ HuggingFace 2. Parse documents qua LlamaIndex PropertyGraphIndex 3. Chuyển đổi sang PyTorch Geometric graph - Nodes: mỗi text chunk = một node - Edges: cosine similarity > threshold (document similarity graph) 4. Trả về pyg_data, nodes, và raw embeddings Lý do dùng document similarity graph: Hai documents có nội dung liên quan (e.g. cùng nói về diabetes) sẽ được kết nối. GNN sẽ propagate thông tin qua các cạnh này, enriching embeddings với neighborhood context. """ import torch import torch.nn.functional as F from torch_geometric.data import Data from torch_geometric.utils import add_self_loops, coalesce from llama_index.core import PropertyGraphIndex, Document, Settings from llama_index.core.indices.property_graph import ImplicitPathExtractor from datasets import load_dataset from typing import List, Tuple def load_healthcare_documents(num_samples: int = 250) -> List[Document]: """ Load dataset y tế từ HuggingFace. Dataset: keivalya/MedQuad-MedicalQnADataset Mỗi Document = 1 câu trả lời y tế, với câu hỏi lưu trong metadata. Args: num_samples: số lượng documents cần load (giới hạn để tiết kiệm RAM) Returns: List[Document] - danh sách LlamaIndex Document objects """ print(f"Đang load {num_samples} medical documents từ HuggingFace...") dataset = load_dataset( "keivalya/MedQuad-MedicalQnADataset", split="train" ) num_samples = min(num_samples, len(dataset)) dataset = dataset.select(range(num_samples)) docs = [] for i, row in enumerate(dataset): # Truncate text dài để tiết kiệm RAM khi embedding answer_text = str(row.get("Answer", "")).strip()[:800] question_text = str(row.get("Question", "")).strip()[:200] if not answer_text: continue doc = Document( text=answer_text, metadata={ "question": question_text, "doc_id": str(i), "source": "MedQuad" } ) docs.append(doc) print(f"✓ Loaded {len(docs)} medical documents") return docs def build_llamaindex_property_graph( documents: List[Document], embed_model, ) -> Tuple[PropertyGraphIndex, List]: """ Xây dựng LlamaIndex PropertyGraphIndex từ documents. Sử dụng ImplicitPathExtractor - không cần LLM, phù hợp để deploy trên free tier. PropertyGraphIndex lưu trữ: - ChunkNodes: các đoạn văn bản từ documents - Implicit edges: dựa trên cấu trúc document Args: documents: danh sách LlamaIndex Document objects embed_model: HuggingFace embedding model Returns: (index, text_nodes) - PropertyGraphIndex và list TextNode objects """ print("Đang xây dựng LlamaIndex PropertyGraphIndex...") # Cấu hình Settings - không dùng LLM của LlamaIndex # vì sẽ dùng llama-cpp trực tiếp cho generation Settings.embed_model = embed_model Settings.llm = None index = PropertyGraphIndex.from_documents( documents, embed_model=embed_model, kg_extractors=[ImplicitPathExtractor()], # no LLM required show_progress=True, ) # Lấy tất cả text nodes từ docstore # index.docstore.docs = {id: TextNode} dict text_nodes = list(index.docstore.docs.values()) print(f"✓ PropertyGraphIndex built: {len(text_nodes)} text nodes") return index, text_nodes def build_pyg_graph( text_nodes: List, embed_model, similarity_threshold: float = 0.70, ) -> Tuple[Data, List, torch.Tensor]: """ Chuyển LlamaIndex text nodes sang PyTorch Geometric graph. Strategy xây dựng graph: - Mỗi text chunk = 1 node trong graph - Edge giữa 2 nodes nếu cosine_similarity > threshold - Thêm self-loops cho GCN stability (A_hat = A + I) Rationale: Documents về cùng chủ đề y tế (e.g. diabetes, hypertension) sẽ có high cosine similarity và được kết nối. GNN sẽ propagate thông tin qua các edges này, giúp mỗi node "biết" về context của các documents liên quan. Args: text_nodes: danh sách LlamaIndex TextNode objects embed_model: HuggingFace embedding model similarity_threshold: ngưỡng cosine similarity để tạo edge Returns: (pyg_data, nodes, raw_embeddings) pyg_data: PyTorch Geometric Data object nodes: ordered list of TextNode objects raw_embeddings: [N, D] tensor - raw text embeddings """ N = len(text_nodes) print(f"Đang embed {N} nodes và xây dựng PyG graph...") # ── Step 1: Embed all nodes ────────────────────────────────────────────── texts = [node.get_content() for node in text_nodes] embeddings = embed_model.get_text_embedding_batch(texts, show_progress=True) x = torch.tensor(embeddings, dtype=torch.float) # [N, D] # ── Step 2: Tính cosine similarity matrix ──────────────────────────────── x_norm = F.normalize(x, dim=-1) # [N, D] - L2 normalized sim_matrix = x_norm @ x_norm.T # [N, N] - cosine similarity # ── Step 3: Tạo edges từ similarity ────────────────────────────────────── # Loại bỏ diagonal (self-similarity = 1.0 không có ý nghĩa cấu trúc) mask = (sim_matrix > similarity_threshold) & \ ~torch.eye(N, dtype=torch.bool) edge_pairs = mask.nonzero(as_tuple=False) # [E, 2] if len(edge_pairs) > 0: edge_index = edge_pairs.T # [2, E] print(f" Tạo được {edge_index.shape[1]} edges (threshold={similarity_threshold})") else: # Fallback: nếu không có edges, kết nối mỗi node với top-1 neighbor print(f" Không có edges trên ngưỡng {similarity_threshold}, " f"dùng top-1 nearest neighbors...") sim_matrix.fill_diagonal_(-1.0) top1 = sim_matrix.argmax(dim=1) # [N] src = torch.arange(N) edge_index = torch.stack([src, top1]) # [2, N] # ── Step 4: Thêm self-loops và loại bỏ duplicate edges ─────────────────── # Self-loops quan trọng cho GCN: # Khi aggregate neighbors, node cũng aggregate chính nó edge_index, _ = add_self_loops(edge_index, num_nodes=N) edge_index = coalesce(edge_index, num_nodes=N) # remove duplicates # ── Step 5: Build PyG Data object ──────────────────────────────────────── pyg_data = Data(x=x, edge_index=edge_index, num_nodes=N) print(f"✓ PyG graph: {N} nodes, {edge_index.shape[1]} edges") print(f" Average degree: {edge_index.shape[1]/N:.1f}") return pyg_data, text_nodes, x