Sk/ept classifier

.gitignore

@@ -51,3 +51,6 @@ tests/integration/attacks/tartan_federer/assets/tartan_federer_attack_results
 # Training Logs
 *.err
 *.out
+
+# CatBoost training artifacts
+catboost_info/

examples/common/utils.py

@@ -4,7 +4,9 @@
 from pathlib import Path
 
 
-def iterate_model_folders(input_data_path: Path, diffusion_model_names: list[str]) -> Generator[tuple[str, Path, str]]:
+def iterate_model_folders(
+    input_data_path: Path, diffusion_model_names: list[str]
+) -> Generator[tuple[str, Path, str, str]]:
     """
     Iterates over the competition's shadow model folder structure and yields model information.
 
@@ -13,7 +15,8 @@ def iterate_model_folders(input_data_path: Path, diffusion_model_names: list[str
         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.
+        A tuple containing the model name (e.g. tabddpm), the path to the model's data,
+        the model folder name (e.g. tabddpm_1), and mode (train, dev, final).
     """
     modes = ["train", "dev", "final"]
     for model_name in diffusion_model_names:
@@ -25,4 +28,21 @@ def iterate_model_folders(input_data_path: Path, diffusion_model_names: list[str
 
             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
+                yield model_name, model_folder_path, model_folder_path.name, mode
+
+
+def directory_checks(directory_path: Path, custom_error_message: str = "") -> None:
+    """
+    Performs checks to ensure that the given directory exists and is not empty.
+
+    Args:
+        directory_path: Path to the directory to check.
+        custom_error_message: Additional message to include in the assertion error.
+
+    Raises:
+        AssertionError: If the directory does not exist or is empty.
+    """
+    assert directory_path.exists() and directory_path.is_dir(), (
+        f"Directory not found: {directory_path}. " + custom_error_message
+    )
+    assert any(directory_path.iterdir()), f"Directory is empty: {directory_path}. " + custom_error_message

examples/ept_attack/config.yaml

@@ -1,20 +1,23 @@
 # Ensemble example configuration
 # Base data directory (can be overridden from command line)
-base_data_dir: examples/ept_attack/data/
+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
-
+  attribute_features_path: ${data_paths.output_data_path}/attribute_prediction_features # Path to save attribute prediction features
 # 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
+  run_feature_extraction: false # Whether to run attribute prediction model training
+  run_attack_classifier_training: true # Whether to run attack classifier training
+
+classifier_settings:
+  results_output_path: ${data_paths.output_data_path}/evaluation_ML
 
 attack_settings:
   single_table: true # Whether the data is single-table

examples/ept_attack/run_ept_attack.py

@@ -6,17 +6,23 @@
 
 """
 
+import itertools
 import json
+from collections import defaultdict
+from datetime import datetime
 from logging import INFO
 from pathlib import Path
 
 import hydra
+import pandas as pd
 from omegaconf import DictConfig
 
-from examples.common.utils import iterate_model_folders
+from examples.common.utils import directory_checks, iterate_model_folders
 from midst_toolkit.attacks.ensemble.data_utils import load_dataframe, save_dataframe
+from midst_toolkit.attacks.ept.classification import ClassifierType, train_attack_classifier
 from midst_toolkit.attacks.ept.feature_extraction import extract_features
 from midst_toolkit.common.logger import log
+from midst_toolkit.common.random import set_all_random_seeds
 
 
 # Step 2 and 3: Attribute prediction model training and feature extraction
@@ -34,7 +40,7 @@ def run_attribute_prediction(config: DictConfig) -> None:
 
     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")
+    output_features_path = Path(config.data_paths.attribute_features_path)
 
     # Load column types specific to the competition dataset
     with open(config.data_paths.data_types_file_path, "r") as f:
@@ -46,8 +52,15 @@ def run_attribute_prediction(config: DictConfig) -> None:
         "categorical": [col for col in column_types.get("categorical", []) if not col.endswith("_id")],
     }
 
+    # Assert that the input data path exists and is not empty
+    directory_checks(input_data_path)
+
     # 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):
+    for model_name, model_data_path, model_folder, mode in iterate_model_folders(
+        input_data_path, diffusion_model_names
+    ):
+        log(INFO, f"Processing model: {model_name}, path: {model_data_path}, folder: {model_folder}, mode: {mode}")
+
         # 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")
@@ -65,17 +78,197 @@ def run_attribute_prediction(config: DictConfig) -> None:
             random_seed=config.random_seed,
         )
 
-        final_output_dir = output_features_path / f"{model_name}_black_box"
+        final_output_dir = output_features_path / f"{model_name}_black_box" / mode
 
         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"
 
+        if mode == "train":
+            file_name = f"attribute_prediction_features_with_labels_{model_folder_number}.csv"
+
+            # Load the challenge labels and add them to the features dataframe
+            df_labels = load_dataframe(model_data_path, "challenge_label.csv")
+
+            # Check that the number of rows align
+            assert len(df_extracted_features) == len(df_labels), (
+                f"The number of rows in the extracted features ({len(df_extracted_features)}) "
+                f"does not match the number of labels ({len(df_labels)})."
+            )
+            df_extracted_features["is_train"] = df_labels.values
+
         save_dataframe(df=df_extracted_features, file_path=final_output_dir, file_name=file_name)
 
 
+def _summarize_and_save_training_results(
+    summary_results: dict, output_summary_path: Path, summary_file_name: str
+) -> pd.DataFrame:
+    """
+    Processes summary results, saves them to a CSV, and returns the summary DataFrame.
+
+    Args:
+        summary_results: A dictionary containing the summary results.
+        output_summary_path: The path where the summary CSV will be saved.
+        summary_file_name: The name of the summary CSV file.
+
+    Returns:
+        A pandas DataFrame containing the summarized results.
+    """
+    processed_results = []
+    for (classifier, columns_lst), model_scores in summary_results.items():
+        row: dict[str, str | float] = {"classifier": classifier, "column_types": columns_lst}
+        for diffusion_model_name, scores in model_scores:
+            for score_name, score_value in scores.items():
+                col_name = (
+                    score_name.lower().replace(" ", "_").replace("-", "_").replace("_at_", "_").replace(".0", "")
+                )
+                row[f"{diffusion_model_name}_{col_name}"] = score_value
+        processed_results.append(row)
+
+    summary_df = pd.DataFrame(processed_results)
+    tpr_10_cols = [col for col in summary_df.columns if col.endswith("_tpr_fpr_10")]
+    if tpr_10_cols:
+        summary_df["final_tpr_fpr_10"] = summary_df[tpr_10_cols].max(axis=1)
+
+    summary_df.to_csv(output_summary_path / summary_file_name, index=False)
+    log(INFO, f"Saved attack classifier summary to {output_summary_path / summary_file_name}")
+    return summary_df
+
+
+# Step 4: Attack classifier training
+def run_attack_classifier_training(config: DictConfig) -> None:
+    """
+    Trains multiple attack classifiers to distinguish between training and
+    non-training data, and selects the best performing configuration based
+    on evaluation metrics.
+
+    This function orchestrates the training of various attack classifiers (XGBoost,
+    CatBoost, MLP) to perform a membership inference attack. It iterates through
+    different diffusion models used to generate synthetic data and all combinations
+    of feature types derived from the attribute prediction task.
+
+    The process involves:
+    1.  Reading pre-computed feature files generated by the feature extraction step.
+    2.  Splitting the feature files into training and testing sets.
+    3.  For each diffusion model, iterating through all possible combinations of
+        feature columns ('actual', 'error', 'error_ratio', 'accuracy', 'prediction').
+    4.  Training each classifier type on these feature combinations.
+    5.  Evaluating the classifier's performance and saving the scores (e.g., AUC, TPR at
+        specific FPR) and prediction results for each configuration.
+    6.  Aggregating all results into a summary CSV file, which includes a final
+        metric ('final_tpr_fpr_10') representing the best TPR at 10% FPR across
+        all diffusion models for a given classifier and feature set.
+    7.  Logging the best-performing attack configuration based on this final metric.
+
+    Args:
+        config: Configuration object set in config.yaml.
+    """
+    log(INFO, "Running attack classifier training.")
+
+    data_format, diffusion_models = (
+        ("single_table", ["tabddpm", "tabsyn"])
+        if config.attack_settings.single_table
+        else ("multi_table", ["clavaddpm"])
+    )
+
+    # Read all the files from the attribute prediction features directory
+    features_data_path = Path(config.data_paths.attribute_features_path)
+
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+
+    # An example of summary_results structure:
+    # {
+    #   ('XGBoost', 'actual error'): [
+    #       ('tabddpm', {'AUC': 0.85, 'TPR at FPR=10%': 0.75, ...}),
+    #       ('tabsyn', {'AUC': 0.80, 'TPR at FPR=10%': 0.70, ...}),
+    #   ],
+    #   ('CatBoost', 'accuracy prediction'): [
+    #       ('tabddpm', {'AUC': 0.82, 'TPR at FPR=10%': 0.72, ...}),
+    #       ('tabsyn', {'AUC': 0.78, 'TPR at FPR=10%    ': 0.68, ...}),
+    #   ],
+    #   ...
+    # }
+
+    summary_results: dict[tuple[str, str], list[tuple[str, dict[str, float]]]] = defaultdict(list)
+
+    for diffusion_model_name in diffusion_models:
+        train_features_path = features_data_path / f"{diffusion_model_name}_black_box" / "train"
+
+        directory_checks(train_features_path, "Make sure to run feature extraction first.")
+
+        sorted_feature_files = sorted(train_features_path.glob("*.csv"))
+        split_index = len(sorted_feature_files) * 5 // 6
+
+        # Get the first 25 feature files
+        train_feature_files = sorted_feature_files[:split_index]
+        # Concatenate all the train feature files into a single dataframe
+        df_train_features = pd.concat([pd.read_csv(f) for f in train_feature_files], ignore_index=True)
+        train_labels = df_train_features["is_train"]
+        df_train_features = df_train_features.drop(columns=["is_train"])
+
+        test_feature_files = sorted_feature_files[split_index:]
+        df_test_features = pd.concat([pd.read_csv(f) for f in test_feature_files], ignore_index=True)
+        test_labels = df_test_features["is_train"]
+        df_test_features = df_test_features.drop(columns=["is_train"])
+
+        classifier_types = ["XGBoost", "CatBoost", "MLP"]
+        column_types = ["actual", "error", "error_ratio", "accuracy", "prediction"]
+
+        output_summary_path = Path(config.classifier_settings.results_output_path) / data_format / f"{timestamp}_train"
+        output_summary_path.mkdir(parents=True, exist_ok=True)
+
+        for classifier in classifier_types:
+            for r in range(1, len(column_types) + 1):
+                for selected_column_types_tuple in itertools.combinations(column_types, r):
+                    columns_str = " ".join(sorted(selected_column_types_tuple))
+                    result_key = (classifier, columns_str)
+
+                    classifier_type = ClassifierType(classifier)
+
+                    results = train_attack_classifier(
+                        classifier_type=classifier_type,
+                        column_types=list(selected_column_types_tuple),
+                        x_train=df_train_features,
+                        y_train=train_labels,
+                        x_test=df_test_features,
+                        y_test=test_labels,
+                    )
+
+                    # Store raw scores for the current diffusion model
+                    summary_results[result_key].append((diffusion_model_name, results["scores"]))
+
+                    training_directory_name = f"{classifier}_{'_'.join(selected_column_types_tuple)}"
+                    training_output_path = output_summary_path / training_directory_name
+                    training_output_path.mkdir(parents=True, exist_ok=True)
+
+                    # Save prediction results
+                    prediction_results_df = results["prediction_results"]
+                    prediction_results_file_name = f"{diffusion_model_name}_prediction_results.csv"
+                    save_dataframe(
+                        df=pd.DataFrame(prediction_results_df),
+                        file_path=training_output_path,
+                        file_name=prediction_results_file_name,
+                    )
+
+                    # Save scores
+                    scores_file_name = f"{diffusion_model_name}_results.txt"
+                    with open(training_output_path / scores_file_name, "w") as f:
+                        for score_name, score_value in results["scores"].items():
+                            f.write(f"{score_name}: {score_value}\n")
+
+    summary_df = _summarize_and_save_training_results(
+        summary_results, output_summary_path, "attack_classifier_summary.csv"
+    )
+
+    summary_df.sort_values(by=["final_tpr_fpr_10"], ascending=False, inplace=True)
+    best_result = summary_df.head(1)
+    log(INFO, f"Best performing attack configuration:\n{best_result}")
+
+    log(INFO, f"Best performing attack configuration:\n{best_result}")
+
+
 @hydra.main(config_path=".", config_name="config", version_base=None)
 def main(config: DictConfig) -> None:
     """
@@ -88,18 +281,22 @@ def main(config: DictConfig) -> None:
     """
     log(INFO, "Running EPT-MIA Attack Example Pipeline.")
 
+    if config.random_seed is not None:
+        set_all_random_seeds(seed=config.random_seed)
+        log(INFO, f"Training phase random seed set to {config.random_seed}.")
+
     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 config.pipeline.run_attack_classifier_training:
+        run_attack_classifier_training(config)
 
 
 if __name__ == "__main__":

mypy.ini

@@ -42,3 +42,6 @@ ignore_missing_imports = True
 
 [mypy-faiss.*]
 ignore_missing_imports = True
+
+[mypy-catboost.*]
+ignore_missing_imports = True

pyproject.toml

@@ -28,6 +28,8 @@ dependencies = [
     "pydantic>=2.12.3",
     "filelock>=3.20.1",
     "sdv>=1.18.0",
+    "catboost>=1.2.8",
+    "scipy==1.16.3",
 ]
 
 [build-system]