Upload source code structural_encoder_v2.py

structural_encoder_v2.py

@@ -0,0 +1,196 @@
+import hashlib
+from collections import defaultdict
+from typing import Dict, List, Tuple, TYPE_CHECKING
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_geometric.data import HeteroData, Batch
+from torch_geometric.nn import HeteroConv, GATConv, global_mean_pool
+from transformers import AutoModel, AutoTokenizer
+from tqdm import tqdm
+import numpy as np
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+# Import Builder from dataloader for inference/eval
+from dataloader import CodeGraphBuilder
+
+class RelationalGraphEncoder(nn.Module):
+    """R-GNN encoder over the AST+CFG heterogeneous graph."""
+
+    EDGE_TYPES = (
+        ("ast", "ast_parent_child", "ast"),
+        ("ast", "ast_child_parent", "ast"),
+        ("ast", "ast_next_sibling", "ast"),
+        ("ast", "ast_prev_sibling", "ast"),
+        ("token", "token_to_ast", "ast"),
+        ("ast", "ast_to_token", "token"),
+        ("stmt", "cfg", "stmt"),
+        ("stmt", "cfg_rev", "stmt"),
+        ("stmt", "stmt_to_ast", "ast"),
+        ("ast", "ast_to_stmt", "stmt"),
+    )
+
+    def __init__(self, hidden_dim: int = 256, out_dim: int = 768, num_layers: int = 2) -> None:
+        super().__init__()
+        self.hidden_dim = hidden_dim
+        self.out_dim = out_dim
+
+        self.ast_encoder = nn.Embedding(2048, hidden_dim)
+        self.token_encoder = nn.Embedding(8192, hidden_dim)
+        self.stmt_encoder = nn.Embedding(512, hidden_dim)
+
+        self.convs = nn.ModuleList()
+        for _ in range(num_layers):
+            hetero_modules = {
+                edge_type: GATConv((-1, -1), hidden_dim, add_self_loops=False)
+                for edge_type in self.EDGE_TYPES
+            }
+            hetero_conv = HeteroConv(hetero_modules, aggr="sum")
+            self.convs.append(hetero_conv)
+
+        self.output_proj = nn.Linear(hidden_dim, out_dim)
+
+    def _encode_nodes(self, data: HeteroData) -> Dict[str, torch.Tensor]:
+        device = self.ast_encoder.weight.device
+        
+        def get_embed(node_type, encoder):
+            if node_type not in data.node_types:
+                return torch.zeros((0, self.hidden_dim), device=device)
+            
+            x = data[node_type].get('x')
+            if x is None:
+                 return torch.zeros((0, self.hidden_dim), device=device)
+            
+            x = x.to(device)
+            return encoder(x)
+
+        x_dict = {
+            "ast": get_embed("ast", self.ast_encoder),
+            "token": get_embed("token", self.token_encoder),
+            "stmt": get_embed("stmt", self.stmt_encoder),
+        }
+        return x_dict
+
+    def forward(self, data: HeteroData) -> torch.Tensor:
+        device = next(self.parameters()).device
+        data = data.to(device)
+        
+        x_dict = self._encode_nodes(data)
+
+        edge_index_dict = {}
+        for edge_type in self.EDGE_TYPES:
+            if edge_type in data.edge_index_dict:
+                edge_index_dict[edge_type] = data.edge_index_dict[edge_type]
+
+        for conv in self.convs:
+            x_dict = conv(x_dict, edge_index_dict)
+            x_dict = {key: F.relu(x) for key, x in x_dict.items()}
+
+        # Global Pooling
+        batch_size = data.num_graphs if hasattr(data, 'num_graphs') else 1
+        
+        pooled_embeddings = []
+        for key, x in x_dict.items():
+            if x.size(0) == 0:
+                 continue
+            
+            if hasattr(data[key], 'batch') and data[key].batch is not None:
+                pool = global_mean_pool(x, data[key].batch, size=batch_size)
+            else:
+                 # Logic for single graph without batch attribute (e.g. inference on one item)
+                 pool = x.mean(dim=0, keepdim=True)
+                 if pool.size(0) != batch_size:
+                    # Should be 1
+                    pass
+            pooled_embeddings.append(pool)
+            
+        if not pooled_embeddings:
+             return torch.zeros((batch_size, self.out_dim), device=device)
+
+        # Average across node types [num_types, B, dim] -> [B, dim]
+        # We need to ensure all pools are [B, dim]. 
+        # If a graph misses a node type, its embedding for that type might be 0 or NaN?
+        # global_mean_pool returns 0 for empty batches.
+        
+        graph_repr = torch.stack(pooled_embeddings).mean(dim=0)
+        return self.output_proj(graph_repr)
+
+
+class GatedFusion(nn.Module):
+    def __init__(self, text_dim: int, graph_dim: int) -> None:
+        super().__init__()
+        self.graph_proj = nn.Linear(graph_dim, text_dim)
+        self.gate = nn.Linear(text_dim * 2, text_dim)
+
+    def forward(self, h_text: torch.Tensor, h_graph: torch.Tensor) -> torch.Tensor:
+        h_graph_proj = self.graph_proj(h_graph)
+        joint = torch.cat([h_text, h_graph_proj], dim=-1)
+        gate = torch.sigmoid(self.gate(joint))
+        return gate * h_text + (1.0 - gate) * h_graph_proj
+
+
+class StructuralEncoderV2(nn.Module):
+    """Structural encoder that fuses GraphCodeBERT text features with AST+CFG graph context."""
+
+    def __init__(self, device: torch.device | str, graph_hidden_dim: int = 256, graph_layers: int = 2):
+        super().__init__()
+        self.device = torch.device(device)
+        self.text_tokenizer = AutoTokenizer.from_pretrained("microsoft/graphcodebert-base")
+        self.text_model = AutoModel.from_pretrained("microsoft/graphcodebert-base")
+        self.text_model.to(self.device)
+
+        self.graph_encoder = RelationalGraphEncoder(hidden_dim=graph_hidden_dim, out_dim=self.text_model.config.hidden_size, num_layers=graph_layers)
+        self.graph_encoder.to(self.device)
+
+        self.fusion = GatedFusion(self.text_model.config.hidden_size, self.text_model.config.hidden_size)
+        self.fusion.to(self.device)
+
+    def encode_text(self, codes: List[str]) -> torch.Tensor:
+        inputs = self.text_tokenizer(
+            codes,
+            padding=True,
+            truncation=True,
+            max_length=512,
+            return_tensors="pt",
+        ).to(self.device)
+        outputs = self.text_model(**inputs)
+        return outputs.last_hidden_state[:, 0, :]
+
+    def forward(self, codes: List[str], graph_batch: Batch | HeteroData) -> torch.Tensor:
+        text_embeddings = self.encode_text(codes)
+        graph_embeddings = self.graph_encoder(graph_batch)
+        return self.fusion(text_embeddings, graph_embeddings)
+
+    def generate_embeddings(self, df: "pd.DataFrame", batch_size: int = 8, save_path: str | None = None, desc: str = "Structural V2 embeddings") -> np.ndarray:
+        # Create local builder for inference
+        builder = CodeGraphBuilder()
+        
+        codes = df["code"].tolist()
+        batches = range(0, len(codes), batch_size)
+        all_embeddings: List[torch.Tensor] = []
+
+        for start in tqdm(batches, desc=desc):
+            batch_codes = codes[start:start + batch_size]
+            
+            # Parallelism here not strictly needed for eval unless slow, but we do it simply
+            data_list = [builder.build(c) for c in batch_codes]
+            graph_batch = Batch.from_data_list(data_list)
+            
+            with torch.no_grad():
+                fused = self.forward(batch_codes, graph_batch)
+            all_embeddings.append(fused.cpu())
+
+        embeddings = torch.cat(all_embeddings, dim=0).numpy().astype("float32")
+        if save_path is not None:
+            np.save(save_path, embeddings)
+        return embeddings
+
+    def load_checkpoint(self, checkpoint_path: str, map_location: str | torch.device = "cpu", strict: bool = True) -> None:
+        if not checkpoint_path:
+            raise ValueError("checkpoint_path must be provided")
+        state = torch.load(checkpoint_path, map_location=map_location)
+        if isinstance(state, dict) and "state_dict" in state:
+            state = state["state_dict"]
+        self.load_state_dict(state, strict=strict)