From 2281041f7223bae2b422f515f642d7ca8d7b4285 Mon Sep 17 00:00:00 2001 From: ChrisW09 <50968720+ChrisW09@users.noreply.github.com> Date: Mon, 27 Jul 2026 22:34:41 +0200 Subject: [PATCH 1/2] fix(models): make encode() deterministic and repair the LSS override encode() never switched the model out of training mode, so embeddings taken after fit() ran with dropout active: two calls on the same rows differed by up to 4.45 and no result was reproducible. It also ran with autograd enabled. The model is now put in eval() under torch.no_grad() and its previous mode restored afterwards. SklearnBaseLSS.encode() additionally unpacked two values from the three-element (num_features, cat_features, embeddings) batches that preprocess_new_data yields, so it raised 'ValueError: too many values to unpack (expected 2)' for every LSS model. It now passes the whole batch to estimator.encode(), matching the base _PredictMixin implementation. Fixes #451 Co-Authored-By: Claude Fable 5 --- deeptab/models/_mixins/predict.py | 17 +++++++-- deeptab/models/lss_base.py | 17 +++++++-- tests/test_encode.py | 61 +++++++++++++++++++++++++++++++ 3 files changed, 87 insertions(+), 8 deletions(-) create mode 100644 tests/test_encode.py diff --git a/deeptab/models/_mixins/predict.py b/deeptab/models/_mixins/predict.py index 464f353f..80d043e1 100644 --- a/deeptab/models/_mixins/predict.py +++ b/deeptab/models/_mixins/predict.py @@ -162,9 +162,18 @@ def encode(self, X, embeddings=None, batch_size=64): encoded_dataset = self._data_module.preprocess_new_data(X, embeddings) data_loader = DataLoader(encoded_dataset, batch_size=batch_size, shuffle=False) - encoded_outputs = [] - for batch in tqdm(data_loader): - emb = self._task_model.estimator.encode(batch) # type: ignore[union-attr] - encoded_outputs.append(emb) + # Embeddings must be deterministic: without eval() dropout stays active + # after fit() and two calls on the same rows disagree. + was_training = self._task_model.training + self._task_model.eval() + try: + encoded_outputs = [] + with torch.no_grad(): + for batch in tqdm(data_loader): + emb = self._task_model.estimator.encode(batch) # type: ignore[union-attr] + encoded_outputs.append(emb) + finally: + if was_training: + self._task_model.train() return torch.cat(encoded_outputs, dim=0) diff --git a/deeptab/models/lss_base.py b/deeptab/models/lss_base.py index 16a8f581..99256577 100644 --- a/deeptab/models/lss_base.py +++ b/deeptab/models/lss_base.py @@ -505,11 +505,20 @@ def encode(self, X, batch_size=64): data_loader = DataLoader(encoded_dataset, batch_size=batch_size, shuffle=False) - # Process data in batches + # Process data in batches. preprocess_new_data yields + # (num_features, cat_features, embeddings) triples, and encode() takes the + # whole batch -- matching the base _PredictMixin.encode implementation. + was_training = self._task_model.training # type: ignore[union-attr] + self._task_model.eval() # type: ignore[union-attr] encoded_outputs = [] - for num_features, cat_features in tqdm(data_loader): - embeddings = self._task_model.estimator.encode(num_features, cat_features) # type: ignore[union-attr] # Call your encode function - encoded_outputs.append(embeddings) + try: + with torch.no_grad(): + for batch in tqdm(data_loader): + embeddings = self._task_model.estimator.encode(batch) # type: ignore[union-attr] + encoded_outputs.append(embeddings) + finally: + if was_training: + self._task_model.train() # type: ignore[union-attr] # Concatenate all encoded outputs encoded_outputs = torch.cat(encoded_outputs, dim=0) diff --git a/tests/test_encode.py b/tests/test_encode.py new file mode 100644 index 00000000..57d13c18 --- /dev/null +++ b/tests/test_encode.py @@ -0,0 +1,61 @@ +"""Regression tests for the public encode() API. + +encode() never switched the model to eval mode, so embeddings taken after fit() +were dropout-corrupted and differed between calls. The LSS override additionally +unpacked 2 values from the 3-element batches preprocess_new_data yields, so it +always raised ValueError. +""" + +import numpy as np +import pandas as pd + +from deeptab.configs import FTTransformerConfig +from deeptab.models import FTTransformerLSS, FTTransformerRegressor + +QUIET = { + "accelerator": "cpu", + "devices": 1, + "enable_progress_bar": False, + "enable_model_summary": False, + "logger": False, +} + + +def _data(n=48, seed=0): + rng = np.random.default_rng(seed) + X = pd.DataFrame({"a": rng.normal(size=n), "b": rng.normal(size=n), "g": rng.choice(list("xyz"), n)}) + return X, rng.normal(size=n) + + +def _config(dropout=0.5): + return FTTransformerConfig(d_model=16, n_layers=1, n_heads=2, attn_dropout=dropout) + + +def test_encode_is_deterministic_after_fit(): + """Dropout must be disabled: two calls on the same rows must agree.""" + X, y = _data() + model = FTTransformerRegressor(model_config=_config()) + model.fit(X, y, max_epochs=1, batch_size=16, **QUIET) + + first = np.asarray(model.encode(X)) + second = np.asarray(model.encode(X)) + np.testing.assert_allclose(first, second, rtol=1e-5, atol=1e-6) + + +def test_encode_restores_training_mode(): + X, y = _data() + model = FTTransformerRegressor(model_config=_config()) + model.fit(X, y, max_epochs=1, batch_size=16, **QUIET) + + model._task_model.train() # type: ignore[union-attr] + model.encode(X) + assert model._task_model.training is True # type: ignore[union-attr] + + +def test_lss_encode_returns_embeddings(): + X, y = _data() + model = FTTransformerLSS(model_config=_config(dropout=0.0)) + model.fit(X, y, family="normal", max_epochs=1, batch_size=16, **QUIET) + + encoded = np.asarray(model.encode(X)) + assert encoded.shape[0] == len(X) From 8b589914c1e280ca7064d525f207f24c785559df Mon Sep 17 00:00:00 2001 From: ChrisW09 <50968720+ChrisW09@users.noreply.github.com> Date: Mon, 27 Jul 2026 22:35:24 +0200 Subject: [PATCH 2/2] test: silence pyright on the ITaskModel training-mode access Co-Authored-By: Claude Fable 5 --- deeptab/models/_mixins/predict.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deeptab/models/_mixins/predict.py b/deeptab/models/_mixins/predict.py index 80d043e1..08fbe809 100644 --- a/deeptab/models/_mixins/predict.py +++ b/deeptab/models/_mixins/predict.py @@ -164,7 +164,7 @@ def encode(self, X, embeddings=None, batch_size=64): # Embeddings must be deterministic: without eval() dropout stays active # after fit() and two calls on the same rows disagree. - was_training = self._task_model.training + was_training = self._task_model.training # type: ignore[attr-defined] self._task_model.eval() try: encoded_outputs = []