From a5984211eb8199818c1b6a643225aef1ed698f45 Mon Sep 17 00:00:00 2001
From: Wesley Gifford <79663411+wgifford@users.noreply.github.com>
Date: Wed, 14 Aug 2024 12:40:04 -0400
Subject: [PATCH] no import of functions requiring optional imports

---
 .../hfdemo/tinytimemixer/ttm_benchmarking_1024_96.ipynb      | 3 ++-
 notebooks/hfdemo/tinytimemixer/ttm_benchmarking_512_96.ipynb | 5 +++--
 notebooks/hfdemo/ttm_getting_started.ipynb                   | 5 +++--
 tsfm_public/__init__.py                                      | 2 --
 tsfm_public/toolkit/__init__.py                              | 1 -
 5 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/notebooks/hfdemo/tinytimemixer/ttm_benchmarking_1024_96.ipynb b/notebooks/hfdemo/tinytimemixer/ttm_benchmarking_1024_96.ipynb
index 9f2e39e7..92c5e867 100644
--- a/notebooks/hfdemo/tinytimemixer/ttm_benchmarking_1024_96.ipynb
+++ b/notebooks/hfdemo/tinytimemixer/ttm_benchmarking_1024_96.ipynb
@@ -32,7 +32,8 @@
     "from torch.optim.lr_scheduler import OneCycleLR\n",
     "from transformers import EarlyStoppingCallback, Trainer, TrainingArguments, set_seed\n",
     "\n",
-    "from tsfm_public import TinyTimeMixerForPrediction, TrackingCallback, count_parameters, load_dataset, plot_predictions"
+    "from tsfm_public import TinyTimeMixerForPrediction, TrackingCallback, count_parameters, load_dataset\n",
+    "from tsfm_public.toolkit.visualization import plot_predictions"
    ]
   },
   {
diff --git a/notebooks/hfdemo/tinytimemixer/ttm_benchmarking_512_96.ipynb b/notebooks/hfdemo/tinytimemixer/ttm_benchmarking_512_96.ipynb
index 55ec8f72..967f805c 100644
--- a/notebooks/hfdemo/tinytimemixer/ttm_benchmarking_512_96.ipynb
+++ b/notebooks/hfdemo/tinytimemixer/ttm_benchmarking_512_96.ipynb
@@ -32,7 +32,8 @@
     "from torch.optim.lr_scheduler import OneCycleLR\n",
     "from transformers import EarlyStoppingCallback, Trainer, TrainingArguments, set_seed\n",
     "\n",
-    "from tsfm_public import TinyTimeMixerForPrediction, TrackingCallback, count_parameters, load_dataset, plot_predictions"
+    "from tsfm_public import TinyTimeMixerForPrediction, TrackingCallback, count_parameters, load_dataset\n",
+    "from tsfm_public.toolkit.visualization import plot_predictions"
    ]
   },
   {
@@ -3842,7 +3843,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.4"
+   "version": "3.10.13"
   }
  },
  "nbformat": 4,
diff --git a/notebooks/hfdemo/ttm_getting_started.ipynb b/notebooks/hfdemo/ttm_getting_started.ipynb
index d9a24858..5d464e21 100644
--- a/notebooks/hfdemo/ttm_getting_started.ipynb
+++ b/notebooks/hfdemo/ttm_getting_started.ipynb
@@ -43,7 +43,8 @@
     "from torch.optim.lr_scheduler import OneCycleLR\n",
     "from transformers import EarlyStoppingCallback, Trainer, TrainingArguments, set_seed\n",
     "\n",
-    "from tsfm_public import TinyTimeMixerForPrediction, TrackingCallback, count_parameters, load_dataset, plot_predictions"
+    "from tsfm_public import TinyTimeMixerForPrediction, TrackingCallback, count_parameters, load_dataset\n",
+    "from tsfm_public.toolkit.visualization import plot_predictions"
    ]
   },
   {
@@ -873,7 +874,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.4"
+   "version": "3.10.13"
   }
  },
  "nbformat": 4,
diff --git a/tsfm_public/__init__.py b/tsfm_public/__init__.py
index 5725aa8d..400357d0 100644
--- a/tsfm_public/__init__.py
+++ b/tsfm_public/__init__.py
@@ -55,7 +55,6 @@
         "load_dataset",
         "TrackingCallback",
         "count_parameters",
-        "plot_predictions",
     ],
 }
 
@@ -90,7 +89,6 @@
         count_parameters,
         get_datasets,
         load_dataset,
-        plot_predictions,
     )
 else:
     # Standard
diff --git a/tsfm_public/toolkit/__init__.py b/tsfm_public/toolkit/__init__.py
index e05868cb..98e75af8 100644
--- a/tsfm_public/toolkit/__init__.py
+++ b/tsfm_public/toolkit/__init__.py
@@ -7,4 +7,3 @@
 from .time_series_forecasting_pipeline import TimeSeriesForecastingPipeline
 from .time_series_preprocessor import TimeSeriesPreprocessor, get_datasets
 from .util import count_parameters
-from .visualization import plot_predictions