Spaces:
Sleeping
Sleeping
| """ | |
| 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']})") | |