redirect to test models

ibm-granite · Aug 16, 2024 · bbab0f5 · bbab0f5
1 parent 0263bd5
commit bbab0f5
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/services/inference/tests/test_inference.py b/services/inference/tests/test_inference.py
@@ -68,7 +68,7 @@ def test_zero_shot_forecast_inference(ts_data):
     test_data_ = test_data[test_data[id_columns[0]] == "a"].copy()
 
     msg = {
-        "model_id": "ibm-granite/granite-timeseries-ttm-v1",
+        "model_id": "ibm/test-ttm-v1",
         "parameters": {
             # "prediction_length": params["prediction_length"],
         },
@@ -89,7 +89,7 @@ def test_zero_shot_forecast_inference(ts_data):
     test_data_ = test_data.copy()
 
     msg = {
-        "model_id": "ibm-granite/granite-timeseries-ttm-v1",
+        "model_id": "ibm/test-ttm-v1",
         "parameters": {
             # "prediction_length": params["prediction_length"],
         },

diff --git a/tests/toolkit/test_time_series_forecasting_pipeline.py b/tests/toolkit/test_time_series_forecasting_pipeline.py
@@ -7,7 +7,7 @@
 import pytest
 from transformers import PatchTSTForPrediction
 
-from tsfm_public import TinyTimeMixerForPrediction
+from tsfm_public import TinyTimeMixerConfig, TinyTimeMixerForPrediction
 from tsfm_public.toolkit.time_series_forecasting_pipeline import (
     TimeSeriesForecastingPipeline,
 )
@@ -25,8 +25,10 @@ def patchtst_model():
 
 @pytest.fixture(scope="module")
 def ttm_model():
-    model_path = "ibm-granite/granite-timeseries-ttm-v1"
-    model = TinyTimeMixerForPrediction.from_pretrained(model_path)
+    # model_path = "ibm-granite/granite-timeseries-ttm-v1"
+
+    conf = TinyTimeMixerConfig()
+    model = TinyTimeMixerForPrediction(conf)
 
     return model