""" SAULSGPT — LAYER 6: RELATIONAL GNN & PATHWAY RANKING ENGINE ============================================================ Lightweight, high-performance PyTorch Relational Knowledge Graph Embedder and Multi-Hop Judicial Pathway Ranker. What it does: 1. Loads 160+ ground-truth Indian law triples (h, r, t) from legal_triples.json. 2. Embeds entities and relations into continuous vector space (TransE / RotatE energy model). 3. Performs Multi-Hop Subgraph Traversal to discover parallel judicial remedies: - Revenue / SDM Restitution Route - Criminal Magistrate FIR Escalation Route - Civil Injunction Stay Order Triad Route - High Court Inherent Quashing / Writ Route 4. Injects scored topological pathways directly into Layer 3 Reasoning. """ import json import os import torch import torch.nn as nn from typing import List, Dict, Any, Tuple TRIPLES_PATH = os.path.join(os.path.dirname(__file__), "data", "legal_triples.json") # Relation Importance Weights RELATION_WEIGHTS = { "overrides": 1.0, "struck_down_free_speech": 1.0, "struck_down_gender_equality": 1.0, "struck_down_mandatory_death": 1.0, "stayed_and_omitted_in_BNS": 0.98, "replaced_by": 0.95, "remedy_under": 0.95, "empowers_SDM": 0.92, "preserves_perishable_crops": 0.92, "apply_to_Magistrate": 0.92, "escalate_to_SP": 0.90, "quashing_petition": 0.95, "temporary_injunction": 0.92, "requires": 0.90, "supreme_court_benchmark": 0.95, "requires_15_day_legal_notice": 0.92, "30_day_magistrate_complaint_limitation": 0.90, "void_in_restraint_of_trade": 0.95, "void_ab_initio_TRF_Perkins": 0.95 } class RelationalGNNEngine(nn.Module): """ Relational Knowledge Graph Embedding & Pathway Discovery Engine. """ def __init__(self, embedding_dim: int = 64): super().__init__() self.embedding_dim = embedding_dim self.entity2id = {} self.id2entity = {} self.relation2id = {} self.id2relation = {} self.triples = [] self.adj = {} self._load_triples() self._build_index() num_entities = max(len(self.entity2id), 1) num_relations = max(len(self.relation2id), 1) self.entity_embeddings = nn.Embedding(num_entities, self.embedding_dim) self.relation_embeddings = nn.Embedding(num_relations, self.embedding_dim) nn.init.xavier_uniform_(self.entity_embeddings.weight) nn.init.xavier_uniform_(self.relation_embeddings.weight) weights_path = os.path.join(os.path.dirname(__file__), "data", "gnn_weights.pt") if os.path.exists(weights_path): try: ckpt = torch.load(weights_path) saved_e = ckpt["entity_embeddings"]["weight"] saved_r = ckpt["relation_embeddings"]["weight"] # Copy overlapping entity weights min_e = min(num_entities, saved_e.shape[0]) self.entity_embeddings.weight.data[:min_e] = saved_e[:min_e] # Copy overlapping relation weights min_r = min(num_relations, saved_r.shape[0]) self.relation_embeddings.weight.data[:min_r] = saved_r[:min_r] print(f"[GNN Engine] Loaded pre-trained embeddings ({min_e}/{num_entities} entities, {min_r}/{num_relations} relations preserved).") except Exception as e: print(f"[GNN Engine] Initialized with Xavier uniform ({e}).") def _load_triples(self): if os.path.exists(TRIPLES_PATH): with open(TRIPLES_PATH, "r", encoding="utf-8") as f: self.triples = json.load(f) else: self.triples = [] def _build_index(self): for item in self.triples: h = item["head"] r = item["relation"] t = item["tail"] if h not in self.entity2id: eid = len(self.entity2id) self.entity2id[h] = eid self.id2entity[eid] = h if t not in self.entity2id: eid = len(self.entity2id) self.entity2id[t] = eid self.id2entity[eid] = t if r not in self.relation2id: rid = len(self.relation2id) self.relation2id[r] = rid self.id2relation[rid] = r if h not in self.adj: self.adj[h] = [] self.adj[h].append((r, t, item.get("detail", ""))) def score_triple(self, head: str, relation: str, tail: str) -> float: """Computes TransE energy score (lower distance = higher legal coherence).""" if head not in self.entity2id or relation not in self.relation2id or tail not in self.entity2id: return 1.0 h_idx = torch.tensor([self.entity2id[head]]) r_idx = torch.tensor([self.relation2id[relation]]) t_idx = torch.tensor([self.entity2id[tail]]) h_emb = self.entity_embeddings(h_idx) r_emb = self.relation_embeddings(r_idx) t_emb = self.entity_embeddings(t_idx) dist = torch.norm(h_emb + r_emb - t_emb, p=2).item() base_weight = RELATION_WEIGHTS.get(relation, 0.85) # Normalized confidence score between 0.0 and 1.0 score = base_weight * (1.0 / (1.0 + dist * 0.1)) return round(score, 4) def find_multi_hop_pathways(self, seed_entities: List[str], max_depth: int = 3) -> List[Dict[str, Any]]: """ Discovers connected multi-hop legal remedies from seed statutes. Returns ranked litigation pathways across distinct court jurisdictions. """ pathways = [] visited = set() for seed in seed_entities: # Check direct entity matches or fuzzy substring matches matching_nodes = [n for n in self.entity2id if seed.lower() in n.lower() or n.lower() in seed.lower()] for start_node in matching_nodes: self._dfs_pathways(start_node, [start_node], [], pathways, visited, max_depth) # Deduplicate and sort by confidence score unique_paths = {} for p in pathways: path_str = " -> ".join(p["nodes"]) if path_str not in unique_paths or p["confidence"] > unique_paths[path_str]["confidence"]: unique_paths[path_str] = p sorted_paths = sorted(unique_paths.values(), key=lambda x: x["confidence"], reverse=True) return sorted_paths[:6] def _dfs_pathways(self, current: str, current_nodes: List[str], current_relations: List[str], pathways: List[Dict], visited: set, depth_left: int): if depth_left == 0 or current not in self.adj: return for r, next_node, detail in self.adj.get(current, []): if next_node not in current_nodes: new_nodes = current_nodes + [next_node] new_rels = current_relations + [r] score = self.score_triple(current, r, next_node) if len(new_nodes) >= 2: pathways.append({ "nodes": new_nodes, "relations": new_rels, "description": detail, "confidence": score, "start_node": new_nodes[0], "end_node": new_nodes[-1] }) self._dfs_pathways(next_node, new_nodes, new_rels, pathways, visited, depth_left - 1) def format_pathways_for_prompt(self, seed_entities: List[str]) -> str: """Formats discovered pathways into clean markdown for Big Model reasoning.""" paths = self.find_multi_hop_pathways(seed_entities) if not paths: return "" lines = ["\n### GNN-DISCOVERED MULTI-HOP LITIGATION PATHWAYS:"] for idx, p in enumerate(paths, 1): arrow_chain = " -> ".join(f"[{n}]" for n in p["nodes"]) lines.append(f"{idx}. {arrow_chain} (Confidence: {p['confidence'] * 100:.1f}%)") if p.get("description"): lines.append(f" Action: {p['description']}") return "\n".join(lines) # Singleton Instance gnn_engine = RelationalGNNEngine() def get_gnn_engine() -> RelationalGNNEngine: global gnn_engine return gnn_engine if __name__ == "__main__": print(f"[GNN] Total Indexed Entities: {len(gnn_engine.entity2id)}") print(f"[GNN] Total Indexed Relations: {len(gnn_engine.relation2id)}") print(f"[GNN] Sample Pathway Discovery for 'BNS_329':") sample_paths = gnn_engine.find_multi_hop_pathways(["BNS_329"]) for p in sample_paths: print(" •", " -> ".join(p["nodes"]), f"({p['confidence']})")