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

tabpfn_time_series/__init__.py

@@ -1,5 +1,6 @@
 from .feature import DefaultFeatures, FeatureTransformer
 from .predictor import TabPFNTimeSeriesPredictor, TabPFNMode
+from .defaults import TABPFN_TS_DEFAULT_QUANTILE_CONFIG
 
 __version__ = "0.1.0"
 
@@ -8,4 +9,5 @@ __all__ = [
     "FeatureTransformer",
     "TabPFNTimeSeriesPredictor",
     "TabPFNMode",
+    "TABPFN_TS_DEFAULT_QUANTILE_CONFIG",
 ]

tabpfn_time_series/predictor.py

@@ -3,11 +3,8 @@ 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_TS_DEFAULT_QUANTILE_CONFIG,
-    TABPFN_TS_DEFAULT_CONFIG,
-)
+from tabpfn_time_series.tabpfn_worker import TabPFNClient, LocalTabPFN, MockTabPFN
+from tabpfn_time_series.defaults import TABPFN_TS_DEFAULT_CONFIG
 
 logger = logging.getLogger(__name__)
 
@@ -15,6 +12,7 @@ logger = logging.getLogger(__name__)
 class TabPFNMode(Enum):
     LOCAL = "tabpfn-local"
     CLIENT = "tabpfn-client"
+    MOCK = "tabpfn-mock"
 
 
 class TabPFNTimeSeriesPredictor:
@@ -30,6 +28,7 @@ class TabPFNTimeSeriesPredictor:
         worker_mapping = {
             TabPFNMode.CLIENT: lambda: TabPFNClient(config),
             TabPFNMode.LOCAL: lambda: LocalTabPFN(config),
+            TabPFNMode.MOCK: lambda: MockTabPFN(config),
         }
         self.tabpfn_worker = worker_mapping[tabpfn_mode]()
 
@@ -37,7 +36,6 @@ class TabPFNTimeSeriesPredictor:
         self,
         train_tsdf: TimeSeriesDataFrame,  # with features and target
         test_tsdf: TimeSeriesDataFrame,  # with features only
-        quantile_config: list[float] = TABPFN_TS_DEFAULT_QUANTILE_CONFIG,
     ) -> TimeSeriesDataFrame:
         """
         Predict on each time series individually (local forecasting).
@@ -47,4 +45,4 @@ class TabPFNTimeSeriesPredictor:
             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)
+        return self.tabpfn_worker.predict(train_tsdf, test_tsdf)

tabpfn_time_series/tabpfn_worker.py

@@ -2,8 +2,10 @@ import logging
 from abc import ABC, abstractmethod
 from joblib import Parallel, delayed
 
+from tqdm import tqdm
 import pandas as pd
 import numpy as np
+import torch
 from scipy.stats import norm
 from autogluon.timeseries import TimeSeriesDataFrame
 
@@ -26,14 +28,7 @@ class TabPFNWorker(ABC):
         self,
         train_tsdf: TimeSeriesDataFrame,
         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",
@@ -42,9 +37,8 @@ class TabPFNWorker(ABC):
                 item_id,
                 train_tsdf.loc[item_id],
                 test_tsdf.loc[item_id],
-                quantile_config,
             )
-            for item_id in train_tsdf.item_ids
+            for item_id in tqdm(train_tsdf.item_ids, desc="Predicting time series")
         )
 
         predictions = pd.concat(predictions)
@@ -59,8 +53,9 @@ class TabPFNWorker(ABC):
         item_id: str,
         single_train_tsdf: TimeSeriesDataFrame,
         single_test_tsdf: TimeSeriesDataFrame,
-        quantile_config: list[float],
     ) -> pd.DataFrame:
+        # logger.debug(f"Predicting on item_id: {item_id}")
+
         test_index = single_test_tsdf.index
         train_X, train_y = split_time_series_to_X_y(single_train_tsdf.copy())
         test_X, _ = split_time_series_to_X_y(single_test_tsdf.copy())
@@ -70,7 +65,7 @@ class TabPFNWorker(ABC):
         if train_y_has_constant_value:
             logger.info("Found time-series with constant target")
             result = self._predict_on_constant_train_target(
-                single_train_tsdf, single_test_tsdf, quantile_config
+                single_train_tsdf, single_test_tsdf
             )
         else:
             tabpfn = self._get_tabpfn_engine()
@@ -81,7 +76,9 @@ class TabPFNWorker(ABC):
             result.update(
                 {
                     q: q_pred
-                    for q, q_pred in zip(quantile_config, full_pred["quantiles"])
+                    for q, q_pred in zip(
+                        TABPFN_TS_DEFAULT_QUANTILE_CONFIG, full_pred["quantiles"]
+                    )
                 }
             )
 
@@ -98,7 +95,6 @@ class TabPFNWorker(ABC):
         self,
         single_train_tsdf: TimeSeriesDataFrame,
         single_test_tsdf: TimeSeriesDataFrame,
-        quantile_config: list[float],
     ) -> pd.DataFrame:
         # If train_y is constant, we return the constant value from the training set
         mean_constant = single_train_tsdf.target.iloc[0]
@@ -106,12 +102,14 @@ class TabPFNWorker(ABC):
 
         # For quantile prediction, we assume that the uncertainty follows a standard normal distribution
         quantile_pred_with_uncertainty = norm.ppf(
-            quantile_config, loc=mean_constant, scale=1
+            TABPFN_TS_DEFAULT_QUANTILE_CONFIG, loc=mean_constant, scale=1
         )
         result.update(
             {
                 q: np.full(len(single_test_tsdf), v)
-                for q, v in zip(quantile_config, quantile_pred_with_uncertainty)
+                for q, v in zip(
+                    TABPFN_TS_DEFAULT_QUANTILE_CONFIG, quantile_pred_with_uncertainty
+                )
             }
         )
 
@@ -141,8 +139,52 @@ class LocalTabPFN(TabPFNWorker):
     def __init__(
         self,
         config: dict = {},
+        num_workers_per_gpu: int = 4,  # per GPU
+    ):
+        self.num_workers_per_gpu = num_workers_per_gpu
+
+        # Only support GPU for now (inference on CPU takes too long)
+        if not torch.cuda.is_available():
+            raise ValueError("GPU is required for local TabPFN inference")
+
+        super().__init__(
+            config, num_workers=torch.cuda.device_count() * self.num_workers_per_gpu
+        )
+
+    def predict(
+        self,
+        train_tsdf: TimeSeriesDataFrame,
+        test_tsdf: TimeSeriesDataFrame,
     ):
-        super().__init__(config, num_workers=1)
+        total_num_workers = torch.cuda.device_count() * self.num_workers_per_gpu
+
+        # Split data into chunks for parallel inference on each GPU
+        #   since the time series are of different lengths, we shuffle
+        #   the item_ids s.t. the workload is distributed evenly across GPUs
+        # Also, using 'min' since num_workers could be larger than the number of time series
+        np.random.seed(0)
+        item_ids_chunks = np.array_split(
+            np.random.permutation(train_tsdf.item_ids),
+            min(total_num_workers, len(train_tsdf.item_ids)),
+        )
+
+        # Run predictions in parallel
+        predictions = Parallel(n_jobs=len(item_ids_chunks), backend="loky")(
+            delayed(self._prediction_routine_per_gpu)(
+                train_tsdf.loc[chunk],
+                test_tsdf.loc[chunk],
+                gpu_id=i
+                % torch.cuda.device_count(),  # Alternate between available GPUs
+            )
+            for i, chunk in enumerate(item_ids_chunks)
+        )
+
+        predictions = pd.concat(predictions)
+
+        # Sort predictions according to original item_ids order
+        predictions = predictions.loc[train_tsdf.item_ids]
+
+        return TimeSeriesDataFrame(predictions)
 
     def _get_tabpfn_engine(self):
         from tabpfn import TabPFNRegressor
@@ -151,7 +193,67 @@ class LocalTabPFN(TabPFNWorker):
             config = self.config["tabpfn_internal"].copy()
             config["model_path"] = self._parse_model_path(config["model_path"])
 
-        return TabPFNRegressor(**config)
+        return TabPFNRegressor(**config, random_state=0)
 
     def _parse_model_path(self, model_name: str) -> str:
         return f"tabpfn-v2-regressor-{model_name}.ckpt"
+
+    def _prediction_routine_per_gpu(
+        self,
+        train_tsdf: TimeSeriesDataFrame,
+        test_tsdf: TimeSeriesDataFrame,
+        gpu_id: int,
+    ):
+        # Set GPU
+        torch.cuda.set_device(gpu_id)
+
+        all_pred = []
+        for item_id in tqdm(train_tsdf.item_ids, desc=f"GPU {gpu_id}:"):
+            predictions = self._prediction_routine(
+                item_id,
+                train_tsdf.loc[item_id],
+                test_tsdf.loc[item_id],
+            )
+            all_pred.append(predictions)
+
+        # Clear GPU cache
+        torch.cuda.empty_cache()
+
+        return pd.concat(all_pred)
+
+
+class MockTabPFN(TabPFNWorker):
+    """
+    Mock TabPFN worker that returns random values for predictions.
+    Can be used for testing or debugging.
+    """
+
+    class MockTabPFNRegressor:
+        TABPFN_QUANTILE = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
+
+        def __init__(self, *args, **kwargs):
+            pass
+
+        def fit(self, *args, **kwargs):
+            pass
+
+        def predict(self, test_X, output_type="main", **kwargs):
+            if output_type != "main":
+                raise NotImplementedError(
+                    "Only main output is supported for mock TabPFN"
+                )
+
+            return {
+                "mean": np.random.rand(len(test_X)),
+                "median": np.random.rand(len(test_X)),
+                "mode": np.random.rand(len(test_X)),
+                "quantiles": [
+                    np.random.rand(len(test_X)) for _ in self.TABPFN_QUANTILE
+                ],
+            }
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def _get_tabpfn_engine(self):
+        return self.MockTabPFNRegressor()

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tabpfn_time_series
-Version: 0.1.2
+Version: 0.1.3
 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,11 +10,12 @@ 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==1.2
-Requires-Dist: gluonts==0.16.0
-Requires-Dist: pandas
-Requires-Dist: tabpfn-client==0.1.1
-Requires-Dist: tabpfn==2.0.0
+Requires-Dist: autogluon-timeseries>=1.2
+Requires-Dist: datasets>=3.3.2
+Requires-Dist: gluonts>=0.16.0
+Requires-Dist: pandas<2.2.0,>=2.1.2
+Requires-Dist: tabpfn-client>=0.1.1
+Requires-Dist: tabpfn>=2.0.0
 Requires-Dist: tqdm
 Provides-Extra: dev
 Requires-Dist: build; extra == 'dev'
@@ -24,13 +25,20 @@ Requires-Dist: ruff; extra == 'dev'
 Requires-Dist: twine; extra == 'dev'
 Description-Content-Type: text/markdown
 
-# Time Series Forecasting with TabPFN
+# Zero-Shot Time Series Forecasting with TabPFN
 
+[![PyPI version](https://badge.fury.io/py/tabpfn-time-series.svg)](https://badge.fury.io/py/tabpfn-time-series)
 [![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-2501.02945-<COLOR>.svg)](https://arxiv.org/abs/2501.02945)
 
+## 📌 News
+- **27-01-2025**: 🚀 Ranked _**1st**_ on [GIFT-EVAL](https://huggingface.co/spaces/Salesforce/GIFT-Eval) benchmark<sup>[1]</sup>!
+- **10-10-2024**: 🚀 TabPFN-TS [paper](https://arxiv.org/abs/2501.02945) accepted to NeurIPS 2024 [TRL](https://table-representation-learning.github.io/NeurIPS2024/) and [TSALM](https://neurips-time-series-workshop.github.io/) workshops!
 
+_[1] Last checked on: 10/03/2025_
+
+## ✨ Introduction
 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?
@@ -59,3 +67,9 @@ On top of that, thanks to **[tabpfn-client](https://github.com/automl/tabpfn-cli
 [![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)
 
 The demo should explain it all. 😉
+
+## 📊 GIFT-EVAL Benchmark
+
+We have submitted our results to the [GIFT-EVAL](https://huggingface.co/spaces/Salesforce/GIFT-Eval) benchmark. Stay tuned for results!
+
+For more details regarding the evaluation setup, please refer to [README.md](gift_eval/README.md).

tabpfn_time_series-0.1.3.dist-info/RECORD

@@ -0,0 +1,11 @@
+tabpfn_time_series/__init__.py,sha256=brJLLVOis4tBGOmNk6PCjyk_RaOvFITZgaYChOTVqSo,353
+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=JzuV34zERf1XDLacGzSFJb-o077qd7GlKC6lvD62EPk,1457
+tabpfn_time_series/tabpfn_worker.py,sha256=zvFwg4Dc01_m5emqmVITBr6W_cNZ04tMyntmj40pyPE,8299
+tabpfn_time_series-0.1.3.dist-info/METADATA,sha256=KQZBVKZgMX4e3uxk2LTCuSwruATLowUmgrP6wbcLMB8,4158
+tabpfn_time_series-0.1.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+tabpfn_time_series-0.1.3.dist-info/licenses/LICENSE.txt,sha256=iwhPL7kIWQG6gyLZZwIMDItGrNgxMDIq9itxkUSMapY,11345
+tabpfn_time_series-0.1.3.dist-info/RECORD,,

tabpfn_time_series-0.1.2.dist-info/RECORD

@@ -1,11 +0,0 @@
-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,,