@@ -43,6 +43,7 @@ def _load_native_module():
4343try :
4444 run_mcp_server = _native_module .run_mcp_server
4545except AttributeError :
46+
4647 def run_mcp_server (* args , ** kwargs ):
4748 raise RuntimeError ("run_mcp_server requires the 'grpc' feature (maturin develop --features full)" )
4849
@@ -866,6 +867,7 @@ def __init__(self, node_feat_dim=None, edge_feat_dim=None, hidden_size=16, n_lay
866867 # Load PyTorch implementation if torch is installed.
867868 try :
868869 from .torch_predecoder import TorchGNNPredecoder
870+
869871 self ._torch_model = TorchGNNPredecoder (nfd , efd , hidden_size , n_layers )
870872 self ._torch_model .double () # Cast to float64 to match Rust f64.
871873 self .use_torch = True
@@ -876,6 +878,7 @@ def _sync_to_rust(self):
876878 if self .use_torch :
877879 import tempfile
878880 import os
881+
879882 with tempfile .NamedTemporaryFile (suffix = ".safetensors" , delete = False ) as tmp :
880883 tmp_path = tmp .name
881884 try :
@@ -889,6 +892,7 @@ def _sync_to_torch(self):
889892 if self .use_torch :
890893 import tempfile
891894 import os
895+
892896 with tempfile .NamedTemporaryFile (suffix = ".safetensors" , delete = False ) as tmp :
893897 tmp_path = tmp .name
894898 try :
@@ -939,6 +943,7 @@ def forward(self, graph):
939943 """Predict adjusted edge weights for a DetectorGraph."""
940944 if self .use_torch :
941945 import torch
946+
942947 self ._torch_model .eval ()
943948 with torch .no_grad ():
944949 node_feats = torch .tensor (graph .node_features , dtype = torch .float64 )
@@ -957,17 +962,18 @@ def train(self, graphs, targets, n_epochs):
957962 raise ImportError ("PyTorch is required to train GNNPredecoder" )
958963 import torch
959964 import torch .optim as optim
965+
960966 optimizer = optim .SGD (self ._torch_model .parameters (), lr = self .learning_rate , weight_decay = self .l2_lambda )
961967 self ._torch_model .train ()
962-
968+
963969 for epoch in range (n_epochs ):
964970 for graph , target in zip (graphs , targets ):
965971 node_feats = torch .tensor (graph .node_features , dtype = torch .float64 )
966972 edge_feats = torch .tensor (graph .edge_features , dtype = torch .float64 )
967973 edge_src = torch .tensor (graph .edge_src , dtype = torch .long )
968974 edge_dst = torch .tensor (graph .edge_dst , dtype = torch .long )
969975 target_t = torch .tensor (target , dtype = torch .float64 )
970-
976+
971977 optimizer .zero_grad ()
972978 preds = self ._torch_model (node_feats , edge_feats , edge_src , edge_dst )
973979 loss = F .mse_loss (preds , target_t )
@@ -982,35 +988,35 @@ def predict_with_node_probs(self, graph):
982988 """Predict edge weights and node error probabilities."""
983989 if self .use_torch :
984990 import torch
991+
985992 self ._torch_model .eval ()
986993 with torch .no_grad ():
987994 node_feats = torch .tensor (graph .node_features , dtype = torch .float64 )
988995 edge_feats = torch .tensor (graph .edge_features , dtype = torch .float64 )
989996 edge_src = torch .tensor (graph .edge_src , dtype = torch .long )
990997 edge_dst = torch .tensor (graph .edge_dst , dtype = torch .long )
991-
998+
992999 weights = self ._torch_model (node_feats , edge_feats , edge_src , edge_dst )
993-
1000+
9941001 N = node_feats .size (0 )
9951002 node_probs = torch .zeros (N , dtype = torch .float64 )
9961003 node_counts = torch .zeros (N , dtype = torch .float64 )
997-
1004+
9981005 node_probs .scatter_add_ (0 , edge_src , weights )
9991006 node_probs .scatter_add_ (0 , edge_dst , weights )
1000-
1007+
10011008 node_counts .scatter_add_ (0 , edge_src , torch .ones_like (edge_src , dtype = torch .float64 ))
10021009 node_counts .scatter_add_ (0 , edge_dst , torch .ones_like (edge_dst , dtype = torch .float64 ))
1003-
1010+
10041011 mask = node_counts > 0
10051012 node_probs [mask ] /= node_counts [mask ]
10061013 node_probs = 1.0 / (1.0 + (- node_probs + 1.0 ).exp ())
1007-
1014+
10081015 return weights .cpu ().numpy ().tolist (), node_probs .cpu ().numpy ().tolist ()
10091016 else :
10101017 return self ._inner .predict_with_node_probs (graph ._inner if isinstance (graph , DetectorGraph ) else graph )
10111018
10121019
1013-
10141020class DetectorGraph :
10151021 """Detector graph used by the GNN and hybrid decoder paths."""
10161022
0 commit comments