Sk/ept classifier

[sarakodeiri]

PR Type

[Feature]

Short Description

Clickup Ticket(s): https://app.clickup.com/t/868g60gtc

Implemented classification training (on three different classifiers) and selection based on the submission repository.
In run_ept_attack.py, the codes iterates through different diffusion models used to generate synthetic data and ALL combinations of feature types derived from the attribute prediction task. I think this part is part of the main design of their attack and should be included in the library rather than the example, but handling saving and writing files would be a nightmare. Let me know how you think this is best handled.

Also currently everything is identical to the submission repository. There's a lot of saving results going on which I find redundant so far and might remove later, if they're proven to be truly redundant.

Tests Added

Added checks and tests at tests/unit/attacks/ept_attack/test_classification.py

[coderabbitai]

📝 Walkthrough

Walkthrough

This pull request introduces an attack classifier training pipeline for membership inference attacks in the EPT framework. A new classification module is added with functions for feature filtering, multiple model implementations (XGBoost, CatBoost, MLP neural network), and scoring logic. The main orchestration script is updated to iterate through model modes and execute classifier training. Configuration files are modified to enable the new training workflow, CatBoost is added as a runtime dependency, and comprehensive unit tests cover the classification functionality. A training log file records a CatBoost execution.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

🚥 Pre-merge checks | ✅ 1 | ❌ 2

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Title check	❓ Inconclusive	The title 'Sk/ept classifier' is vague and does not clearly convey the main purpose of the PR, using only branch naming convention without meaningful description.	Use a more descriptive title that clearly explains the primary change, such as 'Implement classification training for EPT attack with multiple classifier types' or 'Add attack classifier training for EPT membership inference attack'.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Description check	✅ Passed	The description follows the template structure and provides context, including PR type, short description, ticket link, and tests added, meeting core requirements.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

• 📝 Generate docstrings

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 3

🤖 Fix all issues with AI agents

In `@catboost_info/catboost_training.json`:
- Around line 1-1004: This commit includes an auto-generated CatBoost training
artifact (the catboost_info directory); add "catboost_info/" to .gitignore, run
git rm -r --cached catboost_info/ (or git rm -r catboost_info/) to remove the
tracked artifact from the index, and commit the .gitignore update and removal so
future CatBoost runs do not reintroduce the file.

In `@examples/ept_attack/run_ept_attack.py`:
- Around line 210-212: Rename the misspelled variables trianing_directory_name
and trianing_output_path to training_directory_name and training_output_path
respectively, and update all subsequent references in the same scope/function
(e.g., the later uses around where those variables are referenced) to use the
corrected names so the directory creation and downstream code use
training_directory_name and training_output_path consistently.
- Around line 159-168: The split currently skips the file at index 25 because
train_feature_files = sorted_feature_files[:25] uses indices 0–24 while
test_feature_files = sorted_feature_files[26:] starts at 26; update
test_feature_files to use sorted_feature_files[25:] (referencing
train_feature_files, test_feature_files, and sorted_feature_files) so the file
at index 25 is included in the test set and no file is omitted.

🧹 Nitpick comments (4)

examples/ept_attack/run_ept_attack.py (1)

156-168: Consider validating file count before splitting.

If sorted_feature_files contains fewer than 26 files, test_feature_files will be empty, and pd.concat([]) will raise a ValueError. Consider adding a validation check.

♻️ Suggested validation

         sorted_feature_files = sorted(train_features_path.glob("*.csv"))

+        min_required_files = 26  # At least 25 train + 1 test
+        assert len(sorted_feature_files) >= min_required_files, (
+            f"Expected at least {min_required_files} feature files, found {len(sorted_feature_files)}"
+        )
+
         # Get the first 25 feature files
         train_feature_files = sorted_feature_files[:25]

src/midst_toolkit/attacks/ept/classification.py (2)

50-76: LGTM!

The MLP architecture is appropriate for binary classification with Sigmoid output for probability estimation. Minor style note: super().__init__() is the preferred Python 3 idiom over super(MLPClassifier, self).__init__().

---

79-86: Parameter eval shadows Python built-in.

The parameter name eval shadows Python's built-in eval() function. While not causing a runtime error here, it's a best practice violation that could cause confusion.

♻️ Suggested rename

 def train_mlp(
     x_train: np.ndarray,
     y_train: np.ndarray,
     x_test: np.ndarray,
     device: torch.device,
-    eval: bool,
+    evaluate: bool,
     epochs: int = 10,
 ) -> tuple[np.ndarray | None, np.ndarray | None]:

Then update usages at lines 113, 127 to use evaluate instead of eval.

tests/unit/attacks/ept_attack/test_classification.py (1)

183-229: LGTM!

Excellent error handling test coverage for data mismatches and unsupported classifier types. The MLP integration test properly verifies the eval=True argument.

Minor style note: Line 223 is missing a blank line before the function definition per PEP 8.

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between d765bba and 85f03de.

⛔ Files ignored due to path filters (3)

• catboost_info/learn_error.tsv is excluded by !**/*.tsv
• catboost_info/time_left.tsv is excluded by !**/*.tsv
• uv.lock is excluded by !**/*.lock

📒 Files selected for processing (9)

• catboost_info/catboost_training.json
• catboost_info/learn/events.out.tfevents
• examples/common/utils.py
• examples/ept_attack/config.yaml
• examples/ept_attack/run_ept_attack.py
• mypy.ini
• pyproject.toml
• src/midst_toolkit/attacks/ept/classification.py
• tests/unit/attacks/ept_attack/test_classification.py

🧰 Additional context used

🧠 Learnings (1)

📚 Learning: 2025-12-11T16:08:49.024Z

Learnt from: lotif
Repo: VectorInstitute/midst-toolkit PR: 107
File: examples/gan/synthesize.py:1-47
Timestamp: 2025-12-11T16:08:49.024Z
Learning: When using SDV (version >= 1.18.0), prefer loading a saved CTGANSynthesizer with CTGANSynthesizer.load(filepath) instead of sdv.utils.load_synthesizer(). This applies to Python code across the repo (e.g., any script that loads a CTGANSynthesizer). Ensure the SDV version is >= 1.18.0 before using CTGANSynthesizer.load, and fall back to sdv.utils.load_synthesizer() only if a compatible alternative is required.

Applied to files:

• tests/unit/attacks/ept_attack/test_classification.py
• examples/ept_attack/run_ept_attack.py
• examples/common/utils.py
• src/midst_toolkit/attacks/ept/classification.py

🧬 Code graph analysis (1)

examples/ept_attack/run_ept_attack.py (4)

examples/common/utils.py (1)

• iterate_model_folders (7-31)

src/midst_toolkit/attacks/ensemble/data_utils.py (2)

• load_dataframe (31-52)
• save_dataframe (9-28)

src/midst_toolkit/attacks/ept/classification.py (1)

• train_attack_classifier (182-271)

src/midst_toolkit/attacks/ept/feature_extraction.py (1)

• extract_features (106-198)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)

• GitHub Check: integration-tests
• GitHub Check: unit-tests
• GitHub Check: run-code-check

🔇 Additional comments (19)

mypy.ini (1)

45-47: LGTM!

The mypy configuration for CatBoost follows the established pattern for third-party libraries without type stubs, consistent with the other ignore stanzas in this file.

examples/ept_attack/config.yaml (2)

3-3: LGTM!

Removing the trailing slash from base_data_dir improves consistency in path handling across the configuration.

---

11-20: LGTM!

The new configuration paths and classifier settings are well-structured with proper YAML variable interpolation. The pipeline control flags clearly indicate the intended workflow stage.

examples/common/utils.py (1)

7-31: Approved. The function signature change is properly handled—the only caller in the codebase (examples/ept_attack/run_ept_attack.py) already unpacks all four tuple elements correctly. The implementation is clean with defensive path existence checks.

pyproject.toml (1)

31-31: CatBoost 1.2.8 is the current release and has no known vulnerabilities.

The dependency addition is appropriate. Version 1.2.8 (released April 13, 2025) is the latest available, and the >=1.2.8 constraint correctly ensures security updates are included while locking to the stable release.

examples/ept_attack/run_ept_attack.py (4)

9-24: LGTM!

The new imports are appropriate for the added functionality: itertools for combinations, datetime for timestamping results, typing.Any for type hints, and pandas for data manipulation.

---

54-57: LGTM!

Good defensive programming with assertions to validate input data path existence and non-emptiness before processing.

---

89-101: LGTM!

The mode-specific handling for "train" data correctly loads labels and joins them with extracted features, with proper row count validation.

---

268-269: LGTM!

Clean integration into the main pipeline following the existing pattern.

src/midst_toolkit/attacks/ept/classification.py (4)

1-13: LGTM!

Imports are appropriate for the module's functionality. CatBoost and XGBoost for tree-based classifiers, PyTorch for the MLP, and sklearn for metrics.

---

15-47: LGTM!

The suffix-based column filtering logic is clean and the lambda mapping approach is readable. The "actual" filter correctly excludes columns with specific suffixes.

---

139-179: LGTM!

The scoring function correctly computes standard classification metrics and handles edge cases (empty valid_tpr returns 0.0).

---

182-271: LGTM!

The function properly orchestrates classifier training with good input validation and supports multiple classifier types. The defensive assertion at line 262 ensures predictions exist before scoring.

tests/unit/attacks/ept_attack/test_classification.py (6)

1-28: LGTM!

Good test setup with appropriate imports and a well-structured fixture for testing filter_data with various column types.

---

30-51: LGTM!

Excellent parametrized test coverage for filter_data including edge case with empty column list and multiple suffix combinations.

---

54-72: LGTM!

Thorough test of MLPClassifier including layer structure verification and forward pass validation with activation bounds checking.

---

75-118: LGTM!

Good mocking strategy to test train_mlp logic without actual model training. Both eval=True and eval=False paths are covered with appropriate assertions.

---

121-140: LGTM!

Well-documented test with manually computed expected values. The use of pytest.approx for floating-point comparisons is appropriate.

---

143-181: LGTM!

Good fixture design and comprehensive parametrized testing for tree-based classifiers with proper mocking to verify fit/predict calls.

_{✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.}

[coderabbitai]

⚠️ Potential issue | 🟡 Minor

Remove auto-generated training artifact from version control.

The catboost_info/ directory is auto-generated by CatBoost during training runs and should not be committed to the repository. This is similar to other build artifacts like __pycache__ or .pytest_cache.

Add catboost_info/ to .gitignore:

+# CatBoost training artifacts
+catboost_info/

Then remove the committed file:

git rm -r catboost_info/

🤖 Prompt for AI Agents

In `@catboost_info/catboost_training.json` around lines 1 - 1004, This commit
includes an auto-generated CatBoost training artifact (the catboost_info
directory); add "catboost_info/" to .gitignore, run git rm -r --cached
catboost_info/ (or git rm -r catboost_info/) to remove the tracked artifact from
the index, and commit the .gitignore update and removal so future CatBoost runs
do not reintroduce the file.

[coderabbitai]

⚠️ Potential issue | 🟠 Major

Off-by-one error: file at index 25 is skipped.

The train/test split skips the file at index 25:

• train_feature_files = sorted_feature_files[:25] → indices 0-24 (25 files)
• test_feature_files = sorted_feature_files[26:] → indices 26+ (skips index 25)

If all files should be used, this should be sorted_feature_files[25:].

🐛 Proposed fix

         # Get the first 25 feature files
         train_feature_files = sorted_feature_files[:25]
         # 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[26:]
+        test_feature_files = sorted_feature_files[25:]
         df_test_features = pd.concat([pd.read_csv(f) for f in test_feature_files], ignore_index=True)

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	train_feature_files = sorted_feature_files[:25]
	# 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[26:]
	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"])
	train_feature_files = sorted_feature_files[:25]
	# 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[25:]
	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"])

🤖 Prompt for AI Agents

In `@examples/ept_attack/run_ept_attack.py` around lines 159 - 168, The split
currently skips the file at index 25 because train_feature_files =
sorted_feature_files[:25] uses indices 0–24 while test_feature_files =
sorted_feature_files[26:] starts at 26; update test_feature_files to use
sorted_feature_files[25:] (referencing train_feature_files, test_feature_files,
and sorted_feature_files) so the file at index 25 is included in the test set
and no file is omitted.

[coderabbitai]

⚠️ Potential issue | 🟡 Minor

Typo in variable names: "trianing" should be "training".

Minor typo in multiple variable names that affects readability.

✏️ Proposed fix

-                    trianing_directory_name = f"{classifier}_" + "_".join(selected_columns_tuple)
-                    trianing_output_path = output_summary_path / trianing_directory_name
-                    trianing_output_path.mkdir(parents=True, exist_ok=True)
+                    training_directory_name = f"{classifier}_" + "_".join(selected_columns_tuple)
+                    training_output_path = output_summary_path / training_directory_name
+                    training_output_path.mkdir(parents=True, exist_ok=True)

Also update the references on lines 219 and 225:

-                        file_path=trianing_output_path,
+                        file_path=training_output_path,
...
-                    with open(trianing_output_path / scores_file_name, "w") as f:
+                    with open(training_output_path / scores_file_name, "w") as f:

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	trianing_directory_name = f"{classifier}_" + "_".join(selected_columns_tuple)
	trianing_output_path = output_summary_path / trianing_directory_name
	trianing_output_path.mkdir(parents=True, exist_ok=True)
	training_directory_name = f"{classifier}_" + "_".join(selected_columns_tuple)
	training_output_path = output_summary_path / training_directory_name
	training_output_path.mkdir(parents=True, exist_ok=True)

🤖 Prompt for AI Agents

In `@examples/ept_attack/run_ept_attack.py` around lines 210 - 212, Rename the
misspelled variables trianing_directory_name and trianing_output_path to
training_directory_name and training_output_path respectively, and update all
subsequent references in the same scope/function (e.g., the later uses around
where those variables are referenced) to use the corrected names so the
directory creation and downstream code use training_directory_name and
training_output_path consistently.

[emersodb]

Overall the PR looks quite good. Mostly minor comments or suggestions.

[emersodb]

Naive question: Are we distinguishing between training and non-training data or synthetic data here?

[emersodb]

I think this will be fairly unique. You could also just use hashes of some length. Eithter way is fine.

[emersodb]

I think you can be a touch more specific on the typing of this dictionary?

[emersodb]

You do this checking process at least once in another file. Maybe put this into a general utility in the library, either under a new utils folder or under common?

[emersodb]

Perhaps a silly question, but why 25 and 25, i.e. an even split here? Are we also expecting there always to be exactly 50 of these files? That feels tied to the number of shadows that have been trained somehow? Perhaps instead you can do a programmatic split (an even split if that's what you think makes sense?)

[emersodb]

Based on the implementation of filter_data these are actually suffixes of columns I think? I'd suggest updating the doc description here to reflect that and also consider renaming this variable to column_suffixes?

[emersodb]

We actually have a special variable defined for this in the library

from midst_toolkit.common.variables import DEVICE

and just use the variable DEVICE

[emersodb]

Rather than doing this assert, I'd say just fold this into the if-else block, where, at the end, if you get there you throw a value error?

[emersodb]

Super minor, but maybe do this assert before you construct your results.

[emersodb]

You can use our DEVICE variable here too 🙂