Small fixes

Author: fatemetkl

src/midst_toolkit/attacks/black_box_single_table/ensemble_mia/config.py

@@ -1,16 +1,14 @@
 from pathlib import Path
 
+
 BASE_DATA_DIR = Path("midst_toolkit/attacks/black_box_single_table/ensemble_mia/data")
 
 DATA_CONFIG = {
     # Data processing paths and file names
     ## Input directories:
-    "midst_data_path": (
-        BASE_DATA_DIR / "midst_data_all_attacks"
-    ),  # Used only for reading the data
+    "midst_data_path": (BASE_DATA_DIR / "midst_data_all_attacks"),  # Used only for reading the data
     ## Output directories:
-    "population_path": BASE_DATA_DIR
-    / "population_data",  # Path where the population data is stored
+    "population_path": BASE_DATA_DIR / "population_data",  # Path where the population data is stored
     "processed_attack_data_path": (
         BASE_DATA_DIR / "attack_data"
     ),  # Path where the processed attack real train and evaluation data is stored

src/midst_toolkit/attacks/black_box_single_table/ensemble_mia/data_processing/data_utils.py

@@ -1,8 +1,10 @@
 from logging import INFO
 from pathlib import Path
-from midst_toolkit.common.logger import log
+
 import pandas as pd
 
+from midst_toolkit.common.logger import log
+
 
 def save_dataframe(df: pd.DataFrame, file_path: Path, file_name: str) -> None:
     """
@@ -47,6 +49,7 @@ def collect_midst_attack_data(
         data_dir (Path): The path where the data is stored.
         data_split (str): Indicates if this is train, dev, or final data.
         dataset (str): The dataset to be collected. Either "train" or "challenge".
+        data_config (dict): Configuration dictionary containing data paths and file names.
 
     Returns:
         pd.DataFrame: The specified dataset in this setting.
@@ -67,16 +70,14 @@ def collect_midst_attack_data(
         # Multi-table attacks have different file names.
         file_name = (
             data_config["multi_table_train_data_file_name"]
-            if "clavaddpm" == generation_name
+            if generation_name == "clavaddpm"
             else data_config["single_table_train_data_file_name"]
         )
     assert file_name.split(".")[-1] == "csv", "File name should end with .csv."
 
     df_real = pd.DataFrame()
     for i in data_id:
-        data_dir_ith = (
-            data_dir / attack_type / data_split / f"{generation_name}_{i}" / file_name
-        )
+        data_dir_ith = data_dir / attack_type / data_split / f"{generation_name}_{i}" / file_name
         df_real_ith = pd.read_csv(data_dir_ith)
         df_real = df_real_ith if i == 1 else pd.concat([df_real, df_real_ith])
 

src/midst_toolkit/attacks/black_box_single_table/ensemble_mia/data_processing/process_split_data.py

@@ -1,11 +1,9 @@
 from logging import INFO
-from pathlib import Path
 
 import numpy as np
 import pandas as pd
 from sklearn.model_selection import train_test_split
 
-from midst_toolkit.common.logger import log
 from midst_toolkit.attacks.black_box_single_table.ensemble_mia.config import (
     seed,
 )
@@ -15,6 +13,7 @@
 from midst_toolkit.attacks.black_box_single_table.ensemble_mia.data_processing.real_data_collection import (
     collect_population_data_ensemble_mia,
 )
+from midst_toolkit.common.logger import log
 
 
 def split_real_data(
@@ -68,7 +67,7 @@ def generate_val_test(
     seed: int,
 ) -> tuple[pd.DataFrame, np.ndarray, pd.DataFrame, np.ndarray]:
     """
-    Generates the validation and test sets with labels. 
+    Generates the validation and test sets with labels.
     The resulting validation and test sets are used for meta classifier training and evaluation, respectively.
 
     Args:
@@ -208,4 +207,5 @@ def process_split_data(
     from midst_toolkit.attacks.black_box_single_table.ensemble_mia.config import (
         DATA_CONFIG,
     )
+
     process_split_data(data_config=DATA_CONFIG)

src/midst_toolkit/attacks/black_box_single_table/ensemble_mia/data_processing/real_data_collection.py

@@ -21,9 +21,10 @@ def collect_midst_data(
 
     Args:
         attack_types (list[str]): List of attack names to be collected.
-        data_splits (list[str]): A list indicating the data split to be collected. 
+        data_splits (list[str]): A list indicating the data split to be collected.
             Could be any of train, dev, or final data splits.
         dataset (str): The dataset to be collected. Either "train" or "challenge".
+        data_config (dict): Configuration dictionary containing data paths and file names.
 
     Returns:
         pd.DataFrame: Collected train or challenge data as a DataFrame.
@@ -60,10 +61,10 @@ def collect_population_data_ensemble_mia(
     and returned as a dataframe.
 
     Args:
-        data_processing_config (dict): Configuration dictionary containing data paths and file names.
+        data_config (dict): Configuration dictionary containing data paths and file names.
         attack_types (list[str] | None): List of attack names to be collected.
             If None, all the attacks are collected based on ensemble mia implementation.
-    
+
     Returns:
         pd.DataFrame: The collected population data.
     """
@@ -72,17 +73,15 @@ def collect_population_data_ensemble_mia(
     # Collect train data of all the attacks (back box and white box)
     if attack_types is None:
         attack_types = [
-        "tabddpm_black_box",
-        "tabsyn_black_box",
-        "tabddpm_white_box",
-        "tabsyn_white_box", 
-        "clavaddpm_black_box",
-        "clavaddpm_white_box",
-    ]
-
-    df_population = collect_midst_data(
-        attack_types, data_splits=["train"], dataset="train", data_config=data_config
-    )
+            "tabddpm_black_box",
+            "tabsyn_black_box",
+            "tabddpm_white_box",
+            "tabsyn_white_box",
+            "clavaddpm_black_box",
+            "clavaddpm_white_box",
+        ]
+
+    df_population = collect_midst_data(attack_types, data_splits=["train"], dataset="train", data_config=data_config)
     # Drop ids.
     df_population_no_id = df_population.drop(columns=["trans_id", "account_id"])
     # Save the population data

tests/unit/attacks/ensemble_mia/assets/midst_data_all_attacks/tabddpm_black_box/train/tabddpm_2/trans_domain.json

@@ -1 +1 @@
-{"trans_date": {"size": 2191, "type": "continuous"}, "trans_type": {"size": 3, "type": "discrete"}, "operation": {"size": 6, "type": "discrete"}, "amount": {"size": 40400, "type": "continuous"}, "balance": {"size": 542739, "type": "continuous"}, "k_symbol": {"size": 9, "type": "discrete"}, "bank": {"size": 14, "type": "discrete"}, "account": {"size": 7665, "type": "continuous"}}
+{"trans_date": {"size": 2191, "type": "continuous"}, "trans_type": {"size": 3, "type": "discrete"}, "operation": {"size": 6, "type": "discrete"}, "amount": {"size": 40400, "type": "continuous"}, "balance": {"size": 542739, "type": "continuous"}, "k_symbol": {"size": 9, "type": "discrete"}, "bank": {"size": 14, "type": "discrete"}, "account": {"size": 7665, "type": "continuous"}}

tests/unit/attacks/ensemble_mia/config.py

@@ -1,14 +1,14 @@
 from pathlib import Path
 
+
 BASE_DATA_DIR = Path("tests/unit/attacks/ensemble_mia/assets")
 
 DATA_CONFIG = {
     # Data processing paths and file names
     ## Input directories:
     "midst_data_path": BASE_DATA_DIR / "midst_data_all_attacks",  # Used only for reading the data
     ## Output directories:
-    "population_path": BASE_DATA_DIR
-    / "population_data",  # Path where the population data is stored
+    "population_path": BASE_DATA_DIR / "population_data",  # Path where the population data is stored
     "processed_attack_data_path": (
         BASE_DATA_DIR / "attack_data"
     ),  # Path where the processed attack real train and evaluation data is stored

tests/unit/attacks/ensemble_mia/test_data_collection.py

@@ -1,10 +1,12 @@
 from pathlib import Path
+
 from src.midst_toolkit.attacks.black_box_single_table.ensemble_mia.data_processing.real_data_collection import (
     collect_midst_data,
     collect_population_data_ensemble_mia,
 )
 from tests.unit.attacks.ensemble_mia.config import DATA_CONFIG
 
+
 def test_collect_population_data_ensemble_mia(tmp_path: Path) -> None:
     # Comment the next line to update population data stored in DATA_CONFIG["population_path"].
     DATA_CONFIG["population_path"] = tmp_path
@@ -21,13 +23,9 @@ def test_collect_population_data_ensemble_mia(tmp_path: Path) -> None:
 
     assert (DATA_CONFIG["population_path"] / "population_all_no_challenge.csv").exists()
 
-    assert (
-        DATA_CONFIG["population_path"] / "population_all_with_challenge.csv"
-    ).exists()
+    assert (DATA_CONFIG["population_path"] / "population_all_with_challenge.csv").exists()
 
-    assert (
-        DATA_CONFIG["population_path"] / "population_all_with_challenge_no_id.csv"
-    ).exists()
+    assert (DATA_CONFIG["population_path"] / "population_all_with_challenge_no_id.csv").exists()
 
 
 def test_collect_midst_data() -> None:

tests/unit/attacks/ensemble_mia/test_process_data_split.py

@@ -1,10 +1,11 @@
 from pathlib import Path
+
+from src.midst_toolkit.attacks.black_box_single_table.ensemble_mia.data_processing.data_utils import load_dataframe
 from src.midst_toolkit.attacks.black_box_single_table.ensemble_mia.data_processing.process_split_data import (
     process_split_data,
 )
 from tests.unit.attacks.ensemble_mia.config import DATA_CONFIG
 
-from src.midst_toolkit.attacks.black_box_single_table.ensemble_mia.data_processing.data_utils import load_dataframe
 
 def test_process_split_data(tmp_path: Path) -> None:
     # Comment the next line to update processed attack data stored in DATA_CONFIG["processed_attack_data_path"].
@@ -38,15 +39,11 @@ def test_process_split_data(tmp_path: Path) -> None:
     # Recall that `master_challenge_train`` consists of two halves: one half (10k) from `real_val`` data
     # with their "is_train" column set to 0, and the other half (10k) from the real train data (`real_train``)
     # with their "is_train" column set to 1. Note that ["is_train"] column is dropped in the final dataframes.
-    master_challenge_train = load_dataframe(
-        DATA_CONFIG["processed_attack_data_path"], "master_challenge_train.csv"
-    )
+    master_challenge_train = load_dataframe(DATA_CONFIG["processed_attack_data_path"], "master_challenge_train.csv")
     assert master_challenge_train.shape == (20000, 10), f" Shape is {master_challenge_train.shape}"
 
     # Recall that `master_challenge_test`` consists of two halves: one half (10k) from `real_test`` data
     # with their "is_train" column set to 0, and the other half (10k) from the real train data (`real_train``)
     # with their "is_train" column set to 1. Note that ["is_train"] column is dropped in the final dataframes.
-    master_challenge_test = load_dataframe(
-        DATA_CONFIG["processed_attack_data_path"], "master_challenge_test.csv"
-    )
+    master_challenge_test = load_dataframe(DATA_CONFIG["processed_attack_data_path"], "master_challenge_test.csv")
     assert master_challenge_test.shape == (20000, 10), f" Shape is {master_challenge_test.shape}"