no import of functions requiring optional imports

ibm-granite · Aug 14, 2024 · a598421 · a598421
1 parent b188e80
commit a598421
Show file tree

Hide file tree

Showing 5 changed files with 8 additions and 8 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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