| """ |
| AETHER Knowledge Graph Engine. |
| Integrates PyTorch Geometric patterns for relational reasoning: |
| - RGCN for node classification on knowledge graphs |
| - ComplEx for link prediction |
| - Neuro-symbolic bridge: learned attention over symbolic rules |
| """ |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from typing import Dict, List, Any, Optional, Tuple |
| import networkx as nx |
| import numpy as np |
| import logging |
|
|
| logger = logging.getLogger("AETHER.Knowledge") |
|
|
|
|
| class RGCNLayer(nn.Module): |
| """Simplified RGCN layer for knowledge graph reasoning.""" |
| |
| def __init__(self, in_dim: int, out_dim: int, num_relations: int, |
| num_bases: int = 4): |
| super().__init__() |
| self.in_dim = in_dim |
| self.out_dim = out_dim |
| self.num_relations = num_relations |
| self.num_bases = num_bases |
| |
| self.bases = nn.Parameter(torch.Tensor(num_bases, in_dim, out_dim)) |
| self.comp = nn.Parameter(torch.Tensor(num_relations, num_bases)) |
| self.self_loop = nn.Parameter(torch.Tensor(in_dim, out_dim)) |
| self.bias = nn.Parameter(torch.Tensor(out_dim)) |
| |
| self.reset_parameters() |
| |
| def reset_parameters(self): |
| nn.init.xavier_uniform_(self.bases) |
| nn.init.xavier_uniform_(self.comp) |
| nn.init.xavier_uniform_(self.self_loop) |
| nn.init.zeros_(self.bias) |
| |
| def forward(self, x: Optional[torch.Tensor], edge_index: torch.Tensor, |
| edge_type: torch.Tensor) -> torch.Tensor: |
| num_nodes = int(edge_index.max().item()) + 1 if x is None else x.size(0) |
| |
| if x is None: |
| x = torch.eye(num_nodes, self.in_dim, device=edge_index.device) |
| |
| weight = torch.einsum('rb, bio -> rio', self.comp, self.bases) |
| |
| out = torch.zeros(num_nodes, self.out_dim, device=x.device) |
| |
| for rel_id in range(self.num_relations): |
| mask = edge_type == rel_id |
| if mask.sum() == 0: |
| continue |
| |
| rel_edges = edge_index[:, mask] |
| source = rel_edges[0] |
| target = rel_edges[1] |
| |
| messages = torch.mm(x[source], weight[rel_id]) |
| out.index_add_(0, target, messages) |
| |
| out = out + torch.mm(x, self.self_loop) |
| out = out + self.bias |
| return out |
|
|
|
|
| class KnowledgeGraphEncoder(nn.Module): |
| """Multi-layer RGCN encoder for knowledge graph embeddings.""" |
| |
| def __init__(self, num_nodes: int, hidden_dim: int, num_relations: int, |
| num_layers: int = 2, num_bases: int = 4): |
| super().__init__() |
| self.num_nodes = num_nodes |
| self.hidden_dim = hidden_dim |
| self.num_relations = num_relations |
| |
| self.node_embeddings = nn.Embedding(num_nodes, hidden_dim) |
| |
| self.layers = nn.ModuleList([ |
| RGCNLayer( |
| in_dim=hidden_dim if i == 0 else hidden_dim, |
| out_dim=hidden_dim, |
| num_relations=num_relations, |
| num_bases=num_bases, |
| ) |
| for i in range(num_layers) |
| ]) |
| |
| self.norms = nn.ModuleList([ |
| nn.LayerNorm(hidden_dim) |
| for _ in range(num_layers) |
| ]) |
| |
| def forward(self, edge_index: torch.Tensor, |
| edge_type: torch.Tensor) -> torch.Tensor: |
| num_nodes = int(edge_index.max().item()) + 1 |
| x = self.node_embeddings(torch.arange(num_nodes, device=edge_index.device)) |
| |
| for layer, norm in zip(self.layers, self.norms): |
| x_new = layer(x, edge_index, edge_type) |
| x_new = F.relu(norm(x_new)) |
| x = x_new |
| |
| return x |
|
|
|
|
| class ComplExScorer(nn.Module): |
| """ComplEx link prediction scorer for knowledge graph completion.""" |
| |
| def __init__(self, num_nodes: int, num_relations: int, hidden_dim: int = 50): |
| super().__init__() |
| self.num_nodes = num_nodes |
| self.num_relations = num_relations |
| self.hidden_dim = hidden_dim |
| |
| self.head_real = nn.Embedding(num_nodes, hidden_dim) |
| self.head_imag = nn.Embedding(num_nodes, hidden_dim) |
| self.tail_real = nn.Embedding(num_nodes, hidden_dim) |
| self.tail_imag = nn.Embedding(num_nodes, hidden_dim) |
| |
| self.rel_real = nn.Embedding(num_relations, hidden_dim) |
| self.rel_imag = nn.Embedding(num_relations, hidden_dim) |
| |
| self.reset_parameters() |
| |
| def reset_parameters(self): |
| for param in self.parameters(): |
| nn.init.xavier_uniform_(param) |
| |
| def forward(self, head_idx: torch.Tensor, rel_idx: torch.Tensor, |
| tail_idx: torch.Tensor) -> torch.Tensor: |
| hr = self.head_real(head_idx) |
| hi = self.head_imag(head_idx) |
| tr = self.tail_real(tail_idx) |
| ti = self.tail_imag(tail_idx) |
| rr = self.rel_real(rel_idx) |
| ri = self.rel_imag(rel_idx) |
| |
| score = torch.sum( |
| hr * rr * tr + hr * ri * ti + hi * rr * ti - hi * ri * tr, |
| dim=-1 |
| ) |
| return score |
| |
| def loss(self, head_idx: torch.Tensor, rel_idx: torch.Tensor, |
| tail_idx: torch.Tensor, negative_head: torch.Tensor = None, |
| negative_tail: torch.Tensor = None) -> torch.Tensor: |
| pos_score = self.forward(head_idx, rel_idx, tail_idx) |
| |
| if negative_head is not None: |
| neg_score = self.forward(negative_head, rel_idx, tail_idx) |
| elif negative_tail is not None: |
| neg_score = self.forward(head_idx, rel_idx, negative_tail) |
| else: |
| neg_tail = torch.randint(0, self.num_nodes, tail_idx.size(), |
| device=tail_idx.device) |
| neg_score = self.forward(head_idx, rel_idx, neg_tail) |
| |
| pos_loss = F.softplus(-pos_score) |
| neg_loss = F.softplus(neg_score) |
| |
| return (pos_loss + neg_loss).mean() |
|
|
|
|
| class KnowledgeGraphEngine(nn.Module): |
| """ |
| Unified knowledge graph engine combining: |
| - NetworkX for graph construction and symbolic reasoning |
| - RGCN for learned embeddings |
| - ComplEx for link prediction |
| - Neuro-symbolic bridge for AETHER integration |
| """ |
| |
| def __init__(self, embedding_dim: int = 128, num_relations: int = 20, |
| max_nodes: int = 10000): |
| super().__init__() |
| self.embedding_dim = embedding_dim |
| self.num_relations = num_relations |
| self.max_nodes = max_nodes |
| |
| self.graph = nx.DiGraph() |
| self.node_id_map: Dict[str, int] = {} |
| self.relation_map: Dict[str, int] = {} |
| self.next_node_id = 0 |
| self.next_rel_id = 0 |
| |
| self.encoder: Optional[KnowledgeGraphEncoder] = None |
| self.scorer: Optional[ComplExScorer] = None |
| |
| self.symbolic_attention = nn.Parameter(torch.ones(num_relations)) |
| self.rules: List[Tuple[str, str, str]] = [] |
| |
| def _get_or_create_node(self, node_name: str) -> int: |
| if node_name not in self.node_id_map: |
| self.node_id_map[node_name] = self.next_node_id |
| self.graph.add_node(self.next_node_id, name=node_name) |
| self.next_node_id += 1 |
| return self.node_id_map[node_name] |
| |
| def _get_or_create_relation(self, relation: str) -> int: |
| if relation not in self.relation_map: |
| self.relation_map[relation] = self.next_rel_id |
| self.next_rel_id += 1 |
| return self.relation_map[relation] |
| |
| def add_fact(self, head: str, relation: str, tail: str, |
| confidence: float = 1.0): |
| h_id = self._get_or_create_node(head) |
| t_id = self._get_or_create_node(tail) |
| r_id = self._get_or_create_relation(relation) |
| |
| self.graph.add_edge(h_id, t_id, relation=r_id, name=relation, |
| confidence=confidence) |
| self._ensure_model_capacity() |
| |
| def add_rule(self, premise: Tuple[str, str, str], |
| conclusion: Tuple[str, str, str]): |
| self.rules.append((premise, conclusion)) |
| |
| def _ensure_model_capacity(self): |
| if self.encoder is None and self.next_node_id > 0: |
| num_nodes = min(self.next_node_id, self.max_nodes) |
| num_rels = max(self.next_rel_id, self.num_relations) |
| |
| self.encoder = KnowledgeGraphEncoder( |
| num_nodes=num_nodes, |
| hidden_dim=self.embedding_dim, |
| num_relations=num_rels, |
| num_layers=2, |
| ) |
| |
| self.scorer = ComplExScorer( |
| num_nodes=num_nodes, |
| num_relations=num_rels, |
| hidden_dim=self.embedding_dim // 2, |
| ) |
| |
| logger.info(f"Initialized KG models: {num_nodes} nodes, {num_rels} relations") |
| |
| def reason_symbolic(self, query_head: str, query_relation: str) -> List[Dict]: |
| results = [] |
| |
| if query_head not in self.node_id_map: |
| return results |
| |
| h_id = self.node_id_map[query_head] |
| r_name = query_relation |
| |
| if r_name in self.relation_map: |
| r_id = self.relation_map[r_name] |
| for _, target, data in self.graph.out_edges(h_id, data=True): |
| if data.get('relation') == r_id: |
| results.append({ |
| "head": query_head, |
| "relation": r_name, |
| "tail": self.graph.nodes[target].get('name', str(target)), |
| "confidence": data.get('confidence', 1.0), |
| "path": "direct", |
| }) |
| |
| for premise, conclusion in self.rules: |
| p_head, p_rel, p_tail = premise |
| c_head, c_rel, c_tail = conclusion |
| |
| if p_head == query_head and self._check_fact(premise): |
| inferred_tail = c_tail |
| if c_head == "?": |
| c_head = query_head |
| |
| results.append({ |
| "head": c_head, |
| "relation": c_rel, |
| "tail": inferred_tail, |
| "confidence": 0.8, |
| "path": "inferred", |
| "rule": f"{premise} -> {conclusion}", |
| }) |
| |
| for neighbor in nx.bfs_tree(self.graph, h_id, depth_limit=2).nodes(): |
| if neighbor != h_id: |
| for path in nx.all_simple_paths(self.graph, h_id, neighbor, cutoff=2): |
| if len(path) > 1: |
| edge_data = self.graph.edges[path[0], path[1]] |
| results.append({ |
| "head": query_head, |
| "relation": f"multi-hop via {edge_data.get('name', 'unknown')}", |
| "tail": self.graph.nodes[neighbor].get('name', str(neighbor)), |
| "confidence": 0.6 ** (len(path) - 1), |
| "path": "->".join(str(n) for n in path), |
| }) |
| |
| return sorted(results, key=lambda x: x["confidence"], reverse=True) |
| |
| def _check_fact(self, fact: Tuple[str, str, str]) -> bool: |
| h, r, t = fact |
| if h not in self.node_id_map or t not in self.node_id_map: |
| return False |
| |
| h_id = self.node_id_map[h] |
| t_id = self.node_id_map[t] |
| |
| if r not in self.relation_map: |
| return False |
| |
| r_id = self.relation_map[r] |
| return self.graph.has_edge(h_id, t_id) and \ |
| self.graph.edges[h_id, t_id].get('relation') == r_id |
| |
| def reason_learned(self, query_head: str, query_relation: str, |
| top_k: int = 5) -> List[Dict]: |
| if self.scorer is None or query_head not in self.node_id_map: |
| return [] |
| |
| h_id = self.node_id_map[query_head] |
| r_id = self.relation_map.get(query_relation) |
| |
| if r_id is None: |
| return [] |
| |
| h_tensor = torch.tensor([h_id]) |
| r_tensor = torch.tensor([r_id]) |
| |
| all_tails = torch.arange(self.scorer.num_nodes) |
| scores = [] |
| |
| batch_size = 1000 |
| for i in range(0, len(all_tails), batch_size): |
| batch_tails = all_tails[i:i + batch_size] |
| h_batch = h_tensor.repeat(len(batch_tails)) |
| r_batch = r_tensor.repeat(len(batch_tails)) |
| |
| batch_scores = self.scorer(h_batch, r_batch, batch_tails) |
| scores.extend(batch_scores.tolist()) |
| |
| scores = torch.tensor(scores) |
| top_scores, top_indices = torch.topk(scores, min(top_k, len(scores))) |
| |
| results = [] |
| for idx, score in zip(top_indices, top_scores): |
| node_name = self.graph.nodes[idx.item()].get('name', str(idx.item())) |
| results.append({ |
| "head": query_head, |
| "relation": query_relation, |
| "tail": node_name, |
| "confidence": torch.sigmoid(score).item(), |
| "path": "learned", |
| }) |
| |
| return results |
| |
| def query(self, text_query: str, top_k: int = 5) -> Dict[str, Any]: |
| parts = text_query.lower().split() |
| |
| if len(parts) >= 2: |
| head = parts[0].capitalize() |
| relation = " ".join(parts[1:]) |
| else: |
| head = text_query.capitalize() |
| relation = "related_to" |
| |
| symbolic_results = self.reason_symbolic(head, relation) |
| learned_results = self.reason_learned(head, relation, top_k) |
| |
| rel_id = self.relation_map.get(relation, 0) |
| symbolic_weight = torch.sigmoid(self.symbolic_attention[rel_id % self.num_relations]) |
| learned_weight = 1.0 - symbolic_weight.item() |
| |
| all_results = [] |
| |
| for r in symbolic_results[:top_k]: |
| r["source"] = "symbolic" |
| r["fusion_weight"] = symbolic_weight.item() |
| all_results.append(r) |
| |
| for r in learned_results[:top_k]: |
| r["source"] = "learned" |
| r["fusion_weight"] = learned_weight |
| all_results.append(r) |
| |
| all_results.sort(key=lambda x: x.get("confidence", 0), reverse=True) |
| |
| return { |
| "query": text_query, |
| "results": all_results[:top_k], |
| "symbolic_weight": symbolic_weight.item(), |
| "learned_weight": learned_weight, |
| "num_symbolic": len(symbolic_results), |
| "num_learned": len(learned_results), |
| } |
| |
| def to_pyg_data(self) -> Dict[str, torch.Tensor]: |
| edges = [] |
| edge_types = [] |
| |
| for u, v, data in self.graph.edges(data=True): |
| edges.append([u, v]) |
| edge_types.append(data.get('relation', 0)) |
| |
| if not edges: |
| return {} |
| |
| edge_index = torch.tensor(edges, dtype=torch.long).t() |
| edge_type = torch.tensor(edge_types, dtype=torch.long) |
| |
| return { |
| "edge_index": edge_index, |
| "edge_type": edge_type, |
| "num_nodes": self.next_node_id, |
| "num_relations": self.next_rel_id, |
| } |
| |
| def stats(self) -> Dict[str, Any]: |
| return { |
| "num_nodes": self.graph.number_of_nodes(), |
| "num_edges": self.graph.number_of_edges(), |
| "num_relations": len(self.relation_map), |
| "num_rules": len(self.rules), |
| "node_names": len(self.node_id_map), |
| } |
| |
| def export(self) -> Dict[str, Any]: |
| edges = [] |
| for u, v, data in self.graph.edges(data=True): |
| edges.append({ |
| "source": u, |
| "target": v, |
| "relation_id": data.get('relation'), |
| "relation_name": data.get('name'), |
| "confidence": data.get('confidence'), |
| }) |
| |
| return { |
| "nodes": {n: self.graph.nodes[n].get('name', str(n)) |
| for n in self.graph.nodes()}, |
| "edges": edges, |
| "node_id_map": self.node_id_map, |
| "relation_map": self.relation_map, |
| "rules": self.rules, |
| } |
|
|
