fix(models): make encode() deterministic and repair the LSS override

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  # type: ignore[attr-defined]
+        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)

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)

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)