Add TabFMDataGenerator for synthetic tabular data generation - #83
Open
LckyLke wants to merge 2 commits into
Open
Add TabFMDataGenerator for synthetic tabular data generation#83LckyLke wants to merge 2 commits into
LckyLke wants to merge 2 commits into
Conversation
LckyLke
requested review from
abhidas,
erzel,
rajatsen91,
siriuz42,
tamannarayan and
weihaokong
as code owners
July 24, 2026 13:09
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds
TabFMDataGenerator, a scikit-learn style generator that samples entirely new rows mimicking a reference dataset, using the existing pretrained classifier as the conditional-distribution engine. Closes #82.Summary
The joint distribution over columns is factorized with the chain rule,
p(x_1,...,x_d) = prod_j p(x_j | x_{<j}). Columns are visited in a seeded random (or user-supplied) order; each column is sampled for all rows at once from aTabFMClassifierfit with that column as the target and the already-sampled columns as features. This is a pure inference-time composition of the existing sklearn API — no model, architecture, or weight changes, and only the classification checkpoint is used.Numerical columns reach fine resolution despite
max_classes = 10via hierarchical quantile refinement:n_bins ** n_levelsequal-mass bins (100 by default, 1000 withn_levels=3), whose index is sampled digit by digit in basen_binswith already-sampled digits appended to the conditioning features. Every refinement level trains on all reference rows, so resolution grows exponentially while cost grows linearly. Values are drawn uniformly within the sampled bin.Design notes
sample()call: one classifier fit + onepredict_probaper categorical column,n_levelsof each per numeric column, all rows batched. With default settings (n_estimators=4),fit()makes no model forward passes (no NNLS, no calibration), so the cost is exactly thepredict_probacalls.max_classesare modeled top-k with the tail merged into one class that is re-sampled from its empirical frequencies.UniqueFeatureFilter).tsharpens (< 1) or flattens (> 1) each conditional; dtypes (including integer columns and pandas string/categorical dtypes) round-trip.Validation with the released v1.0.0 weights
Testing
tabfm/src/generation_test.py: 25 backend-free tests (absltest) plus a JAX end-to-end test guarded likeclassifier_and_regressor_test.py.tabfm/src/generation_pytorch_test.py: 5 end-to-end tests with a tiny random PyTorch model (no weight downloads); conftest-guarded like its sibling when torch is absent.tabfm/src/BUILD(the torch-only test file is deliberately not registered, matchingclassifier_and_regressor_pytorch_test.py, since torch is not in the Bazel pip lock).pytest -vv -n auto; lint profile matches the existing modules under the repo.pylintrc.examples/synthetic_data_example.py; README section and CHANGELOG entry under[Unreleased](no version bump).Known v1 limits
No extrapolation beyond the observed min/max (TabPFN adds half-normal tails — a candidate follow-up), no causal/DAG column ordering, no missingness simulation.