add progress bar

Author: randommm

changelog/899.added.md

@@ -0,0 +1 @@
+Add an optional `show_progress_bar` flag to TabPFN classifier and regressor inference, defaulting to `False`.

pyproject.toml

@@ -21,6 +21,7 @@ dependencies = [
   # Once Python 3.10 is the minimum version, this can be removed.
   "eval-type-backport>=0.2.2",
   "joblib>=1.2.0",
+  "tqdm>=4.66.0",
   "tabpfn-common-utils[telemetry-interactive]>=0.2.13",
   "filelock>=3.11.0",
 ]

src/tabpfn/classifier.py

@@ -31,6 +31,7 @@
 from sklearn import config_context
 from sklearn.base import BaseEstimator, ClassifierMixin, check_is_fitted
 from tabpfn_common_utils.telemetry import track_model_call
+from tqdm.auto import tqdm
 
 from tabpfn.base import (
     ClassifierModelSpecs,
@@ -230,6 +231,7 @@ def __init__(  # noqa: PLR0913
         differentiable_input: bool = False,
         eval_metric: str | ClassifierEvalMetrics | None = None,
         tuning_config: dict | ClassifierTuningConfig | None = None,
+        show_progress_bar: bool = False,
     ) -> None:
         """Construct a TabPFN classifier.
 
@@ -453,6 +455,9 @@ class in Fine-Tuning. The fit_from_preprocessed() function sets this
                 `eval_metric`. See
                 [tabpfn.inference_tuning.ClassifierTuningConfig][] for details
                 and options.
+
+            show_progress_bar:
+                Whether to show a progress bar during inference. Defaults to False.
         """
         super().__init__()
         self.n_estimators = n_estimators
@@ -467,6 +472,7 @@ class in Fine-Tuning. The fit_from_preprocessed() function sets this
             inference_precision
         )
         self.fit_mode = fit_mode
+        self.show_progress_bar = show_progress_bar
         self.memory_saving_mode: MemorySavingMode = memory_saving_mode
         self.random_state = random_state
         self.inference_config = inference_config
@@ -1442,10 +1448,16 @@ def forward(  # noqa: C901, PLR0912
             self.executor_.use_torch_inference_mode(use_inference=actual_inference_mode)
 
         outputs = []
-        for output, config in self.executor_.iter_outputs(
-            X,
-            autocast=self.use_autocast_,
-            task_type="multiclass",
+        for output, config in tqdm(
+            self.executor_.iter_outputs(
+                X,
+                autocast=self.use_autocast_,
+                task_type="multiclass",
+            ),
+            total=self.n_estimators,
+            desc="TabPFN inference",
+            unit="estimator",
+            disable=not self.show_progress_bar,
         ):
             original_ndim = output.ndim
 

src/tabpfn/regressor.py

@@ -37,6 +37,7 @@
     check_is_fitted,
 )
 from tabpfn_common_utils.telemetry import track_model_call
+from tqdm.auto import tqdm
 
 from tabpfn.architectures.base.bar_distribution import FullSupportBarDistribution
 from tabpfn.base import (
@@ -236,6 +237,7 @@ def __init__(  # noqa: PLR0913
         n_preprocessing_jobs: int = 1,
         inference_config: dict | InferenceConfig | None = None,
         differentiable_input: bool = False,
+        show_progress_bar: bool = False,
     ) -> None:
         """Construct a TabPFN regressor.
 
@@ -436,6 +438,9 @@ class in Fine-Tuning. The fit_from_preprocessed() function sets this
                 If true, preprocessing attempts to be end-to-end differentiable.
                 Less relevant for standard regression fine-tuning compared to
                 prompt-tuning.
+
+            show_progress_bar:
+                Whether to show a progress bar during inference. Defaults to False.
         """
         super().__init__()
         self.n_estimators = n_estimators
@@ -454,6 +459,7 @@ class in Fine-Tuning. The fit_from_preprocessed() function sets this
             "fit_with_cache",
             "batched",
         ] = fit_mode
+        self.show_progress_bar = show_progress_bar
         self.memory_saving_mode: MemorySavingMode = memory_saving_mode
         self.random_state = random_state
         self.inference_config = inference_config
@@ -967,8 +973,12 @@ def predict(  # noqa: C901, PLR0912
         n_estimators = 0
         accumulated_logits: torch.Tensor | None = None
         with handle_oom_errors(self.devices_, X, model_type="regressor"):
-            for borders_t, output in self._iter_forward_executor(
-                X, use_inference_mode=True
+            for borders_t, output in tqdm(
+                self._iter_forward_executor(X, use_inference_mode=True),
+                total=self.n_estimators,
+                desc="TabPFN inference",
+                unit="estimator",
+                disable=not self.show_progress_bar,
             ):
                 transformed = translate_probs_across_borders(
                     output,
@@ -1164,8 +1174,12 @@ def forward(
         outputs: list[torch.Tensor] = []
         borders: list[np.ndarray] = []
 
-        for border, output in self._iter_forward_executor(
-            X, use_inference_mode=use_inference_mode
+        for border, output in tqdm(
+            self._iter_forward_executor(X, use_inference_mode=use_inference_mode),
+            total=self.n_estimators,
+            desc="TabPFN inference",
+            unit="estimator",
+            disable=not self.show_progress_bar,
         ):
             borders.append(border)
             outputs.append(output)

tests/test_classifier_interface.py

@@ -65,6 +65,34 @@ def X_y() -> tuple[np.ndarray, np.ndarray]:
 fit_modes = ["low_memory", "fit_preprocessors"]
 
 
+def test__show_progress_bar__is_configurable() -> None:
+    model = TabPFNClassifier(show_progress_bar=True)
+    assert model.show_progress_bar is True
+    assert model.get_params()["show_progress_bar"] is True
+
+    default_model = TabPFNClassifier()
+    assert default_model.show_progress_bar is False
+    assert default_model.get_params()["show_progress_bar"] is False
+
+
+def test__predict__show_progress_bar_true__tiny_dataset_does_not_crash() -> None:
+    model = TabPFNClassifier(n_estimators=1, show_progress_bar=True, random_state=42)
+    X, y = sklearn.datasets.make_classification(
+        n_samples=9,
+        n_features=3,
+        n_informative=3,
+        n_redundant=0,
+        n_classes=3,
+        random_state=0,
+    )
+
+    model.fit(X, y)
+
+    predictions = model.predict(X)
+
+    assert predictions.shape == (X.shape[0],)
+
+
 @pytest.mark.parametrize(
     ("device", "n_estimators", "fit_mode", "inference_precision"),
     mark_mps_configs_as_slow(

tests/test_regressor_interface.py

@@ -18,6 +18,7 @@
 from sklearn.utils.estimator_checks import parametrize_with_checks
 from torch import nn
 
+import tabpfn.regressor as regressor_module
 from tabpfn import TabPFNRegressor
 from tabpfn.base import RegressorModelSpecs, initialize_tabpfn_model
 from tabpfn.constants import ModelVersion
@@ -40,6 +41,55 @@
 fit_modes = ["low_memory", "fit_preprocessors"]
 
 
+def test__show_progress_bar__is_configurable() -> None:
+    model = TabPFNRegressor(show_progress_bar=True)
+    assert model.show_progress_bar is True
+    assert model.get_params()["show_progress_bar"] is True
+
+    default_model = TabPFNRegressor()
+    assert default_model.show_progress_bar is False
+    assert default_model.get_params()["show_progress_bar"] is False
+
+
+def test__predict__show_progress_bar_true__tiny_dataset_does_not_crash() -> None:
+    model = TabPFNRegressor(n_estimators=1, show_progress_bar=True, random_state=42)
+    X, y = sklearn.datasets.make_regression(
+        n_samples=9,
+        n_features=3,
+        random_state=0,
+        coef=False,
+    )
+
+    model.fit(X, y)
+
+    predictions = model.predict(X)
+
+    assert predictions.shape == (X.shape[0],)
+
+
+def test__forward__passes_progress_bar_flag(monkeypatch: pytest.MonkeyPatch) -> None:
+    model = TabPFNRegressor(n_estimators=1, show_progress_bar=True, random_state=42)
+    X, y = sklearn.datasets.make_regression(
+        n_samples=9, n_features=3, random_state=0, coef=False
+    )
+    model.fit(X, y)
+
+    captured_kwargs: dict[str, object] = {}
+
+    def fake_tqdm(iterable, **kwargs) -> typing.Iterable[object]:
+        captured_kwargs.update(kwargs)
+        return iterable
+
+    monkeypatch.setattr(regressor_module, "tqdm", fake_tqdm)
+
+    averaged_logits, outputs, borders = model.forward(X, use_inference_mode=True)
+
+    assert averaged_logits is not None
+    assert outputs
+    assert borders
+    assert captured_kwargs["disable"] is False
+
+
 @pytest.fixture(scope="module")
 def X_y() -> tuple[np.ndarray, np.ndarray]:
     X, y, _ = sklearn.datasets.make_regression(