ease import of some toolkit items

ibm-granite · Jul 19, 2024 · c9e1337 · c9e1337
1 parent 8e5d467
commit c9e1337
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 1 deletion.
diff --git a/tsfm_public/__init__.py b/tsfm_public/__init__.py
@@ -16,7 +16,13 @@
 _import_structure = {
     "models": [],
     "models.tinytimemixer": ["TINYTIMEMIXER_PRETRAINED_CONFIG_ARCHIVE_MAP", "TinyTimeMixerConfig"],
-    "toolkit": [],
+    "toolkit": [
+        "TimeSeriesPreprocessor",
+        "TimeSeriesForecastingPipeline",
+        "ForecastDFDataset",
+        "PretrainDFDataset",
+        "RegressionDFDataset",
+    ],
 }
 
 
@@ -40,6 +46,13 @@
         TinyTimeMixerModel,
         TinyTimeMixerPreTrainedModel,
     )
+    from .toolkit import (
+        ForecastDFDataset,
+        PretrainDFDataset,
+        RegressionDFDataset,
+        TimeSeriesForecastingPipeline,
+        TimeSeriesPreprocessor,
+    )
 else:
     # Standard
     import sys

diff --git a/tsfm_public/toolkit/__init__.py b/tsfm_public/toolkit/__init__.py
@@ -1,2 +1,6 @@
 # Copyright contributors to the TSFM project
 #
+
+from .dataset import ForecastDFDataset, PretrainDFDataset, RegressionDFDataset
+from .time_series_forecasting_pipeline import TimeSeriesForecastingPipeline
+from .time_series_preprocessor import TimeSeriesPreprocessor