Create model.py

model.py

@@ -0,0 +1,236 @@
+"""
+relgnn/model.py
+Core RelGNN — Atenção sobre Rotas Atômicas (sem grafo estático).
+
+Arquitetura:
+  1. TableEncoder:      embeddings por tabela via MLP sobre features numéricas
+  2. RouteAggregator:   attention ao longo de cada rota (sequência de tabelas)
+  3. HierarchicalAgg:   agrega múltiplas rotas com pesos aprendidos
+  4. FraudHead:         classificador binário final
+"""
+
+from dataclasses import dataclass
+from typing import List, Dict, Tuple, Optional
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from data.routes import AtomicRoute
+
+
+# ─── CONFIG ───────────────────────────────────────────────────────────────────
+
+@dataclass
+class RelGNNConfig:
+    hidden_dim: int   = 64
+    num_epochs: int   = 50
+    learning_rate: float = 1e-3
+    dropout: float    = 0.2
+    num_heads: int    = 4
+    seed: int         = 42
+
+
+# ─── TABLE ENCODER ────────────────────────────────────────────────────────────
+
+class TableEncoder(nn.Module):
+    """
+    Codifica as features de uma tabela em um embedding de tamanho `hidden_dim`.
+    Opera direto nas colunas numéricas — sem conversão para grafo.
+    """
+    def __init__(self, input_dim: int, hidden_dim: int, dropout: float = 0.2):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim * 2),
+            nn.LayerNorm(hidden_dim * 2),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim * 2, hidden_dim),
+            nn.LayerNorm(hidden_dim),
+            nn.ReLU(),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+# ─── ROUTE ATTENTION ──────────────────────────────────────────────────────────
+
+class RouteAttention(nn.Module):
+    """
+    Mecanismo de atenção sobre uma Rota Atômica.
+    Recebe sequência de embeddings [h1, h2, ..., hK] (K = n_hops + 1)
+    e retorna um embedding agregado representando a rota.
+
+    Implementa atenção scaled-dot-product entre os hops.
+    """
+    def __init__(self, hidden_dim: int, num_heads: int = 4, dropout: float = 0.2):
+        super().__init__()
+        self.attn = nn.MultiheadAttention(
+            embed_dim=hidden_dim,
+            num_heads=num_heads,
+            dropout=dropout,
+            batch_first=True,
+        )
+        self.norm = nn.LayerNorm(hidden_dim)
+        self.mlp  = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim * 2),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim * 2, hidden_dim),
+        )
+
+    def forward(self, hop_embeddings: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            hop_embeddings: [batch, n_hops, hidden_dim]
+        Returns:
+            route_emb: [batch, hidden_dim]  — representação da rota
+            alpha:     [batch, n_hops]      — pesos de atenção por hop
+        """
+        # Self-attention entre os hops da rota
+        attn_out, alpha = self.attn(hop_embeddings, hop_embeddings, hop_embeddings)
+
+        # Residual + norm
+        attn_out = self.norm(attn_out + hop_embeddings)
+
+        # Agrega via mean-pooling ponderado (último hop = entidade alvo)
+        # O primeiro token (tabela alvo) agrega informações dos vizinhos
+        route_emb = attn_out[:, 0, :]   # [batch, hidden_dim]
+        route_emb = route_emb + self.mlp(route_emb)
+
+        alpha_weights = alpha.mean(dim=1)[:, 0, :]   # [batch, n_hops]
+        return route_emb, alpha_weights
+
+
+# ─── HIERARCHICAL ROUTE AGGREGATOR ───────────────────────────────────────────
+
+class HierarchicalRouteAgg(nn.Module):
+    """
+    Agrega embeddings de múltiplas rotas com pesos aprendidos.
+    Cada rota contribui de forma diferente para a predição final.
+    """
+    def __init__(self, hidden_dim: int, num_routes: int):
+        super().__init__()
+        self.route_weights = nn.Parameter(torch.ones(num_routes))
+        self.output_proj   = nn.Linear(hidden_dim, hidden_dim)
+
+    def forward(self, route_embeddings: List[torch.Tensor]) -> torch.Tensor:
+        """
+        Args:
+            route_embeddings: lista de [batch, hidden_dim], uma por rota
+        Returns:
+            agg: [batch, hidden_dim]
+        """
+        stacked = torch.stack(route_embeddings, dim=1)   # [batch, R, hidden]
+        weights = F.softmax(self.route_weights, dim=0)   # [R]
+        weighted = (stacked * weights.unsqueeze(0).unsqueeze(-1)).sum(dim=1)
+        return self.output_proj(weighted)
+
+
+# ─── FRAUD HEAD ───────────────────────────────────────────────────────────────
+
+class FraudHead(nn.Module):
+    def __init__(self, hidden_dim: int, dropout: float = 0.2):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim // 2),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim // 2, 1),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x).squeeze(-1)   # [batch]
+
+
+# ─── RELGNN ───────────────────────────────────────────────────────────────────
+
+class RelGNN(nn.Module):
+    """
+    RelGNN completo.
+
+    Fluxo:
+      tabelas SQL
+        → TableEncoder (por tabela)
+        → RouteAttention (por rota atômica)
+        → HierarchicalRouteAgg
+        → FraudHead
+        → sigmoid(logit) = P(fraude)
+    """
+
+    def __init__(self, config: RelGNNConfig):
+        super().__init__()
+        self.config = config
+        torch.manual_seed(config.seed)
+
+    def build(self, feature_dims: Dict[str, int], routes: List[AtomicRoute]):
+        """Instancia os módulos após conhecer as dimensões das features."""
+        H = self.config.hidden_dim
+        D = self.config.dropout
+
+        self.table_encoders = nn.ModuleDict({
+            table: TableEncoder(dim, H, D)
+            for table, dim in feature_dims.items()
+        })
+
+        self.route_attns = nn.ModuleList([
+            RouteAttention(H, self.config.num_heads, D)
+            for _ in routes
+        ])
+
+        self.hierarchical = HierarchicalRouteAgg(H, len(routes))
+        self.fraud_head    = FraudHead(H, D)
+        self.routes        = routes
+
+    def forward(
+        self,
+        table_features: Dict[str, torch.Tensor],
+    ) -> Tuple[torch.Tensor, Dict]:
+        """
+        Args:
+            table_features: {table_name: [batch, feature_dim]}
+        Returns:
+            logits: [batch]
+            attention_info: dict com pesos de atenção por rota
+        """
+        # 1. Encoder por tabela
+        table_embs = {
+            table: encoder(table_features[table])
+            for table, encoder in self.table_encoders.items()
+            if table in table_features
+        }
+
+        # 2. Attention por rota atômica
+        route_embs = []
+        attention_info = {}
+
+        for i, (route, attn_module) in enumerate(zip(self.routes, self.route_attns)):
+            # Coleta embeddings das tabelas na rota
+            available = [t for t in route.path if t in table_embs]
+            if len(available) < 2:
+                # Usa embedding da tabela alvo repetido se rota incompleta
+                e = table_embs.get(route.path[0], list(table_embs.values())[0])
+                route_embs.append(e)
+                continue
+
+            hop_list = [table_embs[t] for t in available]
+            hop_tensor = torch.stack(hop_list, dim=1)   # [batch, K, H]
+
+            route_emb, alpha = attn_module(hop_tensor)
+            route_embs.append(route_emb)
+            attention_info[f"route_{i}"] = alpha.detach().cpu().numpy()
+
+        # 3. Agrega rotas hierarquicamente
+        agg = self.hierarchical(route_embs)
+
+        # 4. Classificador de fraude
+        logits = self.fraud_head(agg)
+
+        return logits, attention_info
+
+    def fit(self, tables, routes, log_fn=print, progress_fn=None):
+        """Wrapper de treinamento completo."""
+        from relgnn.trainer import Trainer
+        trainer = Trainer(self, self.config)
+        return trainer.fit(tables, routes, log_fn=log_fn, progress_fn=progress_fn)