From c9e1337ad4548e30c6a3695a366e18e739f08fba Mon Sep 17 00:00:00 2001
From: "Wesley M. Gifford" <wmgifford@us.ibm.com>
Date: Fri, 19 Jul 2024 09:36:05 -0400
Subject: [PATCH] ease import of some toolkit items

---
 tsfm_public/__init__.py         | 15 ++++++++++++++-
 tsfm_public/toolkit/__init__.py |  4 ++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/tsfm_public/__init__.py b/tsfm_public/__init__.py
index 214a3921..1eb66186 100644
--- 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
index 4f85bd0b..6c019cb7 100644
--- 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