training edge pred ok

mishadr · mishadr · commit 5b01aa472a50 · 2026-02-05T15:39:20.000+03:00
diff --git a/gnn_aid/attacks/evasion_attacks.py b/gnn_aid/attacks/evasion_attacks.py
@@ -25,6 +25,7 @@
 # ReWatt imports
 from .evasion_attacks_collection.rewatt.utils import GraphEnvironment, ReWattPolicyNet, \
     GraphState, ReWattAgent
+from ..data_structures import Task
 
 
 class EvasionAttacker(
@@ -94,7 +95,7 @@ def attack(
             model = model_manager.gnn
             model.eval()
 
-            if task.is_edge_level():
+            if task == Task.EDGE_PREDICTION:
                 data = gen_dataset.data
                 x = data.x
                 x.requires_grad = True
@@ -119,7 +120,7 @@ def attack(
                 perturbed_data_x = torch.clamp(perturbed_data_x, 0, 1)
                 gen_dataset.data.x = perturbed_data_x.detach()
 
-            elif task.is_graph_level() and task.is_classification():  # graph_classification
+            elif task == Task.GRAPH_CLASSIFICATION:
                 graph_idx = self.element_idx
                 x = gen_dataset.dataset[graph_idx].x
                 x.requires_grad = True
@@ -135,7 +136,7 @@ def attack(
                 perturbed_data_x = torch.clamp(perturbed_data_x, 0, 1)
                 gen_dataset.dataset[graph_idx].x = perturbed_data_x.detach()
 
-            elif task.is_node_level() and task.is_classification():  # node_classification
+            elif task == Task.NODE_CLASSIFICATION:
                 node_idx = self.element_idx
                 x = gen_dataset.data.x
                 x.requires_grad = True
@@ -152,7 +153,7 @@ def attack(
                 gen_dataset.data.x = perturbed_data_x.detach()
 
             else:
-                pass
+                raise NotImplementedError
 
             # if task_type:
             #     gni = GlobalNodeIndexer(gen_dataset.dataset)
@@ -542,6 +543,7 @@ def attack(
                         x.copy_(torch.clamp(x, -self.epsilon, self.epsilon))
 
                 self.attack_res = Data(x=x, edge_index=edge_index, y=y)
+                gen_dataset.data.x = x  # FIXME tmp
 
         else:  # structure attack
             if task_type:
diff --git a/gnn_aid/datasets/gen_dataset.py b/gnn_aid/datasets/gen_dataset.py
@@ -204,10 +204,11 @@ def node_features(
         if self.dataset is None:
             raise RuntimeError(f"Cannot get node features: dataset {self} is not built."
                                f" Define {DatasetVarConfig.__name__} and call build() method")
-        if self.is_multi():
-            return [data.x for data in self.dataset]
-        else:
-            return self.dataset[0].x
+        # if self.is_multi():
+        #     return [data.x for data in self.dataset]
+        # else:
+        #     return self.dataset[0].x
+        return self.data.x
 
     @property
     def edge_features(
@@ -219,11 +220,12 @@ def edge_features(
             raise RuntimeError(f"Cannot get edge features: dataset {self} is not built."
                                f" Define {DatasetVarConfig.__name__} and call build() method")
 
-        if self.is_multi():
-            return [data.edge_attr for data in self.dataset]
-        else:
-            return self.dataset[0].edge_attr
+        # if self.is_multi():
+        #     return [data.edge_attr for data in self.dataset]
+        # else:
+        #     return self.dataset[0].edge_attr
         # TODO misha implement
+        return self.data.edge_attr
 
     def __len__(
             self
diff --git a/gnn_aid/models_builder/model_managers/framework_mm.py b/gnn_aid/models_builder/model_managers/framework_mm.py
@@ -4,7 +4,6 @@
 from typing import Union, List, cast, Iterable
 
 import torch
-from torch import tensor
 from torch.cuda import is_available
 from torch.nn.utils import clip_grad_norm
 from torch_geometric.loader import NeighborLoader, DataLoader, LinkNeighborLoader
@@ -16,7 +15,7 @@
 from gnn_aid.data_structures import Task, GraphModificationArtifact
 from gnn_aid.data_structures.configs import ConfigPattern, CONFIG_OBJ
 from gnn_aid.datasets import GeneralDataset
-from gnn_aid.models_builder.models_utils import Metric, predict_top_k_edges
+from gnn_aid.models_builder.models_utils import Metric, predict_top_k_edges, mask_to_tensor
 from . import GNNModelManager
 
 
@@ -751,28 +750,3 @@ def load_train_test_split(
         return gen_dataset
 
 
-def mask_to_tensor(
-        gen_dataset: GeneralDataset,
-        mask: Union[str, List[bool], torch.Tensor] = 'test'
-) -> torch.Tensor:
-    """
-    Convert mask over nodes/edges/graphs to tensor.
-    Mask can be 'train', 'val', 'test', 'all', or Tensor of specific nodes/edges/graphs.
-
-    :param gen_dataset: dataset
-    :param mask: part of the dataset on which the output will be obtained.
-     'train', 'val', 'test', 'all', or Tensor of specific nodes/edges/graphs
-    :return: tensor of nodes/edges/graphs
-    """
-    try:
-        mask_tensor = {
-            'train': gen_dataset.train_mask,
-            'val': gen_dataset.val_mask,
-            'test': gen_dataset.test_mask,
-            'all': tensor([True] * len(gen_dataset.labels)),
-        }[mask]
-    except KeyError:
-        assert isinstance(mask, torch.Tensor)
-        mask_tensor = mask
-
-    return mask_tensor
diff --git a/gnn_aid/models_builder/models_utils.py b/gnn_aid/models_builder/models_utils.py
@@ -1,4 +1,4 @@
-from typing import Any, Optional, List, Callable, Union
+from typing import Any, Optional, List, Callable, Union, Tuple
 from typing import Callable
 
 import sklearn.metrics
@@ -8,6 +8,8 @@
 from torch.utils.hooks import RemovableHandle
 from torch_geometric.nn import MessagePassing
 
+from gnn_aid.datasets import GeneralDataset
+
 
 def apply_message_gradient_capture(
         layer: Any,
@@ -694,4 +696,52 @@ def _predict_full_enumeration(model, data, h, existing_set, k, is_directed, remo
 
     print(f"Top-{final_k} edges found with scores from {top_scores[-1]:.4f} to {top_scores[0]:.4f}")
 
-    return top_edges.cpu(), top_scores.cpu()
+    return top_edges.cpu(), top_scores.cpu()
+
+
+def mask_to_tensor(
+        gen_dataset: GeneralDataset,
+        mask: Union[str, int, Tuple[int], list, torch.Tensor] = 'test'
+) -> torch.Tensor:
+    """
+    Convert a mask over nodes/edges/graphs to tensor of specific size.
+    Mask can be 'train', 'val', 'test', 'all', or id, or a list of ids, or a tensor.
+
+    :param gen_dataset: dataset
+    :param mask: part of the dataset on which the output will be obtained.
+     Can be a node id, graph id, or edge as a tuple (i,j).
+     Can be string: 'train', 'val', 'test', 'all'.
+     Can be Tensor of specific nodes/edges/graphs.
+    :return: tensor of nodes/edges/graphs
+    """
+    task = gen_dataset.dataset_var_config.task
+
+    if isinstance(mask, str):
+        mask_tensor = {
+            'train': gen_dataset.train_mask,
+            'val': gen_dataset.val_mask,
+            'test': gen_dataset.test_mask,
+            'all': tensor([True] * len(gen_dataset.labels)),
+        }[mask]
+
+    elif isinstance(mask, torch.Tensor):
+        mask_tensor = mask
+
+    elif task.is_node_level():  # Node id
+        assert not gen_dataset.is_multi()
+        mask_tensor = tensor([False] * gen_dataset.info.nodes[0])
+        mask_tensor[mask] = True  # for int or list of ints
+
+    elif task.is_graph_level():  # Graph id
+        assert gen_dataset.is_multi()
+        mask_tensor = tensor([False] * len(gen_dataset.info.nodes))
+        mask_tensor[mask] = True  # for int or list of ints
+
+    elif task.is_edge_level():  # Edge
+        # isinstance(mask, tuple)
+        raise NotImplementedError
+
+    else:
+        raise RuntimeError(f"Cannot infer mask tensor for given mask {mask}.")
+
+    return mask_tensor
diff --git a/web_interface/back_front/attack_defense_blocks.py b/web_interface/back_front/attack_defense_blocks.py
@@ -5,15 +5,17 @@
 
 from gnn_aid.attacks import MIAttacker
 from gnn_aid.aux import DataInfo
-from gnn_aid.data_structures.configs import ConfigPattern, PoisonAttackConfig, PoisonDefenseConfig, EvasionAttackConfig, \
-    EvasionDefenseConfig, MIAttackConfig, MIDefenseConfig
+from gnn_aid.data_structures.configs import ConfigPattern, PoisonAttackConfig, PoisonDefenseConfig, \
+    EvasionAttackConfig, \
+    EvasionDefenseConfig, MIAttackConfig, MIDefenseConfig, CONFIG_OBJ
 from gnn_aid.aux.utils import POISON_ATTACK_PARAMETERS_PATH, POISON_DEFENSE_PARAMETERS_PATH, \
     EVASION_ATTACK_PARAMETERS_PATH, EVASION_DEFENSE_PARAMETERS_PATH, MI_ATTACK_PARAMETERS_PATH, \
     MI_DEFENSE_PARAMETERS_PATH
 from gnn_aid.datasets.gen_dataset import GeneralDataset
 from gnn_aid.models_builder import Metric
 from gnn_aid.models_builder.attack_defense_manager import FrameworkAttackDefenseManager
 from gnn_aid.models_builder.model_managers import GNNModelManager
+from gnn_aid.models_builder.models_utils import mask_to_tensor
 from . import VisiblePart
 from .block import Block
 from .utils import WebInterfaceError, send_epoch_results, compute_stats_data
@@ -212,11 +214,12 @@ def do(
 
             # Apply evasion attack
             if self.ad_configs["AD-ea"] is not None:
-                # attack_mask =
+                element_idx = getattr(self.ad_configs["AD-ea"], CONFIG_OBJ).element_idx
+                attack_mask = mask_to_tensor(self.gen_dataset, element_idx)
 
                 self.model_manager.call_evasion_attack(
                     gen_dataset=self.gen_dataset,
-                    mask=torch.empty(1),
+                    mask=attack_mask,
                 )
 
             # Evaluate metrics without attacks
@@ -253,6 +256,7 @@ def do(
 
             # Update dataset features
             dvd = self.visible_part.get_dataset_var_data()
+            print('dvd after attack', dvd.node['features'])
 
             dvd = add_into_dvd(self.gen_dataset, stats_data, dvd)
 
diff --git a/web_interface/main_aiohttp.py b/web_interface/main_aiohttp.py
@@ -16,7 +16,12 @@
 from web_interface.back_front.utils import WebInterfaceError, json_loads, json_dumps, SocketConnect
 
 # Socket.IO server (ASGI not used here)
-sio = socketio.AsyncServer(async_mode="aiohttp")
+sio = socketio.AsyncServer(
+    async_mode="aiohttp",
+    ping_timeout=600,  # wait pong from client
+    ping_interval=25,
+    cors_allowed_origins='*'
+)
 app = web.Application()
 sio.attach(app)
 
diff --git a/web_interface/static/js/controllers/controller.js b/web_interface/static/js/controllers/controller.js
@@ -12,13 +12,19 @@ class Controller {
             reconnectionAttempts: Infinity,
             reconnectionDelay: 1000,
             query: {mode: mode},
+
+            // Longer timeout for backend debug - 10 mins
+            timeout: 600*1000,
+            pingTimeout: 600*1000,
+            pingInterval: 600*1000,
         })
         this.socket.on('connect', () => {
             console.log('socket connected, sid=', this.socket.id)
             this.sid = this.socket.id
             if (this.isActive) {
+                console.log('This is a reconnection')
                 // FIXME seems sometimes it happens when backend is busy - we don't need to reload?
-                // Means re-connection to server. Need to reload the page
+                // // Means re-connection to server. Need to reload the page
                 alert("Web-socket connection is lost. Press OK to reload the page.")
                 this.isActive = false
                 window.location.reload(true)
@@ -49,14 +55,20 @@ class Controller {
             console.log('socket error', err)
         })
 
-        this.socket.on('disconnect', () => {
+        this.socket.on('disconnect', (reason) => {
             // this.isActive = false
-            console.log('Disconnected from server, sid=', this.socket.id);
+            console.log('Disconnected from server, reason=', reason, ', sid=', this.socket.id);
             if (this.isActive) {
-                // Show a notification to the user
-                alert("Web-socket connection is lost. Press OK to reload the page.")
-                this.isActive = false
-                window.location.reload(true)
+                const softDisconnectReasons = ['ping timeout', 'transport error'];
+                if (softDisconnectReasons.includes(reason)) {
+                    console.log('Soft disconnect, will try to reconnect...');
+                    // this.reconnectStartTime = Date.now()
+                } else {
+                    // Show a notification to the user
+                    alert("Web-socket connection is lost. Press OK to reload the page.")
+                    this.isActive = false
+                    window.location.reload(true)
+                }
             }
         })
 
diff --git a/web_interface/static/js/paramsBuilder.js b/web_interface/static/js/paramsBuilder.js
@@ -125,6 +125,52 @@ class ParamsBuilder {
                 $input.change(() => this.kwArgs[name] = $input.is(":checked"))
             }
 
+            else if (type === "int_or_tuple") {
+                // fixme this is a copy. todo tuple for edge task
+                $input.attr("type", "number")
+                $input.val(Number.isFinite(def) ? def : possible["min"])
+                let checkClass = id + "-radio"
+                if ("special" in possible) {
+                    // Add special values as separate checkboxes
+                    let checkId = id + "-check"
+                    for (const variant of possible.special) {
+                        let $checkBox = $("<input>").attr("id", checkId)
+                            .attr("type", "checkbox").prop('checked', variant === def)
+                        $checkBox.addClass(checkClass)
+                        $cb.append($checkBox)
+                        let $label = $("<label></label>").text(variant == null ? "None" : variant)
+                            .attr("for", checkId)
+                        $cb.append($label)
+                        $checkBox.change((e) => { // Uncheck all but this
+                            let wasChecked = $checkBox.is(":checked")
+                            $("." + checkClass).prop("checked", false)
+                            $checkBox.prop("checked", true)
+                            this.kwArgs[name] = variant
+                        })
+                    }
+                    $input.focus(() => {
+                        $("." + checkClass).prop("checked", false)
+                        $input.trigger("change")
+                    })
+                    $input.css("min-width", "60px")
+                    delete possible.special
+                }
+
+                if (type === "int") {
+                    $input.attr("step", 1)
+                    $input.attr("pattern", "\d+")
+                    $input.change(() => this.kwArgs[name] = parseInt($input.val()))
+                }
+                else {
+                    // fixme
+                }
+                for (const [key, value] of Object.entries(possible))
+                    $input.attr(key, value)
+
+                // Check input value when user unfocus it or change it
+                addValueChecker($input, type, def, possible["min"], possible["max"], "change")
+            }
+
             else if (type === "int" || type === "float") {
                 $input.attr("type", "number")
                 $input.val(Number.isFinite(def) ? def : possible["min"])
