Add LP-task for JaccardDefence

Jeratt · Jeratt · commit 878efdf55447 · 2026-01-26T19:37:27.000+03:00
diff --git a/gnn_aid/defenses/jaccard_defense/jaccard_def.py b/gnn_aid/defenses/jaccard_defense/jaccard_def.py
@@ -1,103 +1,134 @@
 import torch
 import numpy as np
+from typing import Optional
 
 from gnn_aid.datasets.gen_dataset import GeneralDataset
 from gnn_aid.data_structures.graph_modification_artifacts import GraphModificationArtifact
 from gnn_aid.defenses.poison_defense import PoisonDefender
+from gnn_aid.data_structures.configs import Task
 
 
-class JaccardDefender(
-    PoisonDefender
-):
+def _is_binary_tensor(X: torch.Tensor) -> bool:
+    return torch.all((X == 0) | (X == 1)).item()
+
+
+class JaccardDefender(PoisonDefender):
     """
-    Poison defense based on removing edges between dissimilar nodes
+    Poison defense based on removing edges between dissimilar nodes.
     """
     name = 'JaccardDefender'
 
-    def __init__(self, threshold):
+    def __init__(self, threshold: float, binarize_threshold: Optional[float] = None):
+        """
+        :param threshold: Jaccard similarity threshold (edges with similarity <= threshold are removed)
+        :param binarize_threshold: Optional threshold to binarize non-binary features
+        """
         super().__init__()
-        self.thrsh = threshold
-        self.remove_edge_index = None
+        self.threshold = threshold
+        self.binarize_threshold = binarize_threshold
+        self.removed_edges_train = None
+        self.original_num_edges = None
 
     def defense(
             self,
             gen_dataset: GeneralDataset,
             **kwargs
     ) -> GeneralDataset:
-        """
-        Modify input graph by removing edges between dissimilar nodes
-        :param gen_dataset: input graph dataset
-        :return: modified graph (only adjacency matrix modified)
-        """
+        task = gen_dataset.dataset_var_config.task
+
+        if task in [Task.EDGE_PREDICTION, Task.EDGE_REGRESSION]:
+            if not hasattr(gen_dataset, 'train_mask') or gen_dataset.train_mask is None:
+                raise RuntimeError("JaccardDefender for link tasks requires train_test_split() to be called first")
+
+        self.original_num_edges = gen_dataset.data.edge_index.size(1)
+
+        x = self._prepare_features(gen_dataset.data.x)
+
+        if task in [Task.EDGE_PREDICTION, Task.EDGE_REGRESSION]:
+            gen_dataset = self._defense_link_task(gen_dataset, x)
+        else:
+            gen_dataset = self._defense_standard_task(gen_dataset, x)
 
-        def is_binary_tensor(X: torch.Tensor) -> bool:
-            return torch.all((X == 0) | (X == 1)).item()
-
-        assert is_binary_tensor(gen_dataset.data.x), "The features should be presented in binary form"
-
-        # TODO need to check whether features binary or not. Consistency required - Cora has 'unknown' features e.g.
-        # self.drop_edges(batch)
-        edge_index = gen_dataset.data.edge_index.tolist()
-        #new_edge_mask = torch.zeros_like(gen_dataset.data.edge_index).bool()
-        new_edge_index = [[],[]]
-        self.remove_edge_index = [[], []]
-        for i in range(len(edge_index[0])):
-            if self.jaccard_index(gen_dataset.data.x, edge_index[0][i], edge_index[1][i]) > self.thrsh:
-                # new_edge_mask[0,i] = True
-                # new_edge_mask[1,i] = True
-                new_edge_index[0].append(edge_index[0][i])
-                new_edge_index[1].append(edge_index[1][i])
-            else:
-                self.remove_edge_index[0].append(edge_index[0][i])
-                self.remove_edge_index[1].append(edge_index[1][i])
-        # gen_dataset.data.edge_index *= new_edge_mask.float()
-        gen_dataset.data.edge_index = torch.tensor(new_edge_index).long()
         return gen_dataset
 
-    def jaccard_index(
+    def _prepare_features(self, x: torch.Tensor) -> torch.Tensor:
+        if self.binarize_threshold is not None:
+            x = (x > self.binarize_threshold).float()
+        elif not _is_binary_tensor(x):
+            raise ValueError(
+                "JaccardDefender requires binary features"
+            )
+        return x
 
+    def _defense_link_task(
             self,
-            x,
-            u,
-            v
-    ) -> float:
-        """
-        Computes jaccard index of 'u' and 'v' objects based on their features
-        :param x: feature matrix
-        :param u: index of object from dataset
-        :param v: index of object from dataset
-        :return:
-        """
-        im1 = x[u,:].detach().cpu().numpy().astype(bool)
-        im2 = x[v,:].detach().cpu().numpy().astype(bool)
-        intersection = np.logical_and(im1, im2)
-        union = np.logical_or(im1, im2)
-        return intersection.sum() / float(union.sum())
-
-    def dataset_diff(
-            self
-    ) -> GraphModificationArtifact:
-        diff = GraphModificationArtifact()
+            gen_dataset: GeneralDataset,
+            x: torch.Tensor
+    ) -> GeneralDataset:
+        train_edge_label_index = gen_dataset.edge_label_index[:, gen_dataset.train_mask]
 
-        try:
-            src_nodes = self.remove_edge_index[0]
-            dst_nodes = self.remove_edge_index[1]
+        filtered_train_edges, removed_edges = self._filter_edges_jaccard(train_edge_label_index, x)
+        self.removed_edges_train = removed_edges
 
-            assert len(src_nodes) == len(dst_nodes), (
-                "Mismatch in source and target edge lengths: "
-                f"{len(src_nodes)} vs {len(dst_nodes)}"
-            )
+        gen_dataset.data.edge_index = filtered_train_edges
 
-            edges_to_remove = [
-                [src, dst] for src, dst in zip(src_nodes, dst_nodes)
-            ]
+        num_removed = removed_edges.size(1) if removed_edges is not None else 0
+        print(f"JaccardDefender: Removed {num_removed}/{train_edge_label_index.size(1)} "
+              f"training edges (threshold={self.threshold})")
 
+        return gen_dataset
+
+    def _defense_standard_task(
+            self,
+            gen_dataset: GeneralDataset,
+            x: torch.Tensor
+    ) -> GeneralDataset:
+        filtered_edges, removed_edges = self._filter_edges_jaccard(
+            gen_dataset.data.edge_index, x
+        )
+        self.removed_edges_train = removed_edges  # Reusing field for simplicity
+
+        gen_dataset.data.edge_index = filtered_edges
+
+        num_removed = removed_edges.size(1) if removed_edges is not None else 0
+        print(f"JaccardDefender: Removed {num_removed}/{self.original_num_edges} edges "
+              f"(threshold={self.threshold})")
+
+        return gen_dataset
+
+    def _filter_edges_jaccard(
+            self,
+            edge_index: torch.Tensor,
+            x: torch.Tensor
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if edge_index.size(1) == 0:
+            return edge_index, None
+
+        src_feats = x[edge_index[0]]
+        dst_feats = x[edge_index[1]]
+
+        intersection = (src_feats * dst_feats).sum(dim=1)  # AND
+        union = ((src_feats + dst_feats) > 0).sum(dim=1).float()  # OR
+
+        union = torch.where(union == 0, torch.ones_like(union), union)
+
+        jaccard_scores = intersection / union
+
+        keep_mask = jaccard_scores > self.threshold
+        filtered_edges = edge_index[:, keep_mask]
+        removed_edges = edge_index[:, ~keep_mask] if (~keep_mask).any() else None
+
+        return filtered_edges, removed_edges
+
+    def dataset_diff(self) -> GraphModificationArtifact:
+        diff = GraphModificationArtifact()
+
+        if self.removed_edges_train is not None and self.removed_edges_train.size(1) > 0:
+            edges_to_remove = self.removed_edges_train.t().tolist()
             diff.remove_edges(edges_to_remove)
             self.defense_diff = diff
-
-        except Exception as e:
-            raise RuntimeError(
-                f"Failed to build dataset diff from remove_edge_index: {e}"
-            ) from e
+        else:
+            # No edges removed
+            self.defense_diff = diff
 
         return self.defense_diff
diff --git a/metainfo/functions_parameters.json b/metainfo/functions_parameters.json
@@ -25,5 +25,8 @@
  },
  "NLLLoss": {
   "reduction": ["Reduction", "string", "mean", ["none","mean","sum"], "Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the weighted mean of the output is taken, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction."]
- }
+ },
+  "BCEWithLogitsLoss": {
+  "reduction": ["Reduction", "string", "mean", ["none","mean","sum"], "Specifies the reduction to apply to the output: 'none' | 'mean' | 'sum'. 'none': no reduction will be applied, 'mean': the weighted mean of the output is taken, 'sum': the output will be summed. Note: size_average and reduce are in the process of being deprecated, and in the meantime, specifying either of those two args will override reduction."]
+  }
 }
diff --git a/tests/defense_test.py b/tests/defense_test.py
@@ -4,11 +4,12 @@
 
 from gnn_aid.attacks.mi_attacks import MIAttacker
 from gnn_aid.aux.utils import POISON_DEFENSE_PARAMETERS_PATH, \
-    OPTIMIZERS_PARAMETERS_PATH, MI_ATTACK_PARAMETERS_PATH, MI_DEFENSE_PARAMETERS_PATH
+    OPTIMIZERS_PARAMETERS_PATH, MI_ATTACK_PARAMETERS_PATH, MI_DEFENSE_PARAMETERS_PATH, FUNCTIONS_PARAMETERS_PATH
 from gnn_aid.data_structures.configs import ModelModificationConfig, DatasetConfig, DatasetVarConfig, \
-    ConfigPattern, FeatureConfig, Task
+    ConfigPattern, FeatureConfig, Task, ModelConfig, ModelStructureConfig
 from gnn_aid.datasets.datasets_manager import DatasetManager
 from gnn_aid.datasets.ptg_datasets import LibPTGDataset
+from gnn_aid.models_builder import FrameworkGNNConstructor
 from gnn_aid.models_builder.models_utils import Metric
 from gnn_aid.models_builder.model_managers import FrameworkGNNModelManager
 from gnn_aid.models_builder.models_zoo import model_configs_zoo
@@ -36,6 +37,13 @@ def setUp(self):
 
         self.gen_dataset_sg_cora.train_test_split(percent_train_class=0.6, percent_test_class=0.4)
 
+        # Single-graph - Cora (Link Prediction)
+        self.gen_dataset_lp_cora = DatasetManager.get_by_config(
+            DatasetConfig((LibPTGDataset.data_folder, "Homogeneous", "Planetoid", "Cora")),
+            LibPTGDataset.default_dataset_var_config.clone_with({"task": Task.EDGE_PREDICTION})
+        )
+        self.gen_dataset_lp_cora.train_test_split(percent_train_class=0.85, percent_test_class=0.15)
+
         self.default_config = ModelModificationConfig(
             model_ver_ind=0,
         )
@@ -53,6 +61,27 @@ def setUp(self):
                 }
             }
         )
+
+        self.manager_config_lp = ConfigPattern(
+            _config_class="ModelManagerConfig",
+            _config_kwargs={
+                "mask_features": [],
+                "optimizer": {
+                    "_config_class": "Config",
+                    "_class_name": "Adam",
+                    "_import_path": OPTIMIZERS_PARAMETERS_PATH,
+                    "_class_import_info": ["torch.optim"],
+                    "_config_kwargs": {"weight_decay": 5e-4},
+                },
+                "loss_function": {
+                    "_config_class": "Config",
+                    "_class_name": "BCEWithLogitsLoss",
+                    "_import_path": FUNCTIONS_PARAMETERS_PATH,
+                    "_class_import_info": ["torch.nn"],
+                    "_config_kwargs": {},
+                },
+            }
+        )
         monkey_patch_dirs()
 
     def tearDown(self):
@@ -136,6 +165,106 @@ def test_noise_mi_defender_cora(self):
             print(f"MI Attack accuracy:"
                   f" {MIAttacker.compute_single_attack_accuracy(mask, res, self.gen_dataset_sg_cora.train_mask)}")
 
+    def test_jaccard_defender_link_prediction(self):
+        """
+        Test JaccardDefender on Link Prediction task (Cora dataset).
+        """
+        poison_defense_config = ConfigPattern(
+            _class_name="JaccardDefender",
+            _import_path=POISON_DEFENSE_PARAMETERS_PATH,
+            _config_class="PoisonDefenseConfig",
+            _config_kwargs={
+                "threshold": 0.03,
+            }
+        )
+
+        gnn = FrameworkGNNConstructor(
+            model_config=ModelConfig(
+                structure=ModelStructureConfig([
+                    # Encoder: 2-layer GCN
+                    {
+                        'label': 'n',
+                        'layer': {
+                            'layer_name': 'GCNConv',
+                            'layer_kwargs': {
+                                'in_channels': self.gen_dataset_lp_cora.num_node_features,
+                                'out_channels': 32,
+                            },
+                        },
+                        'activation': {
+                            'activation_name': 'ReLU',
+                            'activation_kwargs': None,
+                        },
+                    },
+                    {
+                        'label': 'n',
+                        'layer': {
+                            'layer_name': 'GCNConv',
+                            'layer_kwargs': {
+                                'in_channels': 32,
+                                'out_channels': 16,
+                            },
+                        },
+                    },
+                    {
+                        'label': 'd',
+                        'function': {
+                            'function_name': 'CosineSimilarity',
+                            'function_kwargs': None
+                        }
+                    }
+                ])
+            )
+        )
+
+        gnn_model_manager = FrameworkGNNModelManager(
+            gnn=gnn,
+            dataset_path=self.gen_dataset_lp_cora.prepared_dir,
+            modification=self.default_config,
+            manager_config=self.manager_config_lp,
+        )
+
+        gnn_model_manager.set_poison_defender(poison_defense_config=poison_defense_config)
+
+        original_train_edges = self.gen_dataset_lp_cora.edge_label_index[:,
+                               self.gen_dataset_lp_cora.train_mask].size(1)
+
+        gnn_model_manager.train_model(
+            gen_dataset=self.gen_dataset_lp_cora,
+            steps=30,
+            save_model_flag=False,
+            metrics=[Metric("AUC", mask='train')]
+        )
+
+        defense = gnn_model_manager.poison_defender
+
+        removed_edges = defense.defense_diff.edges["remove"]
+        num_removed = len(removed_edges)
+
+        if num_removed > 0:
+            print(f"JaccardDefender removed {num_removed} training edges "
+                  f"({num_removed / original_train_edges * 100:.1f}%)")
+            # Sanity checks
+            self.assertGreater(num_removed, 0, "No edges were removed - threshold may be too low")
+            self.assertLess(num_removed, original_train_edges * 0.5,
+                            "Too many edges removed (>50%) - threshold may be too high")
+        else:
+            print("WARNING: No edges removed (threshold may be too low for this graph)")
+
+        test_metrics = gnn_model_manager.evaluate_model(
+            gen_dataset=self.gen_dataset_lp_cora,
+            metrics=[
+                Metric("AUC", mask='test'),
+                Metric("Recall@k", mask='test', k=50),
+                Metric("Recall@k", mask='test', k=100),
+            ]
+        )
+        print("Link Prediction test metrics:", test_metrics)
+
+        self.assertGreater(test_metrics['test']['AUC'], 0.5, "AUC should be >0.5 (random baseline) after training")
+
+        self.assertGreater(test_metrics['test']['Recall@k{k=50}'], 0.0, "Recall@50 should be >0 after training")
+
 
 if __name__ == '__main__':
     unittest.main()
