EPT Attack: Feature Extraction (#89)

sarakodeiri · web-flow · commit 5782ca851bae · 2025-11-17T15:19:31.000-05:00
* run_ept_attack for feature extraction.

* Add pre_process_and_train in feature extraction

* Add tests

* Packaged directory iterations, switched to the library's ENUM, resolved other comments.
diff --git a/examples/common/utils.py b/examples/common/utils.py
@@ -0,0 +1,28 @@
+# Functions used for attacks across multiple examples.
+
+from collections.abc import Generator
+from pathlib import Path
+
+
+def iterate_model_folders(input_data_path: Path, diffusion_model_names: list[str]) -> Generator[tuple[str, Path, str]]:
+    """
+    Iterates over the competition's shadow model folder structure and yields model information.
+
+    Args:
+        input_data_path: The base path for the input data.
+        diffusion_model_names: A list of diffusion model names to iterate over.
+
+    Yields:
+        A tuple containing the model name, the path to the model's data, and the model folder name.
+    """
+    modes = ["train", "dev", "final"]
+    for model_name in diffusion_model_names:
+        model_path = input_data_path / f"{model_name}_black_box"
+        for mode in modes:
+            current_path = model_path / mode
+            if not current_path.exists():
+                continue
+
+            model_folders = [entry for entry in current_path.iterdir() if entry.is_dir()]
+            for model_folder_path in model_folders:
+                yield model_name, model_folder_path, model_folder_path.name
diff --git a/examples/ept_attack/config.yaml b/examples/ept_attack/config.yaml
@@ -0,0 +1,23 @@
+# Ensemble example configuration
+# Base data directory (can be overridden from command line)
+base_data_dir: examples/ept_attack/data/
+base_example_dir: examples/ept_attack
+
+# Data paths (relative to base_data_dir)
+data_paths:
+  input_data_path: ${base_data_dir}/midst_data_black_box_attacks # Read-only input data directory
+  output_data_path: ${base_data_dir}/output # Directory to save processed data and results
+  data_types_file_path: ${base_data_dir}/data_configs/data_types.json # Path to the JSON file defining column types
+
+# Pipeline control
+pipeline:
+  run_data_processing: false # Whether to run data processing
+  run_shadow_model_training: false # Whether to run shadow model training
+  run_feature_extraction: true # Whether to run attribute prediction model training
+  run_attack_classifier_training: false # Whether to run attack classifier training
+
+attack_settings:
+  single_table: true # Whether the data is single-table
+
+# General settings
+random_seed: 42
diff --git a/examples/ept_attack/run_ept_attack.py b/examples/ept_attack/run_ept_attack.py
@@ -0,0 +1,107 @@
+"""
+This file is an incomplete example script for running the EPT-MIA Attack on MIDST challenge
+provided resources and data.
+Overall workflow and decisions are taken with from the Cyber@BGU team's attack implementation at
+https://github.com/eyalgerman/MIA-EPT.
+
+"""
+
+import json
+from logging import INFO
+from pathlib import Path
+
+import hydra
+from omegaconf import DictConfig
+
+from examples.common.utils import iterate_model_folders
+from midst_toolkit.attacks.ensemble.data_utils import load_dataframe, save_dataframe
+from midst_toolkit.attacks.ept.feature_extraction import extract_features
+from midst_toolkit.common.logger import log
+
+
+# Step 2 and 3: Attribute prediction model training and feature extraction
+def run_attribute_prediction(config: DictConfig) -> None:
+    """
+    Train attribute prediction models and extract features for EPT-MIA attack.
+    The function is specifically designed to work with the MIDST challenge data structure,
+    and the shadow models provided by the competition organizers.
+    All the reading and writing of data is handled within this function.
+
+    Args:
+        config: Configuration object set in config.yaml.
+    """
+    log(INFO, "Running attribute prediction model training.")
+
+    diffusion_model_names = ["tabddpm", "tabsyn"] if config.attack_settings.single_table else ["clavaddpm"]
+    input_data_path = Path(config.data_paths.input_data_path)
+    output_features_path = Path(config.data_paths.output_data_path, "attribute_prediction_features")
+
+    # Load column types specific to the competition dataset
+    with open(config.data_paths.data_types_file_path, "r") as f:
+        column_types = json.load(f)
+
+    # Drop columns that end with '_id' from column_types, as they do not create meaningful features
+    feature_column_types = {
+        "numerical": [col for col in column_types.get("numerical", []) if not col.endswith("_id")],
+        "categorical": [col for col in column_types.get("categorical", []) if not col.endswith("_id")],
+    }
+
+    # TODO: Package iterating over competition structure (maybe into a utility function)
+    # Iterating over directories specific to the shadow models folder structure in the competition
+    for model_name, model_data_path, model_folder in iterate_model_folders(input_data_path, diffusion_model_names):
+        # Load the data files as dataframes
+        df_synthetic_data = load_dataframe(model_data_path, "trans_synthetic.csv")
+        df_challenge_data = load_dataframe(model_data_path, "challenge_with_id.csv")
+
+        # Keep only the columns that are present in feature_column_types
+        columns_to_keep = feature_column_types["numerical"] + feature_column_types["categorical"]
+        df_synthetic_data = df_synthetic_data[columns_to_keep]
+        df_challenge_data = df_challenge_data[columns_to_keep]
+
+        # Run feature extraction
+        df_extracted_features = extract_features(
+            synthetic_data=df_synthetic_data,
+            challenge_data=df_challenge_data,
+            column_types=feature_column_types,
+            random_seed=config.random_seed,
+        )
+
+        final_output_dir = output_features_path / f"{model_name}_black_box"
+
+        final_output_dir.mkdir(parents=True, exist_ok=True)
+
+        # Extract the number at the end of model_folder
+        model_folder_number = int(model_folder.split("_")[-1])
+        file_name = f"attribute_prediction_features_{model_folder_number}.csv"
+
+        save_dataframe(df=df_extracted_features, file_path=final_output_dir, file_name=file_name)
+
+
+@hydra.main(config_path=".", config_name="config", version_base=None)
+def main(config: DictConfig) -> None:
+    """
+    Main orchestrator of the EPT-MIA Attack example pipeline.
+    First step has yet to be implemented: shadow model training.
+    Second and third steps are attribute prediction model training and feature extraction.
+
+    Args:
+        config: Attack configuration as an OmegaConf DictConfig object.
+    """
+    log(INFO, "Running EPT-MIA Attack Example Pipeline.")
+
+    if config.attack_settings.single_table:
+        log(INFO, "Data: Single-table.")
+    else:
+        log(INFO, "Data: Multi-table.")
+
+    # TODO: Implement potential data preprocessing step.
+    # TODO: Implement shadow model training step.
+
+    if config.pipeline.run_feature_extraction:
+        run_attribute_prediction(config)
+
+    # TODO: Implement attack classifier training step.
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/midst_toolkit/attacks/ept/feature_extraction.py b/src/midst_toolkit/attacks/ept/feature_extraction.py
@@ -0,0 +1,198 @@
+"""
+Module to run feature extraction for EPT attack steps 2 and 3.
+Overall workflow and decisions are taken with from the Cyber@BGU team's attack implementation at
+https://github.com/eyalgerman/MIA-EPT.
+
+"""
+
+from logging import INFO
+
+import numpy as np
+import pandas as pd
+from sklearn.compose import ColumnTransformer
+from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
+from sklearn.pipeline import Pipeline
+from sklearn.preprocessing import OneHotEncoder, StandardScaler
+
+from midst_toolkit.common.enumerations import TaskType
+from midst_toolkit.common.logger import log
+
+
+def preprocess_train_predict(
+    train_points: pd.DataFrame,
+    test_points: pd.DataFrame,
+    target_col: str,
+    column_types: dict[str, list[str]],
+    random_seed: int | None = None,
+) -> tuple[np.ndarray, pd.Series, TaskType]:
+    """
+    An attribute prediction model is trained on `train_points` to predict the `target_col`.
+
+    We determine the nature of the prediction task based on the data type of the target column.
+    If the `target_col` is categorical, the model uses a classification approach. Otherwise, if
+    the `target_col` is numerical, a regression model is used. This allows the
+    model to effectively learn the relationship between the `target_col` and the other attributes
+    present in the training data.
+
+    After the model is trained on `train_points`, it is then used to generate predictions for the `target_col`
+    on `test_points`.
+
+    Args:
+        train_points: Data to train the attribute prediction model on. Must include the target column.
+        test_points: Data to test the attribute prediction model on. Must include the target column.
+        target_col: Name of the target column to predict.
+        column_types: Types of columns in the data. Relevant keys are "numerical", "categorical".
+        random_seed: Seed for model reproducibility. Defaults to None.
+
+    Returns:
+            predictions: Predicted values for the target column on the test data.
+            y_test: True values for the target column on the test data.
+            task_type: Whether the attribution prediction model was a classification or regression model.
+    """
+    assert target_col in train_points.columns, f"Target column '{target_col}' not found in train_points."
+    assert target_col in test_points.columns, f"Target column '{target_col}' not found in test_points."
+
+    assert set(train_points.columns) == set(test_points.columns), "Columns in df_train and df_test do not match"
+
+    x_train = train_points.drop([target_col], axis=1)
+    y_train = train_points[target_col]
+    x_test = test_points.drop([target_col], axis=1)
+    y_test = test_points[target_col]
+
+    numeric_columns = column_types["numerical"]
+    categorical_columns = column_types["categorical"]
+
+    # Assert that the target column appears exactly once in numeric_columns + categorical_columns
+    assert (numeric_columns + categorical_columns).count(target_col) == 1, (
+        f"The target column '{target_col}' must appear exactly once in numeric_columns + categorical_columns"
+    )
+
+    # Assert that the union of numeric_columns and categorical_columns matches the columns in train_points
+    assert set(numeric_columns + categorical_columns) == set(train_points.columns), (
+        "The union of numeric_columns and categorical_columns must match the columns in the combined dataframe"
+    )
+
+    task_type = TaskType.MULTICLASS_CLASSIFICATION if target_col in categorical_columns else TaskType.REGRESSION
+
+    # Remove target column from feature columns
+    numeric_columns = [col for col in numeric_columns if col != target_col]
+    categorical_columns = [col for col in categorical_columns if col != target_col]
+
+    numeric_transformer = StandardScaler()
+    categorical_transformer = OneHotEncoder(drop="first", handle_unknown="ignore")
+
+    preprocessor = ColumnTransformer(
+        transformers=[
+            ("num", numeric_transformer, numeric_columns),
+            ("cat", categorical_transformer, categorical_columns),
+        ]
+    )
+
+    model = (
+        RandomForestClassifier(random_state=random_seed)
+        if task_type == TaskType.MULTICLASS_CLASSIFICATION
+        else RandomForestRegressor(random_state=random_seed)
+    )
+
+    model_pipeline = Pipeline(steps=[("preprocessor", preprocessor), ("model", model)])
+
+    model_pipeline.fit(x_train, y_train)
+
+    predictions = model_pipeline.predict(x_test)
+
+    return predictions, y_test, task_type
+
+
+def extract_features(
+    synthetic_data: pd.DataFrame,
+    challenge_data: pd.DataFrame,
+    column_types: dict[str, list[str]],
+    random_seed: int | None = None,
+) -> pd.DataFrame:
+    """
+    Orchestrator function to run feature extraction for EPT attack:
+    1. For each attribute (column) in the synthetic data that is not an ID, train an attribute prediction model
+        using the synthetic data.
+    2. Use the trained model to predict the values of that attribute in the challenge data, which also doesn't
+        contain IDs.
+    3. Compute relevant metrics (accuracy for categorical data, error and error ratio for numerical data).
+    4. Compile the results into a DataFrame.
+
+    Args:
+        synthetic_data: Synthetic data to extract features from. Note: This data should not contain any identifier
+            columns, as the function will attempt to train a prediction model for every column included.
+        challenge_data: The data the predictions are compared against, to compute prediction accuracy/errors.
+        column_types: A dictionary specifying the types of columns (numerical or categorical) in the data.
+        random_seed: Random seed for reproducibility. Defaults to None.
+
+    Returns:
+        A DataFrame containing the extracted features for each attribute in the challenge data.
+        It includes the following columns:
+            - <column_name>: The true values for the attribute.
+            - <column_name>_prediction: The predicted values for the attribute.
+        If the data is categorical:
+            - <column_name>_accuracy: The element-wise accuracy of the predictions. 0 for incorrect prediction,
+                1 for correct.
+        If the data is numerical:
+            - <column_name>_error (if regression): The absolute errors of the predictions.
+            - <column_name>_error_ratio (if regression): The ratio of the errors to the true values, which is
+                derived by dividing the absolute error by the true value in a zero-safe manner.
+    """
+    features = []
+    columns = []
+
+    for column in synthetic_data.columns:
+        log(INFO, f"Extracting features for column: {column}")
+
+        predictions, y_test, task_type = preprocess_train_predict(
+            train_points=synthetic_data,
+            test_points=challenge_data,
+            target_col=column,
+            column_types=column_types,
+            random_seed=random_seed,
+        )
+
+        features.append(y_test)
+        columns.append(column)
+
+        if task_type == TaskType.MULTICLASS_CLASSIFICATION:
+            # TODO: Maybe change the variable name from accuracy to correctness
+            # Calculate accuracy
+            accuracy = predictions == y_test
+            accuracy = accuracy.astype(int)
+            features.append(accuracy)
+            columns.append(f"{column}_accuracy")
+
+        elif task_type == TaskType.REGRESSION:
+            # Calculate errors
+            errors = pd.Series(np.abs(predictions - y_test), index=y_test.index)
+
+            # Calculate the ratio of the error in a zero-safe manner
+            denominator = y_test.replace(0, np.nan)
+            error_ratio = errors / np.abs(denominator)
+
+            # Replace infs and NaNs with a large number. If all values are NaN, replace with 1e9.
+            finite_max = error_ratio[np.isfinite(error_ratio)].max()
+            error_ratio = error_ratio.replace([np.inf, -np.inf], np.nan).fillna(
+                finite_max if pd.notna(finite_max) else 1e9
+            )
+
+            # Save the error and the ratio error
+            features.append(errors)
+            features.append(error_ratio)
+
+            columns.append(f"{column}_error")
+            columns.append(f"{column}_error_ratio")
+
+        else:
+            raise ValueError(f"Unsupported task type: {task_type}")
+
+        # predictions from the model
+        features.append(pd.Series(predictions, index=y_test.index))
+        columns.append(f"{column}_prediction")
+
+        # Create a DataFrame with the results
+    df_results = pd.DataFrame(features).T
+    df_results.columns = columns
+
+    return df_results
diff --git a/tests/unit/attacks/ept_attack/test_feature_extraction.py b/tests/unit/attacks/ept_attack/test_feature_extraction.py
