|
|
import torch |
|
|
import torch.nn as nn |
|
|
from typing import List, Optional, Tuple, Any |
|
|
from tqdm import tqdm |
|
|
import numpy as np |
|
|
import pandas as pd |
|
|
from torch_geometric.data import Batch |
|
|
from transformers import AutoTokenizer |
|
|
|
|
|
|
|
|
from dataloader import CodeGraphBuilder |
|
|
|
|
|
from structural_encoder_v2 import RelationalGraphEncoder, StructuralEncoderV2, GatedFusion |
|
|
|
|
|
class StructuralEncoderOnlyGraph(nn.Module): |
|
|
""" |
|
|
Ablation variant 1: Pure Structural Encoder. |
|
|
Removes GraphCodeBERT and uses only the graph path (R-GNN). |
|
|
""" |
|
|
|
|
|
def __init__(self, device: torch.device | str, graph_hidden_dim: int = 256, graph_layers: int = 2, out_dim: int = 768): |
|
|
super().__init__() |
|
|
self.device = torch.device(device) |
|
|
|
|
|
self.graph_encoder = RelationalGraphEncoder(hidden_dim=graph_hidden_dim, out_dim=out_dim, num_layers=graph_layers) |
|
|
self.graph_encoder.to(self.device) |
|
|
|
|
|
def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, graph_batch: Batch) -> torch.Tensor: |
|
|
|
|
|
return self.graph_encoder(graph_batch) |
|
|
|
|
|
def generate_embeddings(self, df: "pd.DataFrame", batch_size: int = 8, save_path: str | None = None, desc: str = "Structural OnlyGraph embeddings") -> np.ndarray: |
|
|
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] |
|
|
|
|
|
data_list = [builder.build(c) for c in batch_codes] |
|
|
graph_batch = Batch.from_data_list(data_list) |
|
|
|
|
|
|
|
|
dummy_ids = torch.zeros((1,1), device=self.device) |
|
|
dummy_mask = torch.zeros((1,1), device=self.device) |
|
|
|
|
|
with torch.no_grad(): |
|
|
out = self.forward(dummy_ids, dummy_mask, graph_batch) |
|
|
all_embeddings.append(out.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) |
|
|
|
|
|
|
|
|
class StructuralEncoderConcat(StructuralEncoderV2): |
|
|
""" |
|
|
Ablation variant 2: Concatenation Fusion. |
|
|
Keeps both text and graph paths but fuses them via simple concatenation + projection |
|
|
instead of Gated Fusion. |
|
|
""" |
|
|
|
|
|
def __init__(self, device: torch.device | str, graph_hidden_dim: int = 256, graph_layers: int = 2): |
|
|
super().__init__(device, graph_hidden_dim, graph_layers) |
|
|
|
|
|
text_dim = self.text_model.config.hidden_size |
|
|
graph_dim = self.text_model.config.hidden_size |
|
|
|
|
|
self.concat_proj = nn.Linear(text_dim + graph_dim, text_dim) |
|
|
self.concat_proj.to(self.device) |
|
|
|
|
|
del self.fusion |
|
|
|
|
|
def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, graph_batch: Batch) -> torch.Tensor: |
|
|
text_embeddings = self.encode_text(input_ids, attention_mask) |
|
|
graph_embeddings = self.graph_encoder(graph_batch) |
|
|
|
|
|
combined = torch.cat([text_embeddings, graph_embeddings], dim=-1) |
|
|
return self.concat_proj(combined) |
|
|
|
