diff --git a/CHANGELOG.md b/CHANGELOG.md index ca8ca44..c3165de 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,18 @@ To release a new version (e.g. from `1.0.0` -> `2.0.0`): --> +## [Unreleased] + +### Added + +* `TabFMDataGenerator`: synthetic tabular data generation. Samples new rows + mimicking a reference dataset via chain-rule factorization over columns, + with `TabFMClassifier` providing each conditional distribution. Numerical + columns use hierarchical quantile refinement (the fine-bin index is sampled + digit by digit, `n_bins ** n_levels` effective bins — 100 by default); + sampling is temperature-controlled with seeded or user-supplied column + orders. + ## [1.0.1] - 2026-07-09 ### Fixed diff --git a/README.md b/README.md index 0fcffd5..e7dc343 100644 --- a/README.md +++ b/README.md @@ -134,13 +134,56 @@ predictions = reg.predict(X_test) print("Predicted Prices:", predictions) ``` +### 3. Synthetic Data Generation + +TabFM can also act as a generative model: `TabFMDataGenerator` samples +entirely new rows that mimic a reference dataset. It factorizes the joint +distribution over columns with the chain rule and samples each column from a +TabFM classifier's predictive distribution. Numerical columns are sampled at +fine resolution via hierarchical quantile refinement (`n_bins ** n_levels` +equal-mass bins — 100 by default, 1000 with `n_levels=3`), with the value +drawn uniformly within the sampled bin. + +```python +import pandas as pd +from tabfm import TabFMDataGenerator + +# Choose your backend: + +# OPTION A: JAX Backend +from tabfm import tabfm_v1_0_0_jax as tabfm_v1_0_0 + +# OPTION B: PyTorch Backend +# from tabfm import tabfm_v1_0_0_pytorch as tabfm_v1_0_0 + +# Generation only needs the classification model. +model = tabfm_v1_0_0.load() + +X = pd.DataFrame({ + "age": [25.0, 45.0, 35.0, 50.0, 28.0, 41.0], + "job": ["engineer", "manager", "engineer", "manager", "analyst", "analyst"], + "income": [80000, 120000, 90000, 130000, 70000, 100000], +}) + +gen = TabFMDataGenerator(model=model, random_state=42) +gen.fit(X, categorical_features=["job"]) +synthetic = gen.sample(n_samples=100, t=1.0) # t < 1 => closer to the modes +print(synthetic.head()) +``` + +Categorical columns support up to the model's `max_classes` distinct values +directly (rarer values are merged and re-sampled empirically). Generated rows +mimic the reference distribution but carry no formal privacy guarantee. + --- ## Examples Directory -You can find runnable scripts for both classification and regression under the [examples/](examples/) folder: +You can find runnable scripts for classification, regression and synthetic +data generation under the [examples/](examples/) folder: * [classification_example.py](examples/classification_example.py) * [regression_example.py](examples/regression_example.py) +* [synthetic_data_example.py](examples/synthetic_data_example.py) To run them, simply execute: ```bash diff --git a/conftest.py b/conftest.py index a651128..0ae794a 100644 --- a/conftest.py +++ b/conftest.py @@ -36,6 +36,7 @@ collect_ignore = [] if not _has_torch: collect_ignore.append("tabfm/src/classifier_and_regressor_pytorch_test.py") + collect_ignore.append("tabfm/src/generation_pytorch_test.py") if not _has_jax: collect_ignore += [ "tabfm/src/jax/model_test.py", diff --git a/examples/synthetic_data_example.py b/examples/synthetic_data_example.py new file mode 100644 index 0000000..1bd3c86 --- /dev/null +++ b/examples/synthetic_data_example.py @@ -0,0 +1,57 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Example showing how to generate synthetic data with TabFM v1.0.0.""" + +import numpy as np +import pandas as pd +import tabfm + + +def run_example(model=None) -> pd.DataFrame: + """Generates synthetic rows mimicking a small mixed-type dataset.""" + if model is None: + # Option A: JAX Backend (default) + model = tabfm.tabfm_v1_0_0_jax.load(model_type="classification") + + # Option B: PyTorch Backend + # model = tabfm.tabfm_v1_0_0_pytorch.load(model_type="classification") + + # 2. Build a reference dataset with correlated mixed-type columns. + rng = np.random.default_rng(0) + n = 200 + age = rng.uniform(20, 65, n) + job = rng.choice(["engineer", "manager", "analyst"], n) + income = (40000 + age * 1500 + (job == "manager") * 30000 + + rng.normal(0, 5000, n)) + X = pd.DataFrame({"age": age, "job": job, "income": income}) + + # 3. Fit the generator and sample new rows. Generation only needs the + # classification model; t < 1 concentrates near the modes of the data. + gen = tabfm.TabFMDataGenerator(model=model, random_state=42) + gen.fit(X, categorical_features=["job"]) + synthetic = gen.sample(n_samples=100, t=1.0) + + # 4. Compare real and synthetic statistics. + print("Real income by job:\n", X.groupby("job")["income"].mean()) + print("Synthetic income by job:\n", + synthetic.groupby("job")["income"].mean()) + return synthetic + + +if __name__ == "__main__": + print("Running TabFM synthetic data generation... (Note: compilation and " + "model execution may take a few minutes on first run)") + data = run_example() + print("Synthetic sample:\n", data.head()) diff --git a/tabfm/__init__.py b/tabfm/__init__.py index a0844f1..c969915 100644 --- a/tabfm/__init__.py +++ b/tabfm/__init__.py @@ -28,6 +28,7 @@ pass from tabfm.src.classifier_and_regressor import TabFMClassifier, TabFMRegressor +from tabfm.src.generation import TabFMDataGenerator # A new PyPI release will be pushed every time `__version__` is increased. # When changing this, also update the CHANGELOG.md. diff --git a/tabfm/src/BUILD b/tabfm/src/BUILD index 45a0a30..4f0a8ca 100644 --- a/tabfm/src/BUILD +++ b/tabfm/src/BUILD @@ -44,6 +44,32 @@ py_test( ], ) +py_library( + name = "generation", + srcs = ["generation.py"], + visibility = ["//tabfm:__subpackages__"], + deps = [ + ":classifier_and_regressor", + "@pip//jaxtyping", + "@pip//numpy", + "@pip//pandas", + ], +) + +py_test( + name = "generation_test", + srcs = ["generation_test.py"], + deps = [ + ":generation", + "//tabfm/src/jax:model", + "@pip//absl_py", + "@pip//flax", + "@pip//jax", + "@pip//numpy", + "@pip//pandas", + ], +) + py_library( name = "torch_convert", srcs = ["hugging_face/torch_convert.py"], diff --git a/tabfm/src/generation.py b/tabfm/src/generation.py new file mode 100644 index 0000000..b8b053a --- /dev/null +++ b/tabfm/src/generation.py @@ -0,0 +1,425 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Synthetic tabular data generation on top of TabFM. + +Generates new rows that mimic a reference dataset by factorizing the joint +distribution over columns with the chain rule, +p(x_1, ..., x_d) = prod_j p(x_j | x_{ jt.Float[np.ndarray, "N K"]: + """Sharpens (t < 1) or flattens (t > 1) row-wise probabilities.""" + if t == 1.0: + return probs + with np.errstate(divide="ignore"): + logits = np.log(probs) / t # zeros -> -inf, stay zero after re-softmax + logits -= logits.max(axis=-1, keepdims=True) + scaled = np.exp(logits) + return scaled / scaled.sum(axis=-1, keepdims=True) + + +@jt.typed +def sample_rows( + probs: jt.Float[np.ndarray, "N K"], rng: np.random.Generator +) -> jt.Int[np.ndarray, "N"]: + """Draws one class index per row from row-wise probabilities.""" + cdf = np.cumsum(probs, axis=-1) + cdf /= cdf[:, -1:] # guard against floating-point drift + u = rng.random((probs.shape[0], 1)) + return (u > cdf).sum(axis=-1) + + +@jt.typed +def quantile_edges( + values: jt.Float[np.ndarray, "N"], n_bins: int +) -> jt.Float[np.ndarray, "E"]: + """Equal-mass bin edges; duplicate quantiles collapse for skewed data.""" + qs = np.linspace(0.0, 1.0, n_bins + 1) + return np.unique(np.quantile(values, qs)) + + +@jt.typed +def bin_values( + values: jt.Float[np.ndarray, "N"], edges: jt.Float[np.ndarray, "E"] +) -> jt.Int[np.ndarray, "N"]: + """Assigns each value the id of its bin, in [0, len(edges) - 2].""" + ids = np.searchsorted(edges, values, side="right") - 1 + return np.clip(ids, 0, len(edges) - 2) + + +@jt.typed +def sample_within_bins( + bin_ids: jt.Int[np.ndarray, "N"], + edges: jt.Float[np.ndarray, "E"], + rng: np.random.Generator, +) -> jt.Float[np.ndarray, "N"]: + """Draws uniformly inside each row's bin.""" + lo = edges[bin_ids] + hi = edges[bin_ids + 1] + return lo + rng.random(len(bin_ids)) * (hi - lo) + + +@jt.typed +def digits_of( + indices: jt.Int[np.ndarray, "N"], base: int, n_digits: int +) -> jt.Int[np.ndarray, "N D"]: + """Decomposes indices into base-`base` digits, most significant first.""" + out = np.empty((len(indices), n_digits), dtype=np.int64) + rest = indices.astype(np.int64) + for d in range(n_digits - 1, -1, -1): + out[:, d] = rest % base + rest = rest // base + return out + + +@jt.typed +def indices_from_digits( + digits: jt.Int[np.ndarray, "N D"], base: int +) -> jt.Int[np.ndarray, "N"]: + """Recomposes indices from base-`base` digits (inverse of digits_of).""" + out = np.zeros(len(digits), dtype=np.int64) + for d in range(digits.shape[1]): + out = out * base + digits[:, d] + return out + + +# --------------------------------------------------------------------------- +# Generator +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class ColumnSpec: + """Per-column sampling strategy decided at fit() time.""" + + name: str + kind: str # "constant" | "categorical" | "numeric" + dtype: Any + values: np.ndarray # NaN-free reference values for this column + categories: Optional[np.ndarray] = None # categorical: modeled classes + edges: Optional[np.ndarray] = None # numeric: fine quantile bin edges + base: Optional[int] = None # numeric: classes per refinement level + n_digits: Optional[int] = None # numeric: refinement levels + other: Optional[np.ndarray] = None # categorical: merged tail classes + other_freqs: Optional[np.ndarray] = None # empirical freqs of the tail + + +class TabFMDataGenerator: + """Samples synthetic rows that mimic a reference DataFrame using TabFM. + + Attributes: + X_: Reference DataFrame stored at fit() time. + feature_names_: Column names of the reference data, in original order. + columns_: Fitted per-column ``ColumnSpec`` strategies. + """ + + X_: pd.DataFrame + feature_names_: List[str] + columns_: List[ColumnSpec] + + def __init__( + self, + model: Any, + n_bins: int = 10, + n_levels: int = 2, + n_estimators: int = 4, + random_state: Optional[int] = None, + ): + """Initialises the generator. + + Args: + model: Pre-trained TabFM classification model (NNX or PyTorch module). + n_bins: Classes per refinement level for numerical columns, capped at + the model's ``max_classes``. + n_levels: Refinement levels; numerical columns are sampled over up to + ``n_bins ** n_levels`` equal-mass quantile bins. + n_estimators: Ensemble members for each internal ``TabFMClassifier``. + random_state: Seed for column-order and sampling randomness. + """ + self.model = model + self.n_bins = n_bins + self.n_levels = n_levels + self.n_estimators = n_estimators + self.random_state = random_state + + def fit( + self, X: Any, categorical_features: Optional[List[str]] = None + ) -> "TabFMDataGenerator": + """Stores the reference data and decides each column's strategy. + + No model call happens here; like ``TabFMClassifier.fit``, this only + prepares metadata. A column is treated as categorical if its dtype is + object/category/bool, it is listed in ``categorical_features``, or it has + at most ``min(n_bins, model.max_classes)`` unique values. + + Args: + X: Reference data of shape (n_samples, n_features). + categorical_features: Optional names of columns to force categorical. + + Returns: + self. + """ + X = pd.DataFrame(X).reset_index(drop=True) + if X.shape[1] == 0 or X.shape[0] < 2: + raise ValueError("fit() needs at least 2 rows and 1 column.") + max_classes = int(getattr(self.model, "max_classes", 10)) + n_bins = min(self.n_bins, max_classes) + categorical_features = set(categorical_features or []) + + self.X_ = X + self.feature_names_ = list(X.columns) + self.columns_ = [] + for name in self.feature_names_: + values = X[name].dropna().to_numpy() + uniques, counts = np.unique(values, return_counts=True) + dtype = X[name].dtype + # Everything non-numeric (object, string, category, ...) plus bool is + # sampled as categorical; numeric columns may still fold into the + # categorical path below when their cardinality is low enough. + is_cat_dtype = not pd.api.types.is_numeric_dtype(dtype) or dtype == bool + if len(uniques) <= 1: + spec = ColumnSpec(name=name, kind="constant", dtype=dtype, + values=values) + elif (is_cat_dtype or name in categorical_features + or len(uniques) <= n_bins): + if len(uniques) <= max_classes: + spec = ColumnSpec(name=name, kind="categorical", dtype=dtype, + values=values, categories=uniques) + else: + # More classes than the model supports: model the most frequent + # ones directly and merge the tail into one class that is + # re-sampled from its empirical frequencies when drawn. + order = np.argsort(counts)[::-1] + top, tail = order[: max_classes - 1], order[max_classes - 1 :] + tail_counts = counts[tail].astype(float) + spec = ColumnSpec(name=name, kind="categorical", dtype=dtype, + values=values, categories=uniques[top], + other=uniques[tail], + other_freqs=tail_counts / tail_counts.sum()) + else: + edges = quantile_edges(values.astype(float), n_bins ** self.n_levels) + n_fine = len(edges) - 1 + # Digits needed to index the fine bins in base n_bins; duplicate + # quantiles on skewed data may shrink n_fine below the full power. + n_digits = max(1, int(np.ceil(np.log(n_fine) / np.log(n_bins)))) + spec = ColumnSpec(name=name, kind="numeric", dtype=dtype, + values=values, edges=edges, base=n_bins, + n_digits=n_digits) + self.columns_.append(spec) + return self + + def sample( + self, + n_samples: int, + t: float = 1.0, + column_order: Optional[List[str]] = None, + ) -> pd.DataFrame: + """Generates new synthetic rows mimicking the reference data. + + Columns are visited in ``column_order`` (or a seeded random permutation) + and sampled via the chain rule: the first non-constant column from its + empirical marginal, each later column from a TabFM classifier conditioned + on the columns already sampled. Constant columns never join the + conditioning set. + + Args: + n_samples: Number of synthetic rows to generate. + t: Sampling temperature; < 1 concentrates near the modes of the + reference data, > 1 flattens the sampled distributions. + column_order: Optional explicit visitation order; must be a permutation + of the fitted column names. + + Returns: + DataFrame of shape (n_samples, n_features) in the original column + order, with dtypes matching the reference data. + """ + if not hasattr(self, "columns_"): + raise ValueError( + "This TabFMDataGenerator is not fitted yet; call fit(X) first." + ) + if column_order is not None and sorted(column_order) != sorted( + self.feature_names_ + ): + raise ValueError( + "column_order must be a permutation of the fitted columns " + f"{self.feature_names_}, got {column_order}." + ) + rng = np.random.default_rng(self.random_state) + specs = {s.name: s for s in self.columns_} + if column_order is not None: + order = list(column_order) + else: + order = [self.feature_names_[i] + for i in rng.permutation(len(self.feature_names_))] + + synth = pd.DataFrame(index=range(n_samples)) + conditioning = [] + for name in order: + spec = specs[name] + if spec.kind == "constant" or not conditioning: + col = self._sample_marginal(spec, n_samples, t, rng) + else: + col = self._sample_conditional(spec, synth[conditioning], t, rng) + if spec.kind == "numeric" and pd.api.types.is_integer_dtype(spec.dtype): + col = np.round(col) + synth[name] = col + if spec.kind != "constant": + conditioning.append(name) + # Restore the reference dtypes (categoricals were sampled as object, + # integer-dtype numerics as rounded floats). + dtypes = {s.name: s.dtype for s in self.columns_} + return synth[self.feature_names_].astype(dtypes) + + def _encode_target(self, spec: ColumnSpec, values: np.ndarray) -> np.ndarray: + """Maps column values to integer class codes for classifier fitting.""" + code_map = {v: i for i, v in enumerate(spec.categories)} + codes = pd.Series(values).map(code_map) + # Values outside `categories` are the merged tail; they share one code. + return codes.fillna(len(spec.categories)).to_numpy(dtype=np.int64) + + def _decode_categorical( + self, spec: ColumnSpec, codes: np.ndarray, rng: np.random.Generator + ) -> np.ndarray: + """Maps sampled codes back to values; the tail code draws empirically.""" + out = np.empty(len(codes), dtype=object) + top = codes < len(spec.categories) + out[top] = spec.categories[codes[top]] + n_other = int((~top).sum()) + if n_other: + out[~top] = rng.choice(spec.other, size=n_other, p=spec.other_freqs) + return out + + def _sample_marginal( + self, spec: ColumnSpec, n_samples: int, t: float, + rng: np.random.Generator + ) -> np.ndarray: + """Model-free draw from the column's empirical marginal distribution.""" + if spec.kind == "constant": + return np.full(n_samples, spec.values[0]) + if spec.kind == "categorical": + codes = self._encode_target(spec, spec.values) + n_codes = len(spec.categories) + (1 if spec.other is not None else 0) + counts = np.bincount(codes, minlength=n_codes).astype(float) + probs = temperature_scale((counts / counts.sum())[None, :], t)[0] + drawn = rng.choice(len(probs), size=n_samples, p=probs) + return self._decode_categorical(spec, drawn, rng) + ids = bin_values(spec.values.astype(float), spec.edges) + counts = np.bincount(ids, minlength=len(spec.edges) - 1).astype(float) + probs = temperature_scale((counts / counts.sum())[None, :], t)[0] + drawn = rng.choice(len(probs), size=n_samples, p=probs) + return sample_within_bins(drawn, spec.edges, rng) + + def _classify_and_sample( + self, + x_ref: pd.DataFrame, + y_train: np.ndarray, + x_synth: pd.DataFrame, + t: float, + rng: np.random.Generator, + ) -> np.ndarray: + """Fits a fresh TabFMClassifier and samples one class per synthetic row.""" + clf = TabFMClassifier( + model=self.model, + n_estimators=self.n_estimators, + random_state=int(rng.integers(2**31 - 1)), + ) + clf.fit(x_ref, y_train) + probs = np.asarray(clf.predict_proba(x_synth)).astype(np.float64) + probs = probs[:, : len(clf.classes_)] + probs = probs / probs.sum(axis=-1, keepdims=True) + probs = temperature_scale(probs, t) + return np.asarray(clf.classes_)[sample_rows(probs, rng)] + + def _sample_conditional( + self, + spec: ColumnSpec, + x_cond_synth: pd.DataFrame, + t: float, + rng: np.random.Generator, + ) -> np.ndarray: + """Samples one column for all synthetic rows given the prior columns. + + Categorical columns need a single classifier fit, with the integer class + codes as the target. Numeric columns sample their fine-bin index digit by + digit: level l predicts base-`spec.base` digit l with digits 0..l-1 + appended to the conditioning features, so every level trains on all + reference rows while resolution grows as base ** levels. + """ + mask = self.X_[spec.name].notna().to_numpy() + x_cond_ref = self.X_.loc[mask, list(x_cond_synth.columns)] + if spec.kind == "categorical": + y_train = self._encode_target(spec, spec.values) + if len(np.unique(y_train)) < 2: + return self._sample_marginal(spec, len(x_cond_synth), t, rng) + drawn = self._classify_and_sample(x_cond_ref, y_train, x_cond_synth, + t, rng) + return self._decode_categorical(spec, drawn, rng) + + fine_ids = bin_values(spec.values.astype(float), spec.edges) + ref_digits = digits_of(fine_ids, spec.base, spec.n_digits) + synth_digits = np.zeros((len(x_cond_synth), spec.n_digits), dtype=np.int64) + x_ref = x_cond_ref.reset_index(drop=True) + x_synth = x_cond_synth.reset_index(drop=True) + for level in range(spec.n_digits): + y_level = ref_digits[:, level] + if len(np.unique(y_level)) < 2: + synth_digits[:, level] = y_level[0] + else: + drawn = self._classify_and_sample(x_ref, y_level, x_synth, t, rng) + # classes_ round-trips through the label encoder, which may widen + # the integer digits to float; restore ints. + synth_digits[:, level] = drawn.astype(np.int64) + digit_col = f"_digit_{level}" + while digit_col in x_ref.columns: # avoid clobbering a real column + digit_col += "_" + x_ref = x_ref.assign(**{digit_col: ref_digits[:, level]}) + x_synth = x_synth.assign(**{digit_col: synth_digits[:, level]}) + drawn_ids = indices_from_digits(synth_digits, spec.base) + # Digit combinations past the last bin can only arise when duplicate + # quantiles shrank the bin count below a full power of base; clamp them. + drawn_ids = np.minimum(drawn_ids, len(spec.edges) - 2) + return sample_within_bins(drawn_ids, spec.edges, rng) diff --git a/tabfm/src/generation_pytorch_test.py b/tabfm/src/generation_pytorch_test.py new file mode 100644 index 0000000..d50825f --- /dev/null +++ b/tabfm/src/generation_pytorch_test.py @@ -0,0 +1,107 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import numpy as np +import pandas as pd + +from tabfm.src.generation import TabFMDataGenerator +from tabfm.src.pytorch import model as pytorch_model + + +def _tiny_model(): + return pytorch_model.TabFM( + embed_dim=8, + max_classes=3, + col_num_blocks=1, + col_nhead=2, + col_num_inds=8, + row_num_blocks=1, + row_nhead=2, + row_num_cls=2, + icl_num_blocks=1, + icl_nhead=2, + ff_factor=2, + feature_group_size=2, + is_classifier=True, + ) + + +def _reference_frame(): + rng = np.random.default_rng(0) + n = 24 + return pd.DataFrame({ + "num": rng.normal(size=n), + "cat": np.where(rng.random(n) < 0.5, "a", "b"), + "const": np.full(n, 7.0), + }) + + +class TabFMDataGeneratorEndToEndTest(unittest.TestCase): + + def test_sample_shapes_dtypes_and_domains(self): + df = _reference_frame() + gen = TabFMDataGenerator(model=_tiny_model(), n_estimators=2, + random_state=42).fit(df) + out = gen.sample(n_samples=8) + self.assertEqual(list(out.columns), ["num", "cat", "const"]) + self.assertEqual(len(out), 8) + self.assertTrue(set(out["cat"]) <= {"a", "b"}) + self.assertTrue(np.all(out["const"].to_numpy() == 7.0)) + num = out["num"].to_numpy() + self.assertTrue(np.all(num >= df["num"].min()) + and np.all(num <= df["num"].max())) + # Fresh continuous values, not copies of reference rows. + self.assertFalse(bool(set(num) & set(df["num"]))) + + def test_sample_is_deterministic_given_seed(self): + # One shared model instance: _tiny_model() has random init, so two + # instances would differ regardless of the generator's seeding. + df = _reference_frame() + model = _tiny_model() + out1 = TabFMDataGenerator(model=model, n_estimators=2, + random_state=7).fit(df).sample(6) + out2 = TabFMDataGenerator(model=model, n_estimators=2, + random_state=7).fit(df).sample(6) + pd.testing.assert_frame_equal(out1, out2) + + def test_explicit_column_order_and_validation(self): + df = _reference_frame() + gen = TabFMDataGenerator(model=_tiny_model(), n_estimators=2, + random_state=0).fit(df) + out = gen.sample(4, column_order=["cat", "num", "const"]) + self.assertEqual(list(out.columns), ["num", "cat", "const"]) + with self.assertRaises(ValueError): + gen.sample(4, column_order=["cat", "num"]) # not a full permutation + + def test_integer_numeric_columns_round_trip(self): + rng = np.random.default_rng(0) + df = pd.DataFrame({ + "count": rng.integers(0, 1000, size=24), + "cat": np.where(rng.random(24) < 0.5, "a", "b"), + }) + gen = TabFMDataGenerator(model=_tiny_model(), n_estimators=2, + random_state=0).fit(df) + out = gen.sample(6) + self.assertTrue(pd.api.types.is_integer_dtype(out["count"])) + + def test_sample_before_fit_raises(self): + gen = TabFMDataGenerator(model=_tiny_model()) + with self.assertRaises(ValueError): + gen.sample(4) + + +if __name__ == "__main__": + unittest.main() diff --git a/tabfm/src/generation_test.py b/tabfm/src/generation_test.py new file mode 100644 index 0000000..5cb34da --- /dev/null +++ b/tabfm/src/generation_test.py @@ -0,0 +1,284 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +from absl.testing import absltest +import numpy as np +import pandas as pd + +try: + from flax import nnx + from tabfm.src.jax import model as tabfm_jax_model + HAS_JAX = True +except ImportError: + HAS_JAX = False +from tabfm.src.generation import bin_values +from tabfm.src.generation import digits_of +from tabfm.src.generation import indices_from_digits +from tabfm.src.generation import quantile_edges +from tabfm.src.generation import sample_rows +from tabfm.src.generation import sample_within_bins +from tabfm.src.generation import TabFMDataGenerator +from tabfm.src.generation import temperature_scale + + +class _FakeModel: + """Stands in for a TabFM model; fit() only reads max_classes.""" + + max_classes = 10 + + +class SamplingPrimitivesTest(absltest.TestCase): + + def test_temperature_scale_identity_at_t_1(self): + probs = np.array([[0.2, 0.3, 0.5]]) + np.testing.assert_allclose(temperature_scale(probs, 1.0), probs) + + def test_temperature_scale_sharpens_below_1(self): + probs = np.array([[0.4, 0.6]]) + sharp = temperature_scale(probs, 0.1) + self.assertGreater(sharp[0, 1], 0.95) + np.testing.assert_allclose(sharp.sum(axis=-1), 1.0) + + def test_temperature_scale_flattens_above_1(self): + probs = np.array([[0.1, 0.9]]) + flat = temperature_scale(probs, 10.0) + self.assertLess(flat[0, 1], 0.9) + self.assertGreater(flat[0, 1], 0.5) + np.testing.assert_allclose(flat.sum(axis=-1), 1.0) + + def test_temperature_scale_keeps_zeros_zero(self): + probs = np.array([[0.0, 1.0]]) + scaled = temperature_scale(probs, 0.5) + self.assertEqual(scaled[0, 0], 0.0) + self.assertEqual(scaled[0, 1], 1.0) + + def test_sample_rows_respects_degenerate_probs(self): + rng = np.random.default_rng(0) + probs = np.array([[1.0, 0.0, 0.0], [0.0, 0.0, 1.0]] * 5) + drawn = sample_rows(probs, rng) + np.testing.assert_array_equal(drawn, np.array([0, 2] * 5)) + + def test_sample_rows_matches_distribution(self): + rng = np.random.default_rng(0) + probs = np.tile(np.array([[0.2, 0.8]]), (20000, 1)) + drawn = sample_rows(probs, rng) + self.assertAlmostEqual(drawn.mean(), 0.8, delta=0.02) + + def test_quantile_edges_equal_mass(self): + values = np.arange(100, dtype=float) + edges = quantile_edges(values, 4) + self.assertLen(edges, 5) + self.assertEqual(edges[0], 0.0) + self.assertEqual(edges[-1], 99.0) + + def test_quantile_edges_collapses_duplicates_on_skewed_data(self): + values = np.array([0.0] * 90 + [1.0] * 10) + edges = quantile_edges(values, 10) + self.assertLen(edges, len(np.unique(edges))) + self.assertLess(len(edges), 11) + + def test_bin_values_covers_range_inclusively(self): + edges = np.array([0.0, 1.0, 2.0]) + ids = bin_values(np.array([0.0, 0.5, 1.0, 2.0]), edges) + np.testing.assert_array_equal(ids, np.array([0, 0, 1, 1])) + + def test_sample_within_bins_stays_inside(self): + rng = np.random.default_rng(0) + edges = np.array([0.0, 1.0, 10.0]) + ids = np.array([0, 1] * 50) + vals = sample_within_bins(ids, edges, rng) + self.assertTrue(np.all(vals[0::2] >= 0.0) and np.all(vals[0::2] <= 1.0)) + self.assertTrue(np.all(vals[1::2] >= 1.0) and np.all(vals[1::2] <= 10.0)) + + def test_digit_round_trip(self): + idx = np.arange(1000) + digits = digits_of(idx, base=10, n_digits=3) + self.assertEqual(digits.shape, (1000, 3)) + self.assertTrue(np.all(digits >= 0) and np.all(digits < 10)) + np.testing.assert_array_equal(indices_from_digits(digits, base=10), idx) + + def test_digits_most_significant_first(self): + digits = digits_of(np.array([472]), base=10, n_digits=3) + np.testing.assert_array_equal(digits[0], np.array([4, 7, 2])) + + +class FitColumnSpecTest(absltest.TestCase): + + def _fit(self, df, **kwargs): + gen = TabFMDataGenerator(model=_FakeModel(), random_state=0) + return gen.fit(df, **kwargs) + + def _spec(self, gen, name): + return next(s for s in gen.columns_ if s.name == name) + + def test_infers_kinds_from_dtypes(self): + df = pd.DataFrame({ + "num": np.linspace(0.0, 1.0, 40), + "cat": ["a", "b"] * 20, + "const": [7.0] * 40, + }) + gen = self._fit(df) + self.assertEqual(self._spec(gen, "num").kind, "numeric") + self.assertEqual(self._spec(gen, "cat").kind, "categorical") + self.assertEqual(self._spec(gen, "const").kind, "constant") + self.assertEqual(gen.feature_names_, ["num", "cat", "const"]) + + def test_numeric_base_capped_by_max_classes(self): + class TinyModel: + max_classes = 3 + + df = pd.DataFrame({"num": np.linspace(0.0, 1.0, 40), + "num2": np.linspace(0.0, 1.0, 40)}) + gen = TabFMDataGenerator(model=TinyModel(), random_state=0).fit(df) + spec = self._spec(gen, "num") + self.assertEqual(spec.base, 3) + self.assertLessEqual(len(spec.edges) - 1, 3**2) # n_levels=2 default + + def test_hierarchical_levels_and_single_level_fallback(self): + df = pd.DataFrame({"num": np.linspace(0.0, 1.0, 400), + "cat": ["a", "b"] * 200}) + gen = TabFMDataGenerator(model=_FakeModel(), n_levels=2, + random_state=0).fit(df) + spec = self._spec(gen, "num") + self.assertLen(spec.edges, 101) # 10^2 fine bins + self.assertEqual(spec.n_digits, 2) + gen1 = TabFMDataGenerator(model=_FakeModel(), n_levels=1, + random_state=0).fit(df) + self.assertEqual(gen1.columns_[0].n_digits, 1) + self.assertLessEqual(len(gen1.columns_[0].edges) - 1, 10) + + def test_low_cardinality_numeric_becomes_categorical(self): + df = pd.DataFrame({"few": [1.5, 2.5, 3.5] * 10, + "num": np.linspace(0.0, 1.0, 30)}) + gen = self._fit(df) + spec = self._spec(gen, "few") + self.assertEqual(spec.kind, "categorical") + np.testing.assert_array_equal(np.sort(spec.categories), + np.array([1.5, 2.5, 3.5])) + + def test_explicit_categorical_features_override(self): + df = pd.DataFrame({"code": np.arange(40) % 15, + "num": np.linspace(0.0, 1.0, 40)}) + gen = self._fit(df, categorical_features=["code"]) + self.assertEqual(self._spec(gen, "code").kind, "categorical") + + def test_high_cardinality_categorical_merges_tail(self): + labels = (["common%d" % i for i in range(9) for _ in range(10)] + + ["rare%d" % i for i in range(6)]) + df = pd.DataFrame({"cat": labels, "num": np.linspace(0, 1, len(labels))}) + gen = self._fit(df) + spec = self._spec(gen, "cat") + self.assertLen(spec.categories, 9) # max_classes - 1 top classes + self.assertLen(spec.other, 6) + np.testing.assert_allclose(spec.other_freqs.sum(), 1.0) + + def test_nan_rows_excluded_from_values(self): + df = pd.DataFrame({"num": [1.0, np.nan, 3.0, 4.0] * 10, + "cat": ["a", "b", "a", "b"] * 10}) + gen = self._fit(df) + self.assertLen(self._spec(gen, "num").values, 30) + + +class MarginalSamplingTest(absltest.TestCase): + + def _gen(self, df, **kwargs): + return TabFMDataGenerator(model=_FakeModel(), random_state=0).fit( + df, **kwargs) + + def _spec(self, gen, name): + return next(s for s in gen.columns_ if s.name == name) + + def test_constant_marginal(self): + df = pd.DataFrame({"const": [7.0] * 10, "num": np.linspace(0, 1, 10)}) + gen = self._gen(df) + rng = np.random.default_rng(0) + out = gen._sample_marginal(self._spec(gen, "const"), 5, 1.0, rng) + np.testing.assert_array_equal(out, np.full(5, 7.0)) + + def test_categorical_marginal_matches_frequencies(self): + df = pd.DataFrame({"cat": ["a"] * 80 + ["b"] * 20, + "num": np.linspace(0, 1, 100)}) + gen = self._gen(df) + rng = np.random.default_rng(0) + out = gen._sample_marginal(self._spec(gen, "cat"), 5000, 1.0, rng) + self.assertAlmostEqual((out == "a").mean(), 0.8, delta=0.03) + + def test_numeric_marginal_within_range(self): + df = pd.DataFrame({"num": np.linspace(-5.0, 5.0, 100), + "cat": ["a", "b"] * 50}) + gen = self._gen(df) + rng = np.random.default_rng(0) + out = gen._sample_marginal(self._spec(gen, "num"), 200, 1.0, rng) + self.assertTrue(np.all(out >= -5.0) and np.all(out <= 5.0)) + self.assertGreater(len(np.unique(out)), 100) # fresh values, not copies + + def test_encode_decode_round_trip_with_tail(self): + labels = (["common%d" % i for i in range(9) for _ in range(10)] + + ["rare%d" % i for i in range(6)]) + df = pd.DataFrame({"cat": labels, "num": np.linspace(0, 1, len(labels))}) + gen = self._gen(df) + spec = self._spec(gen, "cat") + codes = gen._encode_target(spec, spec.values) + self.assertEqual(codes.max(), len(spec.categories)) # tail code present + rng = np.random.default_rng(0) + decoded = gen._decode_categorical(spec, codes, rng) + top = codes < len(spec.categories) + np.testing.assert_array_equal(decoded[top], spec.values[top]) + for v in decoded[~top]: + self.assertIn(v, set(spec.other.tolist())) + + def test_marginal_temperature_zero_ish_picks_mode(self): + df = pd.DataFrame({"cat": ["a"] * 80 + ["b"] * 20, + "num": np.linspace(0, 1, 100)}) + gen = self._gen(df) + rng = np.random.default_rng(0) + out = gen._sample_marginal(self._spec(gen, "cat"), 50, 1e-6, rng) + self.assertTrue(np.all(out == "a")) + + +@unittest.skipUnless(HAS_JAX, "JAX backend not installed") +class JaxEndToEndTest(absltest.TestCase): + + def test_sample_with_jax_backend(self): + model = tabfm_jax_model.TabFM( + loss="cross_entropy", + max_classes=3, + embed_dim=8, + col_num_blocks=1, + col_nhead=2, + col_num_inds=8, + row_num_blocks=1, + row_nhead=2, + row_num_cls=1, + icl_num_blocks=1, + icl_nhead=2, + rngs=nnx.Rngs(0), + ) + rng = np.random.default_rng(0) + df = pd.DataFrame({ + "num": rng.normal(size=16), + "cat": np.where(rng.random(16) < 0.5, "a", "b"), + }) + gen = TabFMDataGenerator(model=model, n_estimators=2, + random_state=0).fit(df) + out = gen.sample(n_samples=4) + self.assertEqual(list(out.columns), ["num", "cat"]) + self.assertLen(out, 4) + self.assertContainsSubset(set(out["cat"]), {"a", "b"}) + + +if __name__ == "__main__": + absltest.main()