tabpfn-time-series 0.1.1__py3-none-any.whl → 0.1.2__py3-none-any.whl

tabpfn_time_series/__init__.py

@@ -5,7 +5,7 @@ __version__ = "0.1.0"
 
 __all__ = [
     "DefaultFeatures",
-    "FeatureTransformer", 
+    "FeatureTransformer",
     "TabPFNTimeSeriesPredictor",
     "TabPFNMode",
 ]

tabpfn_time_series/defaults.py

@@ -1,5 +1,7 @@
-TABPFN_DEFAULT_QUANTILE_CONFIG = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
-TABPFN_DEFAULT_CONFIG = {
-    "model": "2noar4o2",
-    "optimize_metric": "median",
+TABPFN_TS_DEFAULT_QUANTILE_CONFIG = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
+TABPFN_TS_DEFAULT_CONFIG = {
+    "tabpfn_internal": {
+        "model_path": "2noar4o2",
+    },
+    "tabpfn_output_selection": "median",  # mean or median
 }

tabpfn_time_series/predictor.py

@@ -4,7 +4,10 @@ from enum import Enum
 from autogluon.timeseries import TimeSeriesDataFrame
 
 from tabpfn_time_series.tabpfn_worker import TabPFNClient, LocalTabPFN
-from tabpfn_time_series.defaults import TABPFN_DEFAULT_QUANTILE_CONFIG, TABPFN_DEFAULT_CONFIG
+from tabpfn_time_series.defaults import (
+    TABPFN_TS_DEFAULT_QUANTILE_CONFIG,
+    TABPFN_TS_DEFAULT_CONFIG,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -22,11 +25,11 @@ class TabPFNTimeSeriesPredictor:
     def __init__(
         self,
         tabpfn_mode: TabPFNMode = TabPFNMode.CLIENT,
-        tabpfn_config: dict = TABPFN_DEFAULT_CONFIG,
+        config: dict = TABPFN_TS_DEFAULT_CONFIG,
     ) -> None:
         worker_mapping = {
-            TabPFNMode.CLIENT: lambda: TabPFNClient(tabpfn_config),
-            TabPFNMode.LOCAL: lambda: LocalTabPFN(tabpfn_config),
+            TabPFNMode.CLIENT: lambda: TabPFNClient(config),
+            TabPFNMode.LOCAL: lambda: LocalTabPFN(config),
         }
         self.tabpfn_worker = worker_mapping[tabpfn_mode]()
 
@@ -34,14 +37,14 @@ class TabPFNTimeSeriesPredictor:
         self,
         train_tsdf: TimeSeriesDataFrame,  # with features and target
         test_tsdf: TimeSeriesDataFrame,  # with features only
-        quantile_config: list[float] = TABPFN_DEFAULT_QUANTILE_CONFIG,
+        quantile_config: list[float] = TABPFN_TS_DEFAULT_QUANTILE_CONFIG,
     ) -> TimeSeriesDataFrame:
         """
         Predict on each time series individually (local forecasting).
         """
 
         logger.info(
-            f"Predicting {len(train_tsdf.item_ids)} time series with config{self.tabpfn_worker.tabpfn_config}"
+            f"Predicting {len(train_tsdf.item_ids)} time series with config{self.tabpfn_worker.config}"
         )
 
         return self.tabpfn_worker.predict(train_tsdf, test_tsdf, quantile_config)

tabpfn_time_series/tabpfn_worker.py

@@ -8,7 +8,7 @@ from scipy.stats import norm
 from autogluon.timeseries import TimeSeriesDataFrame
 
 from tabpfn_time_series.data_preparation import split_time_series_to_X_y
-from tabpfn_time_series.defaults import TABPFN_DEFAULT_QUANTILE_CONFIG
+from tabpfn_time_series.defaults import TABPFN_TS_DEFAULT_QUANTILE_CONFIG
 
 logger = logging.getLogger(__name__)
 
@@ -16,10 +16,10 @@ logger = logging.getLogger(__name__)
 class TabPFNWorker(ABC):
     def __init__(
         self,
-        tabpfn_config: dict = {},
+        config: dict = {},
         num_workers: int = 1,
     ):
-        self.tabpfn_config = tabpfn_config
+        self.config = config
         self.num_workers = num_workers
 
     def predict(
@@ -28,6 +28,12 @@ class TabPFNWorker(ABC):
         test_tsdf: TimeSeriesDataFrame,
         quantile_config: list[float],
     ):
+        if not set(quantile_config).issubset(set(TABPFN_TS_DEFAULT_QUANTILE_CONFIG)):
+            raise NotImplementedError(
+                f"We currently only supports {TABPFN_TS_DEFAULT_QUANTILE_CONFIG} for quantile prediction,"
+                f" but got {quantile_config}."
+            )
+
         predictions = Parallel(
             n_jobs=self.num_workers,
             backend="loky",
@@ -67,12 +73,16 @@ class TabPFNWorker(ABC):
                 single_train_tsdf, single_test_tsdf, quantile_config
             )
         else:
-            # Call worker-specific prediction routine
-            result = self._worker_specific_prediction_routine(
-                train_X,
-                train_y,
-                test_X,
-                quantile_config,
+            tabpfn = self._get_tabpfn_engine()
+            tabpfn.fit(train_X, train_y)
+            full_pred = tabpfn.predict(test_X, output_type="main")
+
+            result = {"target": full_pred[self.config["tabpfn_output_selection"]]}
+            result.update(
+                {
+                    q: q_pred
+                    for q, q_pred in zip(quantile_config, full_pred["quantiles"])
+                }
             )
 
         result = pd.DataFrame(result, index=test_index)
@@ -81,13 +91,7 @@ class TabPFNWorker(ABC):
         return result
 
     @abstractmethod
-    def _worker_specific_prediction_routine(
-        self,
-        train_X: pd.DataFrame,
-        train_y: pd.Series,
-        test_X: pd.DataFrame,
-        quantile_config: list[float],
-    ) -> pd.DataFrame:
+    def _get_tabpfn_engine(self):
         pass
 
     def _predict_on_constant_train_target(
@@ -117,108 +121,37 @@ class TabPFNWorker(ABC):
 class TabPFNClient(TabPFNWorker):
     def __init__(
         self,
-        tabpfn_config: dict = {},
+        config: dict = {},
         num_workers: int = 2,
     ):
-        super().__init__(tabpfn_config, num_workers)
+        super().__init__(config, num_workers)
 
         # Initialize the TabPFN client (e.g. sign up, login, etc.)
         from tabpfn_client import init
 
         init()
 
-    def predict(
-        self,
-        train_tsdf: TimeSeriesDataFrame,
-        test_tsdf: TimeSeriesDataFrame,
-        quantile_config: list[float],
-    ):
-        if not set(quantile_config).issubset(set(TABPFN_DEFAULT_QUANTILE_CONFIG)):
-            raise NotImplementedError(
-                f"TabPFNClient currently only supports {TABPFN_DEFAULT_QUANTILE_CONFIG} for quantile prediction,"
-                f" but got {quantile_config}."
-            )
-
-        return super().predict(train_tsdf, test_tsdf, quantile_config)
-
-    def _worker_specific_prediction_routine(
-        self,
-        train_X: pd.DataFrame,
-        train_y: pd.Series,
-        test_X: pd.DataFrame,
-        quantile_config: list[float],
-    ) -> pd.DataFrame:
+    def _get_tabpfn_engine(self):
         from tabpfn_client import TabPFNRegressor
 
-        tabpfn = TabPFNRegressor(**self.tabpfn_config)
-        tabpfn.fit(train_X, train_y)
-        full_pred = tabpfn.predict_full(test_X)
-
-        result = {"target": full_pred[self._get_optimization_mode()]}
-        result.update({q: full_pred[f"quantile_{q:.2f}"] for q in quantile_config})
-
-        return result
-
-    def _get_optimization_mode(self):
-        if (
-            "optimize_metric" not in self.tabpfn_config
-            or self.tabpfn_config["optimize_metric"] is None
-        ):
-            return "mean"
-        elif self.tabpfn_config["optimize_metric"] in ["rmse", "mse", "r2", "mean"]:
-            return "mean"
-        elif self.tabpfn_config["optimize_metric"] in ["mae", "median"]:
-            return "median"
-        elif self.tabpfn_config["optimize_metric"] in ["mode", "exact_match"]:
-            return "mode"
-        else:
-            raise ValueError(f"Unknown metric {self.tabpfn_config['optimize_metric']}")
+        return TabPFNRegressor(**self.config["tabpfn_internal"])
 
 
 class LocalTabPFN(TabPFNWorker):
     def __init__(
         self,
-        tabpfn_config: dict = {},
+        config: dict = {},
     ):
-        # Local TabPFN has a different interface for declaring the model
-        if "model" in tabpfn_config:
-            config = tabpfn_config.copy()
-            config["model_path"] = self._parse_model_path(config["model"])
-            del config["model"]
-            tabpfn_config = config
-
-        super().__init__(tabpfn_config, num_workers=1)
+        super().__init__(config, num_workers=1)
 
-    def _worker_specific_prediction_routine(
-        self,
-        train_X: pd.DataFrame,
-        train_y: pd.Series,
-        test_X: pd.DataFrame,
-        quantile_config: list[float],
-    ) -> pd.DataFrame:
+    def _get_tabpfn_engine(self):
         from tabpfn import TabPFNRegressor
 
-        tabpfn = TabPFNRegressor(**self.tabpfn_config)
-        tabpfn.fit(train_X, train_y)
-        full_pred = tabpfn.predict_full(test_X)
-
-        result = {"target": full_pred[tabpfn.get_optimization_mode()]}
-        if set(quantile_config).issubset(set(TABPFN_DEFAULT_QUANTILE_CONFIG)):
-            result.update({q: full_pred[f"quantile_{q:.2f}"] for q in quantile_config})
-        else:
-            import torch
+        if "model_path" in self.config["tabpfn_internal"]:
+            config = self.config["tabpfn_internal"].copy()
+            config["model_path"] = self._parse_model_path(config["model_path"])
 
-            criterion = full_pred["criterion"]
-            logits = torch.tensor(full_pred["logits"])
-            result.update({q: criterion.icdf(logits, q) for q in quantile_config})
-
-        return result
+        return TabPFNRegressor(**config)
 
     def _parse_model_path(self, model_name: str) -> str:
-        from pathlib import Path
-        import importlib.util
-
-        tabpfn_path = Path(importlib.util.find_spec("tabpfn").origin).parent
-        return str(
-            tabpfn_path / "model_cache" / f"model_hans_regression_{model_name}.ckpt"
-        )
+        return f"tabpfn-v2-regressor-{model_name}.ckpt"

{tabpfn_time_series-0.1.1.dist-info → tabpfn_time_series-0.1.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tabpfn_time_series
-Version: 0.1.1
+Version: 0.1.2
 Summary: Zero-shot time series forecasting with TabPFN
 Project-URL: Homepage, https://github.com/liam-sbhoo/tabpfn-time-series
 Project-URL: Bug Tracker, https://github.com/liam-sbhoo/tabpfn-time-series/issues
@@ -10,13 +10,15 @@ Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.10
-Requires-Dist: autogluon-timeseries
-Requires-Dist: gluonts
+Requires-Dist: autogluon-timeseries==1.2
+Requires-Dist: gluonts==0.16.0
 Requires-Dist: pandas
-Requires-Dist: tabpfn-client
+Requires-Dist: tabpfn-client==0.1.1
+Requires-Dist: tabpfn==2.0.0
 Requires-Dist: tqdm
 Provides-Extra: dev
 Requires-Dist: build; extra == 'dev'
+Requires-Dist: jupyter; extra == 'dev'
 Requires-Dist: pre-commit; extra == 'dev'
 Requires-Dist: ruff; extra == 'dev'
 Requires-Dist: twine; extra == 'dev'
@@ -26,10 +28,10 @@ Description-Content-Type: text/markdown
 
 [![colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/liam-sbhoo/tabpfn-time-series/blob/main/demo.ipynb)
 [![Discord](https://img.shields.io/discord/1285598202732482621?color=7289da&label=Discord&logo=discord&logoColor=ffffff)](https://discord.com/channels/1285598202732482621/)
-[![arXiv](https://img.shields.io/badge/arXiv-<INDEX>-<COLOR>.svg)](https://arxiv.org/abs/2501.02945)
+[![arXiv](https://img.shields.io/badge/arXiv-2501.02945-<COLOR>.svg)](https://arxiv.org/abs/2501.02945)
 
 
-We demonstrate that the tabular foundation model **TabPFN**, when paired with minimal featurization, can perform zero-shot time series forecasting. Its performance on point forecasting matches or even slightly outperforms state-of-the-art methods.
+We demonstrate that the tabular foundation model **[TabPFN](https://github.com/PriorLabs/TabPFN)**, when paired with minimal featurization, can perform zero-shot time series forecasting. Its performance on point forecasting matches or even slightly outperforms state-of-the-art methods.
 
 ## 📖 How does it work?
 
@@ -50,7 +52,7 @@ For more details, please refer to our [paper](https://arxiv.org/abs/2501.02945)
 - **Point and probabilistic forecasting**: it provides accurate point forecasts as well as probabilistic forecasts.
 - **Support for exogenous variables**: if you have exogenous variables, this method can seemlessly incorporate them into the forecasting model.
 
-On top of that, thanks to [tabpfn-client](https://github.com/automl/tabpfn-client) from [Prior Labs](https://priorlabs.ai), you won’t even need your own GPU to run fast inference with TabPFN. 😉 We have included `tabpfn-client` as the default engine in our implementation.
+On top of that, thanks to **[tabpfn-client](https://github.com/automl/tabpfn-client)** from **[Prior Labs](https://priorlabs.ai)**, you won’t even need your own GPU to run fast inference with TabPFN. 😉 We have included `tabpfn-client` as the default engine in our implementation.
 
 ## How to use it?
 

tabpfn_time_series-0.1.2.dist-info/RECORD

@@ -0,0 +1,11 @@
+tabpfn_time_series/__init__.py,sha256=5ruHrmKBQRIZ3WXLA8du4JKttF55ntnI74hkRsHThQ8,256
+tabpfn_time_series/data_preparation.py,sha256=iNW7sAnRkTgmzzOEHBhkkTwm_lQ3p_Q9xgAQ5PbkOts,5416
+tabpfn_time_series/defaults.py,sha256=u2_JnwxiZ5NNibzyNpsE63KuP3TcmOL1iAP8llZ2rJk,238
+tabpfn_time_series/feature.py,sha256=_9FxfQfgPOOO1MiT8hB8523eZ3Nc5oKuoY7vcohKZZc,2531
+tabpfn_time_series/plot.py,sha256=bwSYcWBanzPrUxXKFsbqG8fyGsOJZfgU2v3NsxzTSXo,6571
+tabpfn_time_series/predictor.py,sha256=W9JijaxFaR0chfiW7m4RuDQ0wrRcJezDWVwCBEOQDFk,1502
+tabpfn_time_series/tabpfn_worker.py,sha256=XNpqLEW51PgzrEopNNdtGdYArMCHT4yeBK3BS3z25K0,5021
+tabpfn_time_series-0.1.2.dist-info/METADATA,sha256=hO69b8GN3GDRIetG4DGtxpdMubc8sm8h_aI2RwEto2U,3285
+tabpfn_time_series-0.1.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+tabpfn_time_series-0.1.2.dist-info/licenses/LICENSE.txt,sha256=iwhPL7kIWQG6gyLZZwIMDItGrNgxMDIq9itxkUSMapY,11345
+tabpfn_time_series-0.1.2.dist-info/RECORD,,

tabpfn_time_series-0.1.1.dist-info/RECORD

@@ -1,11 +0,0 @@
-tabpfn_time_series/__init__.py,sha256=atVNap8tQLI0t5COkkCop-wbY3y1FdVxIMfvCf6VDsQ,257
-tabpfn_time_series/data_preparation.py,sha256=iNW7sAnRkTgmzzOEHBhkkTwm_lQ3p_Q9xgAQ5PbkOts,5416
-tabpfn_time_series/defaults.py,sha256=C9HiD7Zm0BzVfE9e2f8nhpiPQSYx79hWozvzb-93L40,165
-tabpfn_time_series/feature.py,sha256=_9FxfQfgPOOO1MiT8hB8523eZ3Nc5oKuoY7vcohKZZc,2531
-tabpfn_time_series/plot.py,sha256=bwSYcWBanzPrUxXKFsbqG8fyGsOJZfgU2v3NsxzTSXo,6571
-tabpfn_time_series/predictor.py,sha256=YfJIe8KsyzkwgX4EFAHR8dDp-mqSv9WK88_qO_EXlws,1505
-tabpfn_time_series/tabpfn_worker.py,sha256=3xInPzzQtmIBPjbc_5TaQsX3-Bl3WOlxttqj3KZlC9Q,7395
-tabpfn_time_series-0.1.1.dist-info/METADATA,sha256=EK_9xSO0-EiE-YzlGietFOwOdD7gt5zFI5gqkzo1IMk,3147
-tabpfn_time_series-0.1.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-tabpfn_time_series-0.1.1.dist-info/licenses/LICENSE.txt,sha256=iwhPL7kIWQG6gyLZZwIMDItGrNgxMDIq9itxkUSMapY,11345
-tabpfn_time_series-0.1.1.dist-info/RECORD,,