fix(optimization): support hold-out CV in Optuna tuning path

src/mother/ml/models/m_tabpfn.py

@@ -536,6 +536,11 @@ def _set_random_state(self) -> None:
         torch.manual_seed(self.random_state)
         np.random.seed(self.random_state)
 
+    def _prepare_prefitted_model_for_embeddings(self) -> None:
+        if self.model is None:
+            raise RuntimeError("A pre-fitted model is required to extract embeddings.")
+        self.model.use_autocast_ = False
+
     def _get_best_embeddings(
         self,
         embeddings: np.ndarray,
@@ -633,6 +638,7 @@ def fit(
             module_logger.info(
                 "A pre-fitted model has been given. The new data will not be used for fitting the model."
             )
+            self._prepare_prefitted_model_for_embeddings()
             self.train_embeddings_ = self.model.get_embeddings(X_array)
             self._embedding_dim = self.train_embeddings_.shape[1]
         else:
@@ -777,6 +783,8 @@ def transform(
             X_array = np.asarray(X, dtype=np.float32)
 
         # Get embeddings for new data using the main model
+        if self.pre_fitted:
+            self._prepare_prefitted_model_for_embeddings()
         embeddings = self.model.get_embeddings(X_array)
         # collapse the additional column caused by estimators (avg)
         if len(embeddings.shape) == 3:

src/mother/optimization/core.py

@@ -13,7 +13,11 @@
 import sklearn.metrics as skl_metrics
 import sklearn.model_selection as skl_model_sel
 from optuna.study import Study, StudyDirection
-from optuna.terminator import TerminatorCallback, report_cross_validation_scores
+from optuna.terminator import (
+    Terminator,
+    TerminatorCallback,
+    report_cross_validation_scores,
+)
 from sklearn.pipeline import Pipeline
 
 from mother import utils as mother_utils
@@ -133,7 +137,7 @@ def __init__(
         tuning_direction: typing.Union[StudyDirection, str] = StudyDirection.MAXIMIZE,
         n_trials_optuna: int = 100,
         n_threads_optuna: int = 1,
-        n_startup_trials: int = 12,
+        n_startup_trials: int = 20,
         seed: int = 42,
         **kwargs,
     ):
@@ -157,17 +161,22 @@ def __init__(
 
         self.study: typing.Optional[Study] = None
 
-    def get_callbacks(self):
+    def get_callbacks(self, cross_validation: typing.Optional[skl_model_sel.BaseCrossValidator] = None):
         """
         Prepares and returns a list of callbacks for early stopping in Optuna optimization.
 
-        If early stopping with Optuna is enabled and PyTorch is available, this method
-        will return a list containing a TerminatorCallback instance. If PyTorch is not
-        available, it will log a warning and return None.
+        If early stopping with Optuna is enabled, PyTorch is available, and the
+        cross-validation strategy supports early termination, this method will return
+        a list containing a TerminatorCallback instance.
+
+        It returns None in any of the following cases:
+        - early stopping with Optuna is disabled
+        - PyTorch is not available (warning is logged)
+        - hold-out cross-validation is detected (fewer than 2 splits)
 
         Returns:
             typing.Optional[typing.List[TerminatorCallback]]: A list of TerminatorCallback
-            instances if early stopping is enabled and PyTorch is available, otherwise None.
+            instances when early stopping can be used, otherwise None.
         """
         callbacks: typing.Optional[typing.List[TerminatorCallback]] = None
         if self.early_stopping_optuna:
@@ -179,7 +188,12 @@ def get_callbacks(self):
                     pip install mother[torch] or uv add mother[torch]"""
                 )
             else:
-                callbacks = [TerminatorCallback()]
+                if cross_validation is not None and cross_validation.get_n_splits() < 2:
+                    module_logger.warning(
+                        "Optuna early termination requires at least 2 CV splits; disabling callback for hold-out setup"
+                    )
+                    return None
+                callbacks = [TerminatorCallback(terminator=Terminator(min_n_trials=40))]
         return callbacks
 
     @handle_metadata_routing
@@ -254,7 +268,8 @@ def objective(trial: optuna.trial.Trial) -> float:
             gc.collect()
             module_logger.info(f"Trial {trial.number}, cv score: {cv_score}")
             cv_score_not_na: np.ndarray = cv_score[~np.isnan(cv_score)]
-            report_cross_validation_scores(trial, list(cv_score_not_na))
+            if len(cv_score_not_na) > 1:
+                report_cross_validation_scores(trial, list(cv_score_not_na))
             mean_cv_score: float = cv_score_not_na.mean()
             return mean_cv_score
 
@@ -284,7 +299,7 @@ def objective(trial: optuna.trial.Trial) -> float:
             objective,
             n_trials=self.n_trials_optuna,
             gc_after_trial=True,
-            callbacks=self.get_callbacks(),
+            callbacks=self.get_callbacks(cross_validation=cross_validation),
         )
 
         if default_parameters != {}:

test/unit/test_model_tuner.py

@@ -7,7 +7,7 @@
 from sklearn.datasets import make_blobs, make_classification
 from sklearn.linear_model import Lasso
 from sklearn.metrics import accuracy_score, make_scorer, mean_squared_error
-from sklearn.model_selection import GroupKFold
+from sklearn.model_selection import GroupKFold, KFold, PredefinedSplit
 from sklearn.pipeline import Pipeline
 from sklearn.preprocessing import MinMaxScaler
 
@@ -18,6 +18,13 @@
 )
 from mother.optimization.core import MotherTuner
 
+try:
+    import torch  # noqa: F401
+
+    _TORCH_AVAILABLE = True
+except (ImportError, OSError):
+    _TORCH_AVAILABLE = False
+
 
 # define a simple model to run the tests with
 @pytest.fixture()
@@ -267,3 +274,63 @@ def test_tune_mother_feature_selection_model_pipeline(
         evaluated = tuner.study.trials[0].params
         enqueued = model_pipeline.default_parameters()
         all(evaluated[k] == enqueued[k] for k in evaluated.keys() & enqueued.keys())
+
+
+@pytest.mark.skipif(
+    not _TORCH_AVAILABLE,
+    reason="Optuna early-stopping callback tests require the optional torch extra (mother[torch]).",
+)
+class TestGetCallbacks:
+    def test_holdout_cv_disables_early_termination_callback(self, caplog):
+        caplog.set_level("WARNING")
+        tuner = MotherTuner(scorer="neg_mean_squared_error", early_stopping_optuna=True)
+        holdout_cv = PredefinedSplit(test_fold=[-1, -1, 0, 0])
+
+        callbacks = tuner.get_callbacks(cross_validation=holdout_cv)
+
+        assert callbacks is None
+        assert "requires at least 2 CV splits" in caplog.text
+
+    def test_multi_split_cv_keeps_early_termination_callback(self):
+        tuner = MotherTuner(scorer="neg_mean_squared_error", early_stopping_optuna=True)
+        kfold_cv = KFold(n_splits=3, shuffle=True, random_state=42)
+
+        callbacks = tuner.get_callbacks(cross_validation=kfold_cv)
+
+        assert callbacks is not None
+        assert len(callbacks) == 1
+
+    def test_no_cv_argument_keeps_existing_behavior(self):
+        tuner = MotherTuner(scorer="neg_mean_squared_error", early_stopping_optuna=True)
+
+        callbacks = tuner.get_callbacks()
+
+        assert callbacks is not None
+        assert len(callbacks) == 1
+
+    def test_optimize_works_with_holdout_cv(self, lasso_pipeline, synthetic_data, caplog):
+        caplog.set_level("WARNING")
+        X, y, _ = synthetic_data
+        test_fold = np.full(len(X), -1)
+        test_fold[-20:] = 0
+        holdout_cv = PredefinedSplit(test_fold=test_fold)
+
+        tuner = MotherTuner(
+            scorer="neg_mean_squared_error",
+            early_stopping_optuna=True,
+            n_trials_optuna=1,
+        )
+
+        model = tuner.optimize(
+            lasso_pipeline,
+            X=X,
+            y=y,
+            hyperparameter_space_function=lasso_pipeline.get_hyperparameter_space,
+            default_parameters=lasso_pipeline.default_parameters(),
+            cross_validation=holdout_cv,
+        )
+
+        assert model is not None
+        assert tuner.study is not None
+        assert len(tuner.study.trials) >= 1
+        assert "requires at least 2 CV splits" in caplog.text

test/unit/test_tabpfn.py

@@ -208,6 +208,7 @@ def test_prefitted_model(self):
         )
 
         result_prefitted = transformer_prefitted.transform(self.X)
+        assert transformer_prefitted.model.use_autocast_ is False
 
         transformer_newfit = TabPFNEmbeddingTransformer(
             model_type="regression", use_kfold=False, n_estimators=1, random_state=0