Skip to content

Adding TabSyn examples and ensemble attack code#143

Open
lotif wants to merge 71 commits into
mainfrom
marcelo/tabsyn-ensemble
Open

Adding TabSyn examples and ensemble attack code#143
lotif wants to merge 71 commits into
mainfrom
marcelo/tabsyn-ensemble

Conversation

@lotif

@lotif lotif commented Jun 3, 2026

Copy link
Copy Markdown
Collaborator

PR Type

Feature

Short Description

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

A few things in this PR:

  • Adding an example for TabSyn full pipeline (train, synthesize and evaluate)
  • Adding an example for Ensemble Attack using TabSyn
  • Adding EnsembleAttackTabSynModelRunner to src/midst_toolkit/attacks/ensemble/models.py so the attack knows how to run the model
  • Renaming test_attack_model.py files to stop pytest thinking they are test files
  • A few small refactorings to simplify things

FYI a lot of the code in the examples was copied from the CTGAN example, and the EnsembleAttackTabSynModelRunner was adapted from the existing TabSyn automated test.

Tests Added

No tests added.

@lotif lotif requested a review from emersodb June 9, 2026 20:43
@lotif lotif marked this pull request as ready for review June 9, 2026 20:43
@coderabbitai

coderabbitai Bot commented Jun 9, 2026

Copy link
Copy Markdown

Caution

Review failed

An error occurred during the review process. Please try again later.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch marcelo/tabsyn-ensemble

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

Comment @coderabbitai help to get the list of available commands.

* Bump actions/checkout from 6.0.2 to 6.0.3 (#145)

Bumps [actions/checkout](https://github.com/actions/checkout) from 6.0.2 to 6.0.3.
- [Release notes](https://github.com/actions/checkout/releases)
- [Changelog](https://github.com/actions/checkout/blob/main/CHANGELOG.md)
- [Commits](actions/checkout@v6.0.2...v6.0.3)

---
updated-dependencies:
- dependency-name: actions/checkout
  dependency-version: 6.0.3
  dependency-type: direct:production
  update-type: version-update:semver-patch
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Marcelo Lotif <lotif@users.noreply.github.com>

* Bump astral-sh/setup-uv from 8.1.0 to 8.2.0 (#144)

Bumps [astral-sh/setup-uv](https://github.com/astral-sh/setup-uv) from 8.1.0 to 8.2.0.
- [Release notes](https://github.com/astral-sh/setup-uv/releases)
- [Commits](astral-sh/setup-uv@v8.1.0...v8.2.0)

---
updated-dependencies:
- dependency-name: astral-sh/setup-uv
  dependency-version: 8.2.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
...

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Marcelo Lotif <lotif@users.noreply.github.com>

* Bump actions/checkout from 6.0.3 to 7.0.0

Bumps [actions/checkout](https://github.com/actions/checkout) from 6.0.3 to 7.0.0.
- [Release notes](https://github.com/actions/checkout/releases)
- [Changelog](https://github.com/actions/checkout/blob/main/CHANGELOG.md)
- [Commits](actions/checkout@v6.0.3...v7.0.0)

---
updated-dependencies:
- dependency-name: actions/checkout
  dependency-version: 7.0.0
  dependency-type: direct:production
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <support@github.com>

---------

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Marcelo Lotif <lotif@users.noreply.github.com>

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 16

🧹 Nitpick comments (4)
examples/tabsyn/ensemble_attack/train_attack_model.py (1)

19-28: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Docstring says "CTGAN model" but this entrypoint uses TabSyn.

Line 62 constructs EnsembleAttackTabSynModelRunner; update the docstring to reference TabSyn to avoid confusion.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@examples/tabsyn/ensemble_attack/train_attack_model.py` around lines 19 - 28,
The docstring in the attack training entrypoint is inconsistent with the actual
implementation: it says “CTGAN model” even though the script constructs
EnsembleAttackTabSynModelRunner. Update the docstring in train_attack_model to
describe the TabSyn pipeline instead, keeping the rest of the step-by-step
summary aligned with the actual runner and attack flow.
examples/tabsyn/ensemble_attack/compute_attack_success.py (1)

21-23: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Stale "CTGAN" reference in comment.

This is the TabSyn example; the comment mentions CTGAN. Minor wording fix to keep the TODO/context accurate.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@examples/tabsyn/ensemble_attack/compute_attack_success.py` around lines 21 -
23, The comment in compute_attack_success is stale because it references CTGAN
even though this TabSyn example is not about that model. Update the wording near
target_ids to describe the TabSyn-specific context accurately, keeping the TODO
focused on the placeholder target ID requirement without mentioning CTGAN.
examples/tabsyn/evaluate.py (1)

131-144: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

cfg.ks_tv.run is listed twice.

The quality any([...]) check includes cfg.ks_tv.run at both Line 133 and Line 135. The duplicate is harmless but should be removed; if one of them was meant to be a different metric flag, please correct it.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@examples/tabsyn/evaluate.py` around lines 131 - 144, The any([...]) metric
gate in evaluate() includes cfg.ks_tv.run twice, so remove the duplicate entry
and verify the remaining list in evaluate() only contains one flag per metric.
If the second occurrence was intended to represent a different check, replace it
with the correct metric symbol from the same group (for example, another
cfg.*.run field) so the set of enabled evaluations is accurate.
examples/tabsyn/synthesize.py (1)

49-86: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win

Avoid running preprocess twice.

preprocess is invoked once for categories/d_numerical and again with inverse=True for num_inverse/cat_inverse, repeating all preprocessing work. A single call with inverse=True returns all six values and halves the cost.

♻️ Suggested refactor
-    _, _, categories, d_numerical = preprocess(  # type: ignore[misc]
-        dataset_path=dataset_path,
-        ref_dataset_path=ref_dataset_path,
-        transforms=tabsyn_config["transforms"],
-        task_type=tabsyn_config["task_type"],
-    )
+    _, _, categories, d_numerical, num_inverse, cat_inverse = preprocess(  # type: ignore[misc]
+        dataset_path=dataset_path,
+        ref_dataset_path=ref_dataset_path,
+        transforms=tabsyn_config["transforms"],
+        task_type=tabsyn_config["task_type"],
+        inverse=True,
+    )

Then drop the second preprocess(..., inverse=True) call at lines 80-86.

Please confirm preprocess(..., inverse=True) returns the 6-tuple in the order (numerical, categorical, categories, d_numerical, num_inverse, cat_inverse).

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@examples/tabsyn/synthesize.py` around lines 49 - 86, `preprocess(...,
inverse=True)` should be used once to avoid duplicating preprocessing in
`synthesize.py`; update the `preprocess` call so it returns the full 6-tuple and
reuse the existing values for `categories`, `d_numerical`, `num_inverse`, and
`cat_inverse`. Make sure the return order from `preprocess` is confirmed as
`(numerical, categorical, categories, d_numerical, num_inverse, cat_inverse)`,
then remove the second `preprocess(..., inverse=True)` invocation and wire the
unpacking in `synthesize()` accordingly.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In @.github/workflows/code_checks.yml:
- Line 32: The workflow step using actions/checkout is still pinned to a mutable
version tag and leaves persisted git credentials enabled. Update the checkout
step in code_checks.yml to use the release commit SHA for actions/checkout
instead of v7.0.0, and add persist-credentials: false since this job does not
require repository auth. Use the actions/checkout step identifier to locate the
change.

In @.github/workflows/docs.yml:
- Around line 45-48: The docs workflow still references mutable action tags and
should be hardened by pinning every action use to a full commit SHA. Update the
`actions/checkout` steps in `docs.yml` and the other workflow action references
(`astral-sh/setup-uv`, `actions/setup-python`, `actions/upload-artifact`,
`actions/download-artifact`, and `JamesIves/github-pages-deploy-action`) to
immutable SHAs, and add `persist-credentials: false` to both `actions/checkout`
invocations unless a later step needs the checkout token.

In @.github/workflows/integration_tests.yml:
- Line 44: The integration_tests workflow still references a mutable
actions/checkout tag and keeps git credentials persisted by default. Update the
checkout step in the workflow to use the release commit SHA for actions/checkout
instead of the version tag, and add persist-credentials set to false unless a
later step explicitly needs repository authentication. Use the existing checkout
step in integration_tests.yml as the target for this change.

In @.github/workflows/publish.yml:
- Line 19: The publish workflow is using a mutable actions reference and leaving
git credentials persisted. Update the checkout step in the publish job to use a
full commit SHA for actions/checkout instead of the version tag, and add
persist-credentials set to false unless this job explicitly needs repository
auth. Keep the change localized to the checkout step so the workflow remains
otherwise unchanged.

In @.github/workflows/unit_tests.yml:
- Line 44: Update the checkout step in the unit test workflow to use a pinned
full commit SHA for actions/checkout instead of the version tag, and configure
the checkout action with persist-credentials set to false since the workflow
does not need git authentication afterward. Use the existing checkout step in
unit_tests.yml as the target and make sure the action reference and its options
are adjusted together.

In `@examples/tabsyn/ensemble_attack/config.yaml`:
- Around line 32-45: The shadow_training section still uses CTGAN-specific
labels that do not apply to the TabSyn ensemble-attack runner. Remove or replace
the unused model_name field in shadow_training, and update the nearby
model_config and inline comments in this config block to use TabSyn/tabddpm
terminology so the configuration reflects what the ensemble_attack flow actually
consumes.

In `@examples/tabsyn/ensemble_attack/make_challenge_dataset.py`:
- Around line 19-24: The challenge dataset sampling in make_challenge_dataset is
non-deterministic and can fail when the untrained pool is too small. Update the
untrained_data sampling to use config.ensemble_attack.random_seed as the
random_state so runs are reproducible, and guard against len(training_data)
exceeding the available untrained rows before calling sample. Keep the fix
localized around the training_data, untrained_data, and challenge_data
construction logic.

In `@examples/tabsyn/ensemble_attack/README.md`:
- Around line 8-15: The README’s data path and config key references are
inconsistent with the actual example configuration. Update the instructions in
README so the documented folder location matches the path used by config.yaml,
and replace the mention of base_data_dir with the real data_dir attribute. Make
sure the setup steps for the ensemble_attack example point users to the same
data directory that the config expects, using the config.yaml data_dir setting
as the source of truth.
- Around line 94-95: The README’s documented output path is incorrect and does
not match the actual results location used by compute_attack_success.py. Update
the path referenced in the ensemble_attack README to use the same results_dir
configured in config.yaml and passed to experiment_directory in
compute_attack_success.py, so the saved
attack_success_for_xgb_metaclassifier_model.txt path reflects
examples/tabsyn/ensemble_attack/results instead of the old
examples/tabsyn/results location.

In `@examples/tabsyn/evaluate.py`:
- Around line 54-57: The metric logging in evaluate() is malformed because the
raw f-string makes the backslash literal, so the output shows \Score instead of
a properly separated label. Update the log message in the results loop to use
normal string escaping/formatting so the metric name and “Score” appear on
separate, readable parts, and keep the surrounding SEPARATOR/header logging
consistent.

In `@examples/tabsyn/README.md`:
- Around line 83-86: The README has a filename mismatch and a typo: it documents
the sampled training data as `{table_name}_samples.csv` and says “sampleed”, but
the actual writer in `train.py` uses `_sample_data_if_needed` to produce
`{table_name}_sampled.csv`. Update the documentation in
`examples/tabsyn/README.md` to match the real filename used by
`_sample_data_if_needed` and correct the spelling so the data path instructions
are consistent.
- Around line 68-73: The Alpha Precision example command is still using CTGAN
paths instead of TabSyn paths. Update the command in the TabSyn README so the
midsts_alpha_precision_eval invocation uses the TabSyn-specific meta info and
output directory locations, matching the existing examples/tabsyn/ structure
rather than examples/gan/. Keep the rest of the command unchanged and verify the
referenced synthetic and real data paths remain correct for the TabSyn example.

In `@examples/tabsyn/train.py`:
- Around line 169-181: The _sample_data_if_needed helper in train.py samples
data non-reproducibly and can fail when sample_size is larger than the available
rows. Update the DataFrame.sample call to use a fixed random_state for
deterministic results, and guard the n argument by clamping sample_size to the
number of rows (or otherwise handling oversized requests) before sampling. Keep
the fix localized to _sample_data_if_needed and preserve the existing behavior
of writing the sampled CSVs.

In `@src/midst_toolkit/attacks/ensemble/models.py`:
- Around line 405-411: The docstring for the training method in models.py
overstates how synthetic sample count is configured, since `_synthesize` does
not use `number_of_points_to_synthesize` and instead derives `num_samples` from
`train_z_att["num_samples"]` before calling `tabsyn.sample(...)`. Update the
implementation to pass the configured synthesis count through the
`_synthesize`/`tabsyn.sample` flow, or revise the docstring to match the actual
behavior; use the `number_of_points_to_synthesize`, `_synthesize`, and
`tabsyn.sample` symbols to locate the affected code.

In `@src/midst_toolkit/attacks/ensemble/rmia/rmia_calculation.py`:
- Around line 126-128: Keep the synthetic and input schemas aligned before
calling gower_matrix in rmia_calculation. The current drop of "_id" columns in
df_synthetic leaves df_input and categorical_features based on a different
column layout, so update the same filtering logic for the input frame or
recompute the categorical mask after both frames are aligned. Make the fix in
the code path around the df_synthetic drop and the subsequent gower.gower_matrix
call so both matrices share identical columns and mask positions.

In `@src/midst_toolkit/models/tabsyn/preprocessing.py`:
- Around line 154-156: The test-file loading path in preprocessing.py is using
the wrong reader for .xls inputs, so the branch that sets data_path/test_path
for train.xls must also make the test handling in the same preprocessing flow
respect test_path.suffix. Update the logic around the tabsyn preprocessing
loader to mirror the train-side behavior: use pd.read_excel for .xls files, keep
pd.read_csv for CSVs, and apply the same ID-column drop to the test dataframe as
already done for the train dataframe in the relevant preprocessing function.

---

Nitpick comments:
In `@examples/tabsyn/ensemble_attack/compute_attack_success.py`:
- Around line 21-23: The comment in compute_attack_success is stale because it
references CTGAN even though this TabSyn example is not about that model. Update
the wording near target_ids to describe the TabSyn-specific context accurately,
keeping the TODO focused on the placeholder target ID requirement without
mentioning CTGAN.

In `@examples/tabsyn/ensemble_attack/train_attack_model.py`:
- Around line 19-28: The docstring in the attack training entrypoint is
inconsistent with the actual implementation: it says “CTGAN model” even though
the script constructs EnsembleAttackTabSynModelRunner. Update the docstring in
train_attack_model to describe the TabSyn pipeline instead, keeping the rest of
the step-by-step summary aligned with the actual runner and attack flow.

In `@examples/tabsyn/evaluate.py`:
- Around line 131-144: The any([...]) metric gate in evaluate() includes
cfg.ks_tv.run twice, so remove the duplicate entry and verify the remaining list
in evaluate() only contains one flag per metric. If the second occurrence was
intended to represent a different check, replace it with the correct metric
symbol from the same group (for example, another cfg.*.run field) so the set of
enabled evaluations is accurate.

In `@examples/tabsyn/synthesize.py`:
- Around line 49-86: `preprocess(..., inverse=True)` should be used once to
avoid duplicating preprocessing in `synthesize.py`; update the `preprocess` call
so it returns the full 6-tuple and reuse the existing values for `categories`,
`d_numerical`, `num_inverse`, and `cat_inverse`. Make sure the return order from
`preprocess` is confirmed as `(numerical, categorical, categories, d_numerical,
num_inverse, cat_inverse)`, then remove the second `preprocess(...,
inverse=True)` invocation and wire the unpacking in `synthesize()` accordingly.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 3e600db0-a23a-4204-a9d7-db085388e76f

📥 Commits

Reviewing files that changed from the base of the PR and between 9f0e89c and 67b1f03.

📒 Files selected for processing (30)
  • .github/workflows/code_checks.yml
  • .github/workflows/docs.yml
  • .github/workflows/integration_tests.yml
  • .github/workflows/publish.yml
  • .github/workflows/unit_tests.yml
  • examples/gan/README.md
  • examples/gan/ensemble_attack/README.md
  • examples/gan/ensemble_attack/config.yaml
  • examples/gan/ensemble_attack/make_challenge_dataset.py
  • examples/gan/ensemble_attack/run_test_attack_model.py
  • examples/gan/ensemble_attack/utils.py
  • examples/tabsyn/README.md
  • examples/tabsyn/config.yaml
  • examples/tabsyn/ensemble_attack/README.md
  • examples/tabsyn/ensemble_attack/compute_attack_success.py
  • examples/tabsyn/ensemble_attack/config.yaml
  • examples/tabsyn/ensemble_attack/make_challenge_dataset.py
  • examples/tabsyn/ensemble_attack/run_test_attack_model.py
  • examples/tabsyn/ensemble_attack/train_attack_model.py
  • examples/tabsyn/evaluate.py
  • examples/tabsyn/synthesize.py
  • examples/tabsyn/tabsyn_config.toml
  • examples/tabsyn/train.py
  • src/midst_toolkit/attacks/ensemble/models.py
  • src/midst_toolkit/attacks/ensemble/rmia/rmia_calculation.py
  • src/midst_toolkit/models/tabsyn/dataset.py
  • src/midst_toolkit/models/tabsyn/pipeline.py
  • src/midst_toolkit/models/tabsyn/preprocessing.py
  • tests/integration/assets/tabsyn/trans_info.json
  • tests/integration/models/tabsyn/test_train.py
💤 Files with no reviewable changes (1)
  • examples/gan/ensemble_attack/config.yaml

Comment thread .github/workflows/code_checks.yml
Comment thread .github/workflows/docs.yml
Comment thread .github/workflows/integration_tests.yml
Comment thread .github/workflows/publish.yml
Comment thread .github/workflows/unit_tests.yml
Comment thread examples/tabsyn/README.md
Comment thread examples/tabsyn/train.py Outdated
Comment thread src/midst_toolkit/attacks/ensemble/models.py Outdated
Comment thread src/midst_toolkit/attacks/ensemble/rmia/rmia_calculation.py Outdated
Comment thread src/midst_toolkit/models/tabsyn/preprocessing.py

@emersodb emersodb left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Most of my comments are fairly small. I think the more fundamental ones that we want to fix were already highlighted by CodeRabbit, which is nice.

Comment thread .github/workflows/code_checks.yml
Comment thread .github/workflows/docs.yml
Comment thread examples/gan/ensemble_attack/utils.py Outdated
Comment thread examples/tabsyn/ensemble_attack/compute_attack_success.py Outdated
Comment thread examples/tabsyn/ensemble_attack/config.yaml Outdated
Comment thread src/midst_toolkit/attacks/ensemble/models.py Outdated

Args:
n_head: The number of heads.
factor: The factor for the dimension of the hidden layer.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think I understand what factor is here?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Will add some extra explanation to it.

Comment thread src/midst_toolkit/models/tabsyn/preprocessing.py
Comment thread src/midst_toolkit/models/tabsyn/preprocessing.py
Comment thread src/midst_toolkit/models/tabsyn/preprocessing.py
@lotif lotif requested a review from emersodb July 9, 2026 18:54

@emersodb emersodb left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Your changes in response to my comments and those of CodeRabbit look good to me.

For the integration tests, I'm going to have to trust that you've ensured that no changes of significance have happened here.

[
cfg.ks_tv.run,
cfg.ci_overlap.run,
cfg.ks_tv.run,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Where did this guy go?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is repeated a couple of lines above this one, so it's an useless line.

Comment thread examples/tabsyn/train.py
train_data = TabularDataset(numerical_features_train.float(), categorical_features_train)

# move test data to gpu if available
numerical_features_test = numerical_features_test.float().to(DEVICE)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe we just do both for consistency? The move is idempotent. So if it's already on the device it will no-op anyway.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I thought the single table integration tests were passing?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

After the upgrades, the values still match on local but not on github anymore, so I have to make new expected values for github actions :/

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants