Revert "dataset v3 try to do inference"

This reverts commit 14ef5f8.

phishGNN/cross_validation.py

@@ -1,5 +1,4 @@
 import time
-from typing import Tuple
 
 import torch
 from sklearn.model_selection import StratifiedKFold
@@ -16,7 +15,7 @@
 
 def cross_validation_with_val_set(dataset, model, loss_fn, folds, epochs, batch_size,
                                   lr, lr_decay_factor, lr_decay_step_size,
-                                  weight_decay, logger=None) -> Tuple[float, float, float]:
+                                  weight_decay, logger=None) -> tuple[float, float, float]:
 
     val_losses, accs, durations = [], [], []
     for fold, (train_idx, test_idx, val_idx) in enumerate(zip(*k_fold(dataset, folds))):

phishGNN/dataprep.py

@@ -10,7 +10,7 @@
 NAN_VALUE = -1
 
 
-def read_csv(path: str, train_test_equilibrum: bool=True) -> pd.DataFrame:
+def read_csv(path: str) -> pd.DataFrame:
     """Opens the csv dataset as DataFrame and cast types.
     """
     date_parser = lambda c: pd.to_datetime(c, format='%Y-%m-%dT%H:%M:%SZ', errors='coerce')
@@ -22,11 +22,10 @@ def read_csv(path: str, train_test_equilibrum: bool=True) -> pd.DataFrame:
     )
 
     # equilibrate dataset classes as 50/50% benign/phishing
-    if train_test_equilibrum:
-        nb_phishing = len(df[df['is_phishing'] == 1])
-        benign = df.index[(df['is_phishing'] == 0)][:nb_phishing]
-        other = df.index[~(df['is_phishing'] == 0)]
-        df = pd.concat([df.iloc[benign], df.iloc[other]])
+    nb_phishing = len(df[df['is_phishing'] == 1])
+    benign = df.index[(df['is_phishing'] == 0)][:nb_phishing]
+    other = df.index[~(df['is_phishing'] == 0)]
+    df = pd.concat([df.iloc[benign], df.iloc[other]])
 
     # cast object dtypes
     df['url'] = df['url'].astype('string')
@@ -116,7 +115,7 @@ def load_every_urls_with_features(df: pd.DataFrame, path: str) -> Tuple[List, Li
     return every_urls, X
 
 
-def load_train_set(csv_file: str, train_test_equilibrum: bool=True) -> Tuple[pd.DataFrame, List[List], List[int]]:
+def load_train_set(csv_file: str) -> Tuple[pd.DataFrame, List[List], List[int]]:
     """Opens the csv file in `csv_file` and returns every
     features and label of each root url in the dataset.
 
@@ -125,7 +124,7 @@ def load_train_set(csv_file: str, train_test_equilibrum: bool=True) -> Tuple[pd.
         X: the list of features (list) of each root url
         y: the list of labels (int) of each root url
     """
-    df = read_csv(csv_file, train_test_equilibrum=train_test_equilibrum)
+    df = read_csv(csv_file)
     df = normalize_features(df)
 
     root_urls = df[~df['is_phishing'].isin([NAN_VALUE])]['url']

phishGNN/dataset_v1.py

@@ -11,7 +11,6 @@
 import dataprep
 from utils.compute_device import COMPUTE_DEVICE
 from utils.utils import normalize_www_prefix
-from typing import Tuple, List
 
 print(f'Torch version: {torch.__version__}')
 print(f'Compute device: {COMPUTE_DEVICE}')
@@ -43,12 +42,12 @@ def __init__(
         super(PhishingDataset, self).__init__(root, transform, pre_transform)
 
     @property
-    def raw_file_names(self) -> List[str]:
+    def raw_file_names(self) -> list[str]:
         """File name of the csv dataset. """
         return glob.glob(os.path.join(self.raw_dir, '*'))
 
     @property
-    def processed_file_names(self) -> List[str]:
+    def processed_file_names(self) -> list[str]:
         return [file + '.pt' for file in self.raw_file_names]
 
     @property
@@ -90,7 +89,7 @@ def process(self) -> None:
     def len(self):
         return (len(os.listdir(self.processed_dir)) - 4) // 2
 
-    def _build_tensors(self, root_url: str, df_to_dict, existing_urls) -> Tuple[Tensor, Tensor, Tensor, Tensor, dict]:
+    def _build_tensors(self, root_url: str, df_to_dict, existing_urls) -> tuple[Tensor, Tensor, Tensor, Tensor, dict]:
         """Builds the required tensors for one graph.
         These matrices will be then used for training the GNN.
 

phishGNN/dataset_v2.py

@@ -1,6 +1,5 @@
 import glob
 import os
-from typing import Tuple, List
 
 import pandas as pd
 import torch
@@ -44,12 +43,12 @@ def __init__(
         super(PhishingDataset2, self).__init__(root, transform, pre_transform)
 
     @property
-    def raw_file_names(self) -> List[str]:
+    def raw_file_names(self) -> list[str]:
         """File name of the csv dataset. """
         return glob.glob(os.path.join(self.raw_dir, '*'))
 
     @property
-    def processed_file_names(self) -> List[str]:
+    def processed_file_names(self) -> list[str]:
         return [file + '.pt' for file in self.raw_file_names]
 
     @property
@@ -105,7 +104,7 @@ def process(self) -> None:
     def len(self):
         return (len(os.listdir(self.processed_dir)) - 4) // 2
 
-    def _build_tensors(self, root_url: str, df_to_dict, existing_urls) -> Tuple[Tensor, Tensor, Tensor, Tensor, dict]:
+    def _build_tensors(self, root_url: str, df_to_dict, existing_urls) -> tuple[Tensor, Tensor, Tensor, Tensor, dict]:
         """Builds the required tensors for one graph.
         These matrices will be then used for training the GNN.