add progress bar

Author: randommm

changelog/899.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,
@@ -222,6 +223,7 @@ def __init__(  # noqa: PLR0913
             "fit_with_cache",
             "batched",
         ] = "fit_preprocessors",
+        show_progress_bar: bool = False,
         memory_saving_mode: MemorySavingMode = "auto",
         random_state: int | np.random.RandomState | np.random.Generator | None = 0,
         n_jobs: Annotated[int | None, deprecated("Use n_preprocessing_jobs")] = None,
@@ -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
@@ -1428,10 +1434,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 getattr(self, "show_progress_bar", False),
         ):
             original_ndim = output.ndim
 

src/tabpfn/regressor.py

@@ -36,6 +36,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 (
@@ -229,6 +230,7 @@ def __init__(  # noqa: PLR0913
             "fit_with_cache",
             "batched",
         ] = "fit_preprocessors",
+        show_progress_bar: bool = False,
         memory_saving_mode: MemorySavingMode = "auto",
         random_state: int | np.random.RandomState | np.random.Generator | None = 0,
         n_jobs: Annotated[int | None, deprecated("Use n_preprocessing_jobs")] = None,
@@ -435,6 +437,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
@@ -453,6 +458,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
@@ -952,8 +958,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 getattr(self, "show_progress_bar", False),
             ):
                 transformed = translate_probs_across_borders(
                     output,

tests/test_classifier_interface.py

@@ -65,6 +65,16 @@ 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
+
+
 @pytest.mark.parametrize(
     ("device", "n_estimators", "fit_mode", "inference_precision"),
     mark_mps_configs_as_slow(

tests/test_regressor_interface.py

@@ -40,6 +40,16 @@
 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
+
+
 @pytest.fixture(scope="module")
 def X_y() -> tuple[np.ndarray, np.ndarray]:
     X, y, _ = sklearn.datasets.make_regression(