Upload grn/model.py

grn/model.py

@@ -0,0 +1,246 @@
+"""
+Graph Reasoning Network (GRN) - Core Model
+
+An LLM alternative that operates on knowledge graphs instead of text tokens.
+See the full docstrings in the training script for architecture details.
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_geometric.nn import MessagePassing
+from typing import Optional, Tuple, Dict
+import math
+
+
+class QueryEncoder(nn.Module):
+    def __init__(self, vocab_size=32000, embed_dim=256, num_heads=8, num_layers=4, max_seq_len=512):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.token_embedding = nn.Embedding(vocab_size, embed_dim)
+        self.position_embedding = nn.Embedding(max_seq_len, embed_dim)
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=embed_dim, nhead=num_heads, dim_feedforward=embed_dim * 4,
+            dropout=0.1, activation='gelu', batch_first=True, norm_first=True)
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+        self.output_proj = nn.Linear(embed_dim, embed_dim)
+        self.layer_norm = nn.LayerNorm(embed_dim)
+
+    def forward(self, token_ids, attention_mask=None):
+        B, L = token_ids.shape
+        positions = torch.arange(L, device=token_ids.device).unsqueeze(0).expand(B, -1)
+        x = self.token_embedding(token_ids) + self.position_embedding(positions)
+        mask = ~attention_mask.bool() if attention_mask is not None else None
+        x = self.transformer(x, src_key_padding_mask=mask)
+        if attention_mask is not None:
+            m = attention_mask.unsqueeze(-1).float()
+            pooled = (x * m).sum(dim=1) / m.sum(dim=1).clamp(min=1)
+        else:
+            pooled = x.mean(dim=1)
+        return self.layer_norm(self.output_proj(pooled))
+
+
+class RelationAwareMessagePassing(MessagePassing):
+    def __init__(self, hidden_dim, edge_dim, num_relation_types=256):
+        super().__init__(aggr='add')
+        self.message_mlp = nn.Sequential(
+            nn.Linear(hidden_dim + edge_dim, hidden_dim * 2), nn.GELU(),
+            nn.Linear(hidden_dim * 2, hidden_dim))
+        self.gate_mlp = nn.Sequential(nn.Linear(hidden_dim * 2, hidden_dim), nn.Sigmoid())
+        self.layer_norm = nn.LayerNorm(hidden_dim)
+
+    def forward(self, x, edge_index, edge_attr):
+        msg = self.propagate(edge_index, x=x, edge_attr=edge_attr)
+        gate = self.gate_mlp(torch.cat([x, msg], dim=-1))
+        return self.layer_norm(gate * msg + (1 - gate) * x)
+
+    def message(self, x_j, edge_attr):
+        return self.message_mlp(torch.cat([x_j, edge_attr], dim=-1))
+
+
+class GraphNavigator(nn.Module):
+    def __init__(self, hidden_dim=256, edge_dim=64, num_layers=6, num_relation_types=256):
+        super().__init__()
+        self.num_layers = num_layers
+        self.query_proj = nn.Linear(hidden_dim, hidden_dim)
+        self.mp_layers = nn.ModuleList([
+            RelationAwareMessagePassing(hidden_dim, edge_dim, num_relation_types)
+            for _ in range(num_layers)])
+        self.query_attention = nn.ModuleList([
+            nn.MultiheadAttention(hidden_dim, num_heads=8, batch_first=True)
+            for _ in range(num_layers)])
+        self.query_norms = nn.ModuleList([nn.LayerNorm(hidden_dim) for _ in range(num_layers)])
+        self.relevance_head = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, 1))
+
+    def forward(self, node_features, edge_index, edge_attr, query, start_node_mask=None):
+        h = node_features.clone()
+        if start_node_mask is not None:
+            qi = self.query_proj(query)
+            if qi.dim() == 2: qi = qi.squeeze(0)
+            h[start_node_mask.bool()] = h[start_node_mask.bool()] + qi
+        for i in range(self.num_layers):
+            h = self.mp_layers[i](h, edge_index, edge_attr)
+            hu = h.unsqueeze(0)
+            qu = query.unsqueeze(1) if query.dim() == 2 else query.unsqueeze(0).unsqueeze(1)
+            att, _ = self.query_attention[i](hu, qu, qu)
+            h = self.query_norms[i](h + att.squeeze(0))
+        return h, self.relevance_head(h)
+
+
+class NodeCreator(nn.Module):
+    def __init__(self, hidden_dim=256, edge_dim=64):
+        super().__init__()
+        self.coverage_head = nn.Sequential(
+            nn.Linear(hidden_dim * 2, hidden_dim), nn.GELU(),
+            nn.Linear(hidden_dim, 1), nn.Sigmoid())
+        self.node_generator = nn.Sequential(
+            nn.Linear(hidden_dim * 2, hidden_dim * 2), nn.GELU(),
+            nn.Linear(hidden_dim * 2, hidden_dim), nn.LayerNorm(hidden_dim))
+        self.edge_generator = nn.Sequential(
+            nn.Linear(hidden_dim * 2, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, edge_dim))
+        self.connection_scorer = nn.Sequential(
+            nn.Linear(hidden_dim * 2, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, 1))
+
+    def forward(self, node_features, query, relevance_scores):
+        if query.dim() == 2: query = query.squeeze(0)
+        w = torch.softmax(relevance_scores.squeeze(-1), dim=0)
+        gs = (node_features * w.unsqueeze(-1)).sum(dim=0)
+        coverage = self.coverage_head(torch.cat([gs, query], dim=-1))
+        new_node = self.node_generator(torch.cat([gs, query], dim=-1)).unsqueeze(0)
+        qe = query.unsqueeze(0).expand(node_features.shape[0], -1)
+        ci = torch.cat([node_features, qe], dim=-1)
+        return coverage, new_node, self.connection_scorer(ci).squeeze(-1), self.edge_generator(ci)
+
+
+class EdgePredictor(nn.Module):
+    def __init__(self, hidden_dim=256, edge_dim=64, num_relation_types=256, gamma=12.0):
+        super().__init__()
+        self.gamma = gamma
+        self.complex_dim = hidden_dim // 2
+        self.head_proj = nn.Linear(hidden_dim, hidden_dim)
+        self.tail_proj = nn.Linear(hidden_dim, hidden_dim)
+        self.relation_phases = nn.Embedding(num_relation_types, self.complex_dim)
+        nn.init.uniform_(self.relation_phases.weight, 0, 2 * math.pi)
+        self.edge_feat_gen = nn.Sequential(
+            nn.Linear(self.complex_dim, edge_dim * 2), nn.GELU(), nn.Linear(edge_dim * 2, edge_dim))
+
+    def score_edges(self, head, tail, relation_ids):
+        h, t = self.head_proj(head), self.tail_proj(tail)
+        cd = self.complex_dim
+        re_h, im_h = h[:, :cd], h[:, cd:]
+        re_t, im_t = t[:, :cd], t[:, cd:]
+        phase = self.relation_phases(relation_ids)
+        re_r, im_r = torch.cos(phase), torch.sin(phase)
+        re_s = re_h * re_r - im_h * im_r - re_t
+        im_s = re_h * im_r + im_h * re_r - im_t
+        return self.gamma - torch.norm(torch.stack([re_s, im_s], dim=0), dim=0).sum(dim=-1)
+
+    def forward(self, node_features, candidate_edges, relation_ids):
+        scores = self.score_edges(node_features[candidate_edges[0]],
+                                  node_features[candidate_edges[1]], relation_ids)
+        return scores, self.edge_feat_gen(self.relation_phases(relation_ids))
+
+
+class SubgraphExtractor(nn.Module):
+    def __init__(self, hidden_dim=256):
+        super().__init__()
+        self.node_selector = nn.Sequential(
+            nn.Linear(hidden_dim * 2, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, 1))
+        self.edge_selector = nn.Sequential(
+            nn.Linear(hidden_dim * 3, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, 1))
+        self.order_head = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, 1))
+
+    def forward(self, node_features, edge_index, edge_attr, query, relevance_scores):
+        if query.dim() == 2: query = query.squeeze(0)
+        N = node_features.shape[0]
+        qe = query.unsqueeze(0).expand(N, -1)
+        nl = self.node_selector(torch.cat([node_features, qe], dim=-1)).squeeze(-1)
+        np_ = torch.sigmoid(nl + relevance_scores.squeeze(-1))
+        if edge_index.shape[1] > 0:
+            sf = node_features[edge_index[0]]
+            df = node_features[edge_index[1]]
+            qee = query.unsqueeze(0).expand(edge_index.shape[1], -1)
+            el = self.edge_selector(torch.cat([sf, df, qee], dim=-1)).squeeze(-1)
+            ep = torch.sigmoid(el) * np_[edge_index[0]] * np_[edge_index[1]]
+        else:
+            ep = torch.zeros(0, device=node_features.device)
+        return np_, ep, self.order_head(node_features).squeeze(-1)
+
+
+class GraphReasoningNetwork(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        hd = config.get('hidden_dim', 256)
+        ed = config.get('edge_dim', 64)
+        self.hidden_dim = hd
+        self.edge_dim = ed
+        self.query_encoder = QueryEncoder(config.get('vocab_size', 32000), hd, 8, config.get('num_encoder_layers', 4))
+        self.navigator = GraphNavigator(hd, ed, config.get('num_nav_layers', 6), config.get('num_relation_types', 256))
+        self.node_creator = NodeCreator(hd, ed)
+        self.edge_predictor = EdgePredictor(hd, ed, config.get('num_relation_types', 256))
+        self.subgraph_extractor = SubgraphExtractor(hd)
+        self.node_input_proj = nn.Linear(hd, hd)
+        self.edge_input_proj = nn.Linear(ed, ed)
+
+    def forward(self, token_ids, attention_mask, node_features, edge_index, edge_attr,
+                start_node_mask, target_node_mask=None, target_edge_mask=None,
+                target_new_nodes=None, candidate_edges=None, candidate_edge_labels=None,
+                candidate_edge_relations=None):
+        results = {}
+        query = self.query_encoder(token_ids, attention_mask)
+        h = self.node_input_proj(node_features)
+        e = self.edge_input_proj(edge_attr) if edge_attr.shape[0] > 0 else edge_attr
+        h_nav, rel = self.navigator(h, edge_index, e, query, start_node_mask)
+        results.update({'node_features': h_nav, 'relevance_scores': rel})
+        cov, nn_, cs, nef = self.node_creator(h_nav, query, rel)
+        results.update({'coverage': cov, 'new_node_features': nn_, 'connection_scores': cs})
+        if candidate_edges is not None and candidate_edges.shape[1] > 0:
+            es, pef = self.edge_predictor(h_nav, candidate_edges, candidate_edge_relations)
+            results['edge_scores'] = es
+        np_, ep, to = self.subgraph_extractor(h_nav, edge_index, e, query, rel)
+        results.update({'node_selection_probs': np_, 'edge_selection_probs': ep, 'topological_order': to})
+        losses = {}
+        if target_node_mask is not None:
+            losses['node_selection_loss'] = F.binary_cross_entropy(np_, target_node_mask.float())
+        if target_edge_mask is not None and ep.numel() > 0:
+            losses['edge_selection_loss'] = F.binary_cross_entropy(ep, target_edge_mask.float())
+        if target_new_nodes is not None:
+            losses['node_creation_loss'] = F.mse_loss(nn_, target_new_nodes)
+        if candidate_edge_labels is not None and 'edge_scores' in results:
+            losses['edge_prediction_loss'] = F.binary_cross_entropy_with_logits(results['edge_scores'], candidate_edge_labels.float())
+        if target_node_mask is not None:
+            losses['coverage_loss'] = F.mse_loss(cov.squeeze(), target_node_mask.float().mean())
+        if edge_index.shape[1] > 0 and target_edge_mask is not None:
+            ov = F.relu(to[edge_index[0]] - to[edge_index[1]] + 0.1)
+            losses['dag_ordering_loss'] = (ov * target_edge_mask.float()).mean()
+        results['losses'] = losses
+        if losses: results['total_loss'] = sum(losses.values())
+        return results
+
+    def count_parameters(self):
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+
+    @torch.no_grad()
+    def reason(self, token_ids, attention_mask, node_features, edge_index, edge_attr,
+               start_node_mask, node_threshold=0.5, edge_threshold=0.3, create_threshold=0.5):
+        self.eval()
+        r = self.forward(token_ids, attention_mask, node_features, edge_index, edge_attr, start_node_mask)
+        sn = r['node_selection_probs'] > node_threshold
+        se = r['edge_selection_probs'] > edge_threshold if r['edge_selection_probs'].numel() > 0 else torch.zeros(0, dtype=torch.bool)
+        order = r['topological_order']
+        if se.any():
+            v = sn[edge_index[0][se]] & sn[edge_index[1][se]]
+            sf = se.clone(); sf[se] = v
+        else: sf = se
+        if sf.any():
+            em = sf.nonzero(as_tuple=True)[0]
+            fwd = order[edge_index[0][em]] < order[edge_index[1][em]]
+            sd = torch.zeros_like(sf); sd[em[fwd]] = True
+        else: sd = sf
+        return {'selected_nodes': sn, 'selected_edges': sd, 'node_scores': r['node_selection_probs'],
+                'edge_scores': r['edge_selection_probs'], 'topological_order': order,
+                'relevance_scores': r['relevance_scores'], 'coverage': r['coverage'],
+                'new_node_features': r['new_node_features'], 'connection_scores': r['connection_scores'],
+                'should_create_node': r['coverage'].item() < create_threshold}