Spaces:
Sleeping
Sleeping
| """ | |
| Knowledge Graph Builder | |
| ======================= | |
| Wraps a NetworkX DiGraph to construct, query, and export | |
| knowledge graphs from extracted entities and relationships. | |
| """ | |
| from typing import List, Dict, Any, Optional | |
| import networkx as nx | |
| from pyvis.network import Network | |
| import json | |
| # Colour palette for entity types | |
| ENTITY_COLORS = { | |
| "PERSON": "#00ff88", | |
| "ORG": "#00d4ff", | |
| "LOCATION": "#a855f7", | |
| "TECHNOLOGY": "#f59e0b", | |
| "DATE": "#ec4899", | |
| } | |
| ENTITY_SHAPES = { | |
| "PERSON": "dot", | |
| "ORG": "diamond", | |
| "LOCATION": "triangle", | |
| "TECHNOLOGY": "square", | |
| "DATE": "star", | |
| } | |
| class KnowledgeGraph: | |
| """ | |
| A directed knowledge graph backed by ``networkx.DiGraph``. | |
| """ | |
| def __init__(self): | |
| self.graph = nx.DiGraph() | |
| # ------------------------------------------------------------------ | |
| # Construction | |
| # ------------------------------------------------------------------ | |
| def add_entities(self, entities: List[Dict]) -> None: | |
| """ | |
| Add entity nodes to the graph. | |
| Parameters | |
| ---------- | |
| entities : list of dict | |
| Each dict must contain at least ``text`` and ``label`` keys. | |
| """ | |
| for ent in entities: | |
| node_id = ent["text"] | |
| if self.graph.has_node(node_id): | |
| # Increment mention count | |
| self.graph.nodes[node_id]["mentions"] = ( | |
| self.graph.nodes[node_id].get("mentions", 1) + 1 | |
| ) | |
| continue | |
| self.graph.add_node( | |
| node_id, | |
| label=ent["label"], | |
| color=ENTITY_COLORS.get(ent["label"], "#888888"), | |
| shape=ENTITY_SHAPES.get(ent["label"], "dot"), | |
| mentions=1, | |
| ) | |
| def add_relationships(self, relationships: List[Dict]) -> None: | |
| """ | |
| Add directed edges (relationships) to the graph. | |
| Parameters | |
| ---------- | |
| relationships : list of dict | |
| Each dict needs ``source``, ``target``, ``relation`` keys. | |
| """ | |
| for rel in relationships: | |
| src, tgt = rel["source"], rel["target"] | |
| # Ensure nodes exist | |
| if not self.graph.has_node(src): | |
| self.graph.add_node( | |
| src, | |
| label=rel.get("source_label", "UNKNOWN"), | |
| color=ENTITY_COLORS.get(rel.get("source_label"), "#888888"), | |
| shape=ENTITY_SHAPES.get(rel.get("source_label"), "dot"), | |
| mentions=1, | |
| ) | |
| if not self.graph.has_node(tgt): | |
| self.graph.add_node( | |
| tgt, | |
| label=rel.get("target_label", "UNKNOWN"), | |
| color=ENTITY_COLORS.get(rel.get("target_label"), "#888888"), | |
| shape=ENTITY_SHAPES.get(rel.get("target_label"), "dot"), | |
| mentions=1, | |
| ) | |
| if self.graph.has_edge(src, tgt): | |
| self.graph.edges[src, tgt]["weight"] = ( | |
| self.graph.edges[src, tgt].get("weight", 1) + 1 | |
| ) | |
| else: | |
| self.graph.add_edge( | |
| src, | |
| tgt, | |
| relation=rel["relation"], | |
| weight=1, | |
| sentence=rel.get("sentence", ""), | |
| ) | |
| # ------------------------------------------------------------------ | |
| # Queries & Analytics | |
| # ------------------------------------------------------------------ | |
| def get_stats(self) -> Dict[str, Any]: | |
| """Return summary statistics of the knowledge graph.""" | |
| G = self.graph | |
| label_counts: Dict[str, int] = {} | |
| for _, data in G.nodes(data=True): | |
| lbl = data.get("label", "UNKNOWN") | |
| label_counts[lbl] = label_counts.get(lbl, 0) + 1 | |
| relation_counts: Dict[str, int] = {} | |
| for _, _, data in G.edges(data=True): | |
| rel = data.get("relation", "UNKNOWN") | |
| relation_counts[rel] = relation_counts.get(rel, 0) + 1 | |
| communities = self.get_communities() | |
| # Degree centrality for top nodes | |
| if G.number_of_nodes() > 0: | |
| centrality = nx.degree_centrality(G) | |
| top_nodes = sorted(centrality.items(), key=lambda x: x[1], reverse=True)[:10] | |
| else: | |
| top_nodes = [] | |
| return { | |
| "total_nodes": G.number_of_nodes(), | |
| "total_edges": G.number_of_edges(), | |
| "entity_type_counts": label_counts, | |
| "relation_type_counts": relation_counts, | |
| "num_communities": len(communities), | |
| "communities": communities, | |
| "top_central_nodes": top_nodes, | |
| "density": nx.density(G) if G.number_of_nodes() > 1 else 0, | |
| } | |
| def get_communities(self) -> List[List[str]]: | |
| """ | |
| Detect communities using the greedy modularity algorithm | |
| on the undirected projection. | |
| """ | |
| if self.graph.number_of_nodes() == 0: | |
| return [] | |
| undirected = self.graph.to_undirected() | |
| try: | |
| from networkx.algorithms.community import greedy_modularity_communities | |
| communities = greedy_modularity_communities(undirected) | |
| return [sorted(list(c)) for c in communities] | |
| except Exception: | |
| # Fallback: connected components | |
| return [sorted(list(c)) for c in nx.connected_components(undirected)] | |
| def get_node_details(self, node_id: str) -> Optional[Dict]: | |
| """Return all attributes for a single node.""" | |
| if not self.graph.has_node(node_id): | |
| return None | |
| data = dict(self.graph.nodes[node_id]) | |
| data["id"] = node_id | |
| data["in_degree"] = self.graph.in_degree(node_id) | |
| data["out_degree"] = self.graph.out_degree(node_id) | |
| data["neighbors"] = list(self.graph.successors(node_id)) + list( | |
| self.graph.predecessors(node_id) | |
| ) | |
| return data | |
| # ------------------------------------------------------------------ | |
| # Export | |
| # ------------------------------------------------------------------ | |
| def to_pyvis(self, height: str = "600px", width: str = "100%") -> Network: | |
| """ | |
| Convert the graph to a PyVis ``Network`` for interactive | |
| HTML visualisation. | |
| """ | |
| net = Network( | |
| height=height, | |
| width=width, | |
| directed=True, | |
| bgcolor="#0a0a0a", | |
| font_color="white", | |
| select_menu=False, | |
| filter_menu=False, | |
| ) | |
| net.barnes_hut( | |
| gravity=-8000, | |
| central_gravity=0.3, | |
| spring_length=200, | |
| spring_strength=0.05, | |
| damping=0.09, | |
| ) | |
| for node_id, data in self.graph.nodes(data=True): | |
| mentions = data.get("mentions", 1) | |
| size = 15 + mentions * 5 | |
| net.add_node( | |
| node_id, | |
| label=node_id, | |
| color=data.get("color", "#888888"), | |
| shape=data.get("shape", "dot"), | |
| size=min(size, 50), | |
| title=f"{data.get('label', 'UNKNOWN')}\nMentions: {mentions}", | |
| font={"size": 14, "color": "white"}, | |
| ) | |
| for src, tgt, data in self.graph.edges(data=True): | |
| relation = data.get("relation", "") | |
| weight = data.get("weight", 1) | |
| net.add_edge( | |
| src, | |
| tgt, | |
| title=relation, | |
| label=relation, | |
| width=min(weight * 1.5, 6), | |
| color={"color": "#444444", "highlight": "#00ff88"}, | |
| font={"size": 10, "color": "#888888", "align": "middle"}, | |
| arrows={"to": {"enabled": True, "scaleFactor": 0.5}}, | |
| smooth={"type": "curvedCW", "roundness": 0.2}, | |
| ) | |
| return net | |
| def to_dict(self) -> Dict: | |
| """Serialise the graph to a JSON-safe dictionary.""" | |
| return nx.node_link_data(self.graph) | |
| def from_dict(self, data: Dict) -> None: | |
| """Load a graph from a dictionary produced by ``to_dict``.""" | |
| self.graph = nx.node_link_graph(data) | |