""" GraphRAG Builder with Community Detection and Hierarchical Summarization Implements Microsoft GraphRAG approach for knowledge graphs """ import json import networkx as nx import numpy as np from typing import List, Dict, Set, Tuple from collections import defaultdict, Counter from tqdm import tqdm from sklearn.feature_extraction.text import TfidfVectorizer from sklearn.metrics.pairwise import cosine_similarity import pickle class GraphRAGBuilder: """Build GraphRAG with community detection and hierarchical summaries""" def __init__(self, chunks_file: str, output_dir: str = "dataset/wikipedia_ireland"): self.chunks_file = chunks_file self.output_dir = output_dir self.graph = nx.Graph() self.entity_graph = nx.DiGraph() self.chunks = [] self.entity_to_chunks = defaultdict(list) self.chunk_to_entities = defaultdict(list) def load_chunks(self): """Load processed chunks""" print(f"[INFO] Loading chunks from {self.chunks_file}") with open(self.chunks_file, 'r', encoding='utf-8') as f: self.chunks = json.load(f) print(f"[SUCCESS] Loaded {len(self.chunks)} chunks") def build_entity_graph(self): """Build graph from entities across chunks""" print("[INFO] Building entity graph from chunks...") # Extract all entities and their co-occurrences for chunk_idx, chunk in enumerate(tqdm(self.chunks, desc="Processing chunks")): chunk_id = chunk['chunk_id'] entities = chunk.get('entities', []) # Track which chunks contain which entities for entity in entities: entity_key = f"{entity['text']}|{entity['label']}" self.entity_to_chunks[entity_key].append(chunk_id) self.chunk_to_entities[chunk_id].append(entity_key) # Add entity as node if not exists if not self.entity_graph.has_node(entity_key): self.entity_graph.add_node( entity_key, text=entity['text'], label=entity['label'], chunk_count=0 ) # Update chunk count self.entity_graph.nodes[entity_key]['chunk_count'] += 1 # Create edges between co-occurring entities in same chunk for i, entity1 in enumerate(entities): for entity2 in entities[i+1:]: key1 = f"{entity1['text']}|{entity1['label']}" key2 = f"{entity2['text']}|{entity2['label']}" if self.entity_graph.has_edge(key1, key2): self.entity_graph[key1][key2]['weight'] += 1 else: self.entity_graph.add_edge(key1, key2, weight=1) print(f"[SUCCESS] Entity graph: {self.entity_graph.number_of_nodes()} nodes, " f"{self.entity_graph.number_of_edges()} edges") def build_semantic_chunk_graph(self, similarity_threshold: float = 0.3): """Build graph of semantically similar chunks""" print("[INFO] Building semantic similarity graph...") # Extract chunk texts chunk_texts = [chunk['text'] for chunk in self.chunks] chunk_ids = [chunk['chunk_id'] for chunk in self.chunks] # Compute TF-IDF vectors vectorizer = TfidfVectorizer(max_features=1000, stop_words='english') tfidf_matrix = vectorizer.fit_transform(chunk_texts) # Compute pairwise cosine similarity (in batches to save memory) batch_size = 500 for i in tqdm(range(0, len(chunk_texts), batch_size), desc="Computing similarity"): end_i = min(i + batch_size, len(chunk_texts)) batch_similarities = cosine_similarity(tfidf_matrix[i:end_i], tfidf_matrix) for local_idx, chunk_idx in enumerate(range(i, end_i)): chunk_id = chunk_ids[chunk_idx] # Add chunk as node if not self.graph.has_node(chunk_id): self.graph.add_node( chunk_id, text=chunk_texts[chunk_idx], source_title=self.chunks[chunk_idx]['source_title'], source_url=self.chunks[chunk_idx]['source_url'], section=self.chunks[chunk_idx]['section'], word_count=self.chunks[chunk_idx]['word_count'] ) # Add edges to similar chunks for other_idx, similarity in enumerate(batch_similarities[local_idx]): if chunk_idx != other_idx and similarity > similarity_threshold: other_chunk_id = chunk_ids[other_idx] if not self.graph.has_edge(chunk_id, other_chunk_id): self.graph.add_edge(chunk_id, other_chunk_id, weight=float(similarity)) print(f"[SUCCESS] Chunk graph: {self.graph.number_of_nodes()} nodes, " f"{self.graph.number_of_edges()} edges") def detect_communities(self, resolution: float = 1.0) -> Dict[str, int]: """Detect communities using Louvain algorithm""" print("[INFO] Detecting communities with Louvain algorithm...") from networkx.algorithms import community as nx_comm # Use Louvain for community detection communities = nx_comm.louvain_communities(self.graph, resolution=resolution, seed=42) # Create node to community mapping node_to_community = {} for comm_id, community_nodes in enumerate(communities): for node in community_nodes: node_to_community[node] = comm_id print(f"[SUCCESS] Detected {len(communities)} communities") # Add community attribute to nodes for node, comm_id in node_to_community.items(): self.graph.nodes[node]['community'] = comm_id return node_to_community def generate_community_summaries(self, node_to_community: Dict[str, int], max_chunks_per_summary: int = 20) -> Dict[int, Dict]: """Generate hierarchical summaries for each community""" print("[INFO] Generating community summaries...") communities = defaultdict(list) for node, comm_id in node_to_community.items(): communities[comm_id].append(node) community_summaries = {} for comm_id, chunk_ids in tqdm(communities.items(), desc="Summarizing communities"): # Gather all text from chunks in this community (limit to avoid huge summaries) sample_chunk_ids = chunk_ids[:max_chunks_per_summary] chunk_texts = [] sources = set() for chunk_id in sample_chunk_ids: chunk_data = self.graph.nodes.get(chunk_id, {}) chunk_texts.append(chunk_data.get('text', '')) sources.add(chunk_data.get('source_title', 'Unknown')) # Extract most common entities in this community community_entities = [] for chunk_id in chunk_ids: community_entities.extend(self.chunk_to_entities.get(chunk_id, [])) entity_counter = Counter(community_entities) top_entities = entity_counter.most_common(20) # Generate summary metadata (would use LLM for actual summary in production) combined_text = " ".join(chunk_texts) summary = { "community_id": comm_id, "num_chunks": len(chunk_ids), "num_sources": len(sources), "sources": list(sources)[:10], "top_entities": [{"entity": ent[0].split('|')[0], "count": ent[1]} for ent in top_entities], "combined_text_sample": combined_text[:2000], # First 2000 chars as preview "total_text_length": len(combined_text), "chunk_ids": chunk_ids[:100] # Limit stored chunk IDs } community_summaries[comm_id] = summary print(f"[SUCCESS] Generated {len(community_summaries)} community summaries") return community_summaries def build_hierarchical_index(self) -> Dict: """Build complete hierarchical index for GraphRAG""" print("=" * 80) print("BUILDING GRAPHRAG HIERARCHICAL INDEX") print("=" * 80) # Step 1: Load chunks self.load_chunks() # Step 2: Build entity graph self.build_entity_graph() # Step 3: Build semantic chunk graph self.build_semantic_chunk_graph(similarity_threshold=0.25) # Step 4: Detect communities node_to_community = self.detect_communities(resolution=1.0) # Step 5: Generate community summaries community_summaries = self.generate_community_summaries(node_to_community) # Step 6: Build complete index graphrag_index = { "metadata": { "total_chunks": len(self.chunks), "total_entities": self.entity_graph.number_of_nodes(), "total_communities": len(set(node_to_community.values())), "chunk_graph_edges": self.graph.number_of_edges(), "entity_graph_edges": self.entity_graph.number_of_edges() }, "communities": community_summaries, "entity_to_chunks": dict(self.entity_to_chunks), "chunk_to_entities": dict(self.chunk_to_entities), "node_to_community": node_to_community } return graphrag_index def save_graphrag_index(self, graphrag_index: Dict): """Save GraphRAG index and graphs""" print("[INFO] Saving GraphRAG index...") # Save main index as JSON index_path = f"{self.output_dir}/graphrag_index.json" with open(index_path, 'w', encoding='utf-8') as f: json.dump(graphrag_index, f, ensure_ascii=False, indent=2) print(f"[SUCCESS] Saved GraphRAG index to {index_path}") # Save graphs as pickle (more efficient for networkx graphs) graphs_path = f"{self.output_dir}/graphrag_graphs.pkl" with open(graphs_path, 'wb') as f: pickle.dump({ 'chunk_graph': self.graph, 'entity_graph': self.entity_graph }, f) print(f"[SUCCESS] Saved graphs to {graphs_path}") # Save human-readable statistics stats = { "total_chunks": graphrag_index["metadata"]["total_chunks"], "total_entities": graphrag_index["metadata"]["total_entities"], "total_communities": graphrag_index["metadata"]["total_communities"], "communities": [] } for comm_id, comm_data in graphrag_index["communities"].items(): stats["communities"].append({ "id": comm_id, "num_chunks": comm_data["num_chunks"], "num_sources": comm_data["num_sources"], "top_sources": comm_data["sources"][:5], "top_entities": [e["entity"] for e in comm_data["top_entities"][:10]] }) stats_path = f"{self.output_dir}/graphrag_stats.json" with open(stats_path, 'w') as f: json.dump(stats, f, indent=2) print(f"[SUCCESS] Saved statistics to {stats_path}") print("=" * 80) print("GRAPHRAG INDEX BUILDING COMPLETE!") print("=" * 80) if __name__ == "__main__": builder = GraphRAGBuilder( chunks_file="dataset/wikipedia_ireland/chunks.json" ) graphrag_index = builder.build_hierarchical_index() builder.save_graphrag_index(graphrag_index)