fix(models): save and continue with the fitted preprocessor

deeptab/core/serialization.py

@@ -557,12 +557,19 @@ def build_save_bundle(
     else:
         task = "regression" if estimator._data_module.regression else "classification"
 
+    # set_params() may already have swapped estimator._preprocessor for an
+    # unfitted one configured for the next fit(); the artifact needs the
+    # preprocessor the data module was actually fitted with.
+    fitted_preprocessor = getattr(estimator._data_module, "preprocessor", None)
+    if fitted_preprocessor is None:
+        fitted_preprocessor = estimator._preprocessor
+
     artifact_metadata = build_artifact_metadata(
         estimator=estimator,
         model_class=type(estimator._estimator),
         config=estimator.config,
         data_module=estimator._data_module,
-        preprocessor=estimator._preprocessor,
+        preprocessor=fitted_preprocessor,
         preprocessor_kwargs=getattr(estimator, "_preprocessor_kwargs", {}),
         task=task,
         regression=estimator._data_module.regression,
@@ -589,7 +596,7 @@ def build_save_bundle(
         "trainer_config": estimator.trainer_config,
         "random_state": estimator.random_state,
         "preprocessor_kwargs": getattr(estimator, "_preprocessor_kwargs", {}),
-        "preprocessor": estimator._preprocessor,
+        "preprocessor": fitted_preprocessor,
         "feature_info": {
             "num": estimator._data_module.num_feature_info,
             "cat": estimator._data_module.cat_feature_info,

deeptab/models/_mixins/fit.py

@@ -346,8 +346,12 @@ def _prepare_continuation(
         from deeptab.core.default_factories import DefaultDataModuleFactory
 
         factory = getattr(self, "_data_module_factory", None) or DefaultDataModuleFactory()
+        # Continuation reuses the fitted representation, which lives on the
+        # previous data module; self._preprocessor may already be an unfitted
+        # replacement created by set_params() for the next full fit.
+        fitted_preprocessor = getattr(previous, "preprocessor", None)
         prepared = factory.create(
-            preprocessor=self._preprocessor,
+            preprocessor=fitted_preprocessor if fitted_preprocessor is not None else self._preprocessor,
             batch_size=batch_size,
             shuffle=shuffle,
             regression=regression,

tests/integration/test_preprocessing_changes_after_fit.py

@@ -0,0 +1,45 @@
+"""Changing preprocessing settings on a fitted estimator prepares the *next* fit().
+
+The fitted model must keep working: saving it, and continuing to train it with
+rebuild=False, both need the preprocessor the data was actually fitted with.
+"""
+
+import numpy as np
+import pytest
+
+from deeptab.configs import MLPConfig, PreprocessingConfig, TrainerConfig
+from deeptab.models import MLPLSS, MLPRegressor
+
+FIT_KWARGS = {"accelerator": "cpu", "logger": False, "enable_progress_bar": False, "enable_model_summary": False}
+
+
+def _fitted(model_class, tmp_path, **fit_options):
+    rng = np.random.default_rng(0)
+    features = rng.normal(size=(80, 3))
+    model = model_class(
+        model_config=MLPConfig(layer_sizes=[16]),
+        preprocessing_config=PreprocessingConfig(),
+        trainer_config=TrainerConfig(max_epochs=1, checkpoint_path=str(tmp_path)),
+        random_state=0,
+    )
+    model.fit(features, features[:, 0] * 2, **fit_options, **FIT_KWARGS)
+    return model, features
+
+
+@pytest.mark.parametrize("model_class, fit_options", [(MLPRegressor, {}), (MLPLSS, {"family": "normal"})])
+def test_save_after_set_params_keeps_the_fitted_preprocessor(model_class, fit_options, tmp_path):
+    model, features = _fitted(model_class, tmp_path, **fit_options)
+    expected = model.predict(features)
+    model.set_params(preprocessing_config__cat_cutoff=0.2)
+    path = str(tmp_path / "model.deeptab")
+    model.save(path)
+    np.testing.assert_allclose(model_class.load(path).predict(features), expected, rtol=1e-6)
+
+
+@pytest.mark.parametrize("model_class, fit_options", [(MLPRegressor, {}), (MLPLSS, {"family": "normal"})])
+def test_continuation_after_set_params_uses_the_fitted_preprocessor(model_class, fit_options, tmp_path):
+    model, features = _fitted(model_class, tmp_path, **fit_options)
+    model.set_params(preprocessing_config__cat_cutoff=0.2)
+    model.fit(features, features[:, 0] * 2, rebuild=False, **fit_options, **FIT_KWARGS)
+    assert model.is_fitted_
+    assert np.isfinite(model.predict(features)).all()