tabpfn-time-series 1.0.0__py3-none-any.whl → 1.0.1__py3-none-any.whl

tabpfn_time_series/__init__.py

@@ -1,9 +1,18 @@
+from pathlib import Path
+
+try:
+    import tomllib  # Python 3.11+
+except ImportError:
+    import tomli as tomllib  # Python <3.11, requires 'tomli' package
+
+with (Path(__file__).parent.parent / "pyproject.toml").open("rb") as f:
+    __version__ = tomllib.load(f)["project"]["version"]
+
+
 from .features import FeatureTransformer
 from .predictor import TabPFNTimeSeriesPredictor, TabPFNMode
 from .defaults import TABPFN_TS_DEFAULT_QUANTILE_CONFIG
 
-__version__ = "0.1.0"
-
 __all__ = [
     "FeatureTransformer",
     "TabPFNTimeSeriesPredictor",

tabpfn_time_series/defaults.py

@@ -1,7 +1,5 @@
 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_internal": {"model_path": "tabpfn-v2-regressor-2noar4o2.ckpt"},
     "tabpfn_output_selection": "median",  # mean or median
 }

tabpfn_time_series/features/feature_transformer.py

@@ -31,9 +31,9 @@ class FeatureTransformer:
         train_tsdf = tsdf.iloc[: len(train_tsdf)]
         test_tsdf = tsdf.iloc[len(train_tsdf) :]
 
-        assert (
-            not train_tsdf[target_column].isna().any()
-        ), "All target values in train_tsdf should be non-NaN"
+        assert not train_tsdf[target_column].isna().any(), (
+            "All target values in train_tsdf should be non-NaN"
+        )
         assert test_tsdf[target_column].isna().all()
 
         return train_tsdf, test_tsdf

tabpfn_time_series/tabpfn_worker.py

@@ -122,18 +122,39 @@ class TabPFNClient(TabPFNWorker):
         config: dict = {},
         num_workers: int = 2,
     ):
-        super().__init__(config, num_workers)
-
         # Initialize the TabPFN client (e.g. sign up, login, etc.)
         from tabpfn_client import init
 
         init()
 
+        # Parse the model name (only needed for TabPFNClient)
+        config = config.copy()
+        config["tabpfn_internal"]["model_path"] = self._parse_model_name(
+            config["tabpfn_internal"]["model_path"]
+        )
+
+        super().__init__(config, num_workers)
+
     def _get_tabpfn_engine(self):
         from tabpfn_client import TabPFNRegressor
 
         return TabPFNRegressor(**self.config["tabpfn_internal"])
 
+    def _parse_model_name(self, model_name: str) -> str:
+        from tabpfn_client import TabPFNRegressor
+
+        available_models = TabPFNRegressor.list_available_models()
+
+        for m in available_models:
+            # Model names from tabpfn_client are abbreviated
+            # e.g. "tabpfn-v2-regressor-2noar4o2.ckpt" -> "2noar4o2"
+            if m in model_name:
+                return m
+        raise ValueError(
+            f"Model {model_name} not found. Available models: {available_models}."
+            "Note that model names from tabpfn_client are abbreviated (e.g. 'tabpfn-v2-regressor-2noar4o2.ckpt' -> '2noar4o2')"
+        )
+
 
 class LocalTabPFN(TabPFNWorker):
     def __init__(
@@ -151,6 +172,9 @@ class LocalTabPFN(TabPFNWorker):
             config, num_workers=torch.cuda.device_count() * self.num_workers_per_gpu
         )
 
+        # Download the model specified in the config
+        self._download_model(self.config["tabpfn_internal"]["model_path"])
+
     def predict(
         self,
         train_tsdf: TimeSeriesDataFrame,
@@ -186,17 +210,29 @@ class LocalTabPFN(TabPFNWorker):
 
         return TimeSeriesDataFrame(predictions)
 
-    def _get_tabpfn_engine(self):
-        from tabpfn import TabPFNRegressor
+    @staticmethod
+    def _download_model(model_name: str):
+        from tabpfn.model.loading import resolve_model_path, download_model
 
-        if "model_path" in self.config["tabpfn_internal"]:
-            config = self.config["tabpfn_internal"].copy()
-            config["model_path"] = self._parse_model_path(config["model_path"])
+        # Resolve the model path
+        # If the model path is not specified, this resolves to the default model path
+        model_path, _, model_name, which = resolve_model_path(
+            model_name,
+            which="regressor",
+        )
 
-        return TabPFNRegressor(**config, random_state=0)
+        if not model_path.exists():
+            download_model(
+                to=model_path,
+                which=which,
+                version="v2",
+                model_name=model_name,
+            )
+
+    def _get_tabpfn_engine(self):
+        from tabpfn import TabPFNRegressor
 
-    def _parse_model_path(self, model_name: str) -> str:
-        return f"tabpfn-v2-regressor-{model_name}.ckpt"
+        return TabPFNRegressor(**self.config["tabpfn_internal"], random_state=0)
 
     def _prediction_routine_per_gpu(
         self,

{tabpfn_time_series-1.0.0.dist-info → tabpfn_time_series-1.0.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tabpfn_time_series
-Version: 1.0.0
+Version: 1.0.1
 Summary: Zero-shot time series forecasting with TabPFNv2
 Project-URL: Homepage, https://github.com/liam-sbhoo/tabpfn-time-series
 Project-URL: Bug Tracker, https://github.com/liam-sbhoo/tabpfn-time-series/issues
@@ -11,19 +11,22 @@ Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.10
 Requires-Dist: autogluon-timeseries>=1.2
-Requires-Dist: datasets>=3.3.2
+Requires-Dist: datasets>=4.0
 Requires-Dist: gluonts>=0.16.0
 Requires-Dist: pandas<2.2.0,>=2.1.2
 Requires-Dist: python-dotenv>=1.1.0
 Requires-Dist: pyyaml>=6.0.1
 Requires-Dist: tabpfn-client>=0.1.7
 Requires-Dist: tabpfn>=2.0.9
+Requires-Dist: tomli>=2.2.1
 Requires-Dist: tqdm
 Provides-Extra: dev
 Requires-Dist: build; extra == 'dev'
+Requires-Dist: ipykernel>=6.29.5; extra == 'dev'
 Requires-Dist: jupyter; extra == 'dev'
 Requires-Dist: pre-commit; extra == 'dev'
-Requires-Dist: ruff; extra == 'dev'
+Requires-Dist: pytest; extra == 'dev'
+Requires-Dist: ruff~=0.12.0; extra == 'dev'
 Requires-Dist: submitit>=1.5.2; extra == 'dev'
 Requires-Dist: twine; extra == 'dev'
 Requires-Dist: wandb>=0.19.8; extra == 'dev'
@@ -34,7 +37,7 @@ Description-Content-Type: text/markdown
 > Zero-Shot Time Series Forecasting with TabPFNv2
 
 [![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)
+[![colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/PriorLabs/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.02945v3)
 
@@ -63,15 +66,33 @@ Concretely, we:
 For more details, please refer to our [paper](https://arxiv.org/abs/2501.02945v3).
  <!-- and our [poster](docs/tabpfn-ts-neurips-poster.pdf) (presented at NeurIPS 2024 TRL and TSALM workshops). -->
 
-## 👉 **Why gives us a try?**
+## 👉 **Why give us a try?**
 - **Zero-shot forecasting**: this method is extremely fast and requires no training, making it highly accessible for experimenting with your own problems.
 - **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 TabPFNv2. 😉 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 TabPFNv2. 😉 We have included `tabpfn-client` as the default engine in our implementation.
+
+## ⚙️ Installation
+
+You can install the package via pip:
+
+```bash
+pip install tabpfn-time-series
+```
+
+### For Developers
+
+To install the package in editable mode with all development dependencies, run the following command in your terminal:
+
+```bash
+pip install -e ".[dev]"
+# or with uv
+uv pip install -e ".[dev]"
+```
 
 ## How to use it?
 
-[![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)
+[![colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/PriorLabs/tabpfn-time-series/blob/main/demo.ipynb)
 
 The demo should explain it all. 😉

{tabpfn_time_series-1.0.0.dist-info → tabpfn_time_series-1.0.1.dist-info}/RECORD

@@ -1,15 +1,15 @@
-tabpfn_time_series/__init__.py,sha256=3XGvQieVbONwhVtn1rITet6HNiTMWQTxHm2xLlGI5ew,314
+tabpfn_time_series/__init__.py,sha256=cuoY5WBdBZM9RUM8bX4-Gyf5M9uQ-905lkQIDSyojTw,578
 tabpfn_time_series/data_preparation.py,sha256=iNW7sAnRkTgmzzOEHBhkkTwm_lQ3p_Q9xgAQ5PbkOts,5416
-tabpfn_time_series/defaults.py,sha256=u2_JnwxiZ5NNibzyNpsE63KuP3TcmOL1iAP8llZ2rJk,238
+tabpfn_time_series/defaults.py,sha256=ki1y38FR4zmbHWgRjcryA5T88GzNMwhlZC-sTRjuK2U,248
 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/tabpfn_worker.py,sha256=NeYPX1XcPe3jF5yFPMUR_3Lq5hw4RpxVFTD7z_ikahA,9608
 tabpfn_time_series/features/__init__.py,sha256=lzdZWkEfntfg3ZHqNNbfbg-3o_VIzju0tebdRu3AzF4,421
 tabpfn_time_series/features/auto_features.py,sha256=3OqqY2h7umcoLjLx4hOXypLTjwzrMtd6cQKTNi83vrU,11561
 tabpfn_time_series/features/basic_features.py,sha256=OV3B__S30-CX88vGjwYQDWqAbJajQw80PxcnvJVUbm4,2955
 tabpfn_time_series/features/feature_generator_base.py,sha256=jtySWLJyX4E31v6CbX44EHa8cdz7OMyauf4ltNEQeAQ,534
-tabpfn_time_series/features/feature_transformer.py,sha256=mUsbnPUhJ4lPcnGWk8Ag1hgCOE1V5I0iQRT4VFgQEso,1763
-tabpfn_time_series-1.0.0.dist-info/METADATA,sha256=CvXqIOHNTKyd-zpCednsqa3FloPk6lFJ4ISG0eSEWx4,4434
-tabpfn_time_series-1.0.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-tabpfn_time_series-1.0.0.dist-info/licenses/LICENSE.txt,sha256=iwhPL7kIWQG6gyLZZwIMDItGrNgxMDIq9itxkUSMapY,11345
-tabpfn_time_series-1.0.0.dist-info/RECORD,,
+tabpfn_time_series/features/feature_transformer.py,sha256=JgPjAUiNBJGfx5PUOFVOZc467AW89Fe-JE08dKQ0AjY,1763
+tabpfn_time_series-1.0.1.dist-info/METADATA,sha256=mSPGoTbkVFjvJswthF0DAUvql7NgHD9wRzVkHu1ia1I,4873
+tabpfn_time_series-1.0.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+tabpfn_time_series-1.0.1.dist-info/licenses/LICENSE.txt,sha256=iwhPL7kIWQG6gyLZZwIMDItGrNgxMDIq9itxkUSMapY,11345
+tabpfn_time_series-1.0.1.dist-info/RECORD,,