Merge pull request #116 from ibm-granite/pipeline_updates

Pipeline updates

tests/toolkit/test_time_series_forecasting_pipeline.py

@@ -84,7 +84,7 @@ def etth_data():
 def test_forecasting_pipeline_defaults():
     model = PatchTSTForPrediction(PatchTSTConfig(prediction_length=3, context_length=33))
 
-    tspipe = TimeSeriesForecastingPipeline(model=model)
+    tspipe = TimeSeriesForecastingPipeline(model)
 
     assert tspipe._preprocess_params["prediction_length"] == 3
     assert tspipe._preprocess_params["context_length"] == 33

tsfm_public/toolkit/time_series_forecasting_pipeline.py

@@ -8,6 +8,7 @@
 import pandas as pd
 import torch
 from torch.utils.data import DataLoader
+from transformers import PreTrainedModel
 from transformers.data.data_collator import default_data_collator
 from transformers.pipelines.base import (
     GenericTensor,
@@ -112,6 +113,7 @@ class TimeSeriesForecastingPipeline(TimeSeriesPipeline):
 
     def __init__(
         self,
+        model: Union["PreTrainedModel"],
         *args,
         freq: Optional[str] = None,
         explode_forecasts: bool = False,
@@ -143,10 +145,6 @@ def __init__(
             if "freq" not in kwargs:
                 kwargs["freq"] = kwargs["feature_extractor"].freq
 
-        model = kwargs.get("model", None)
-        if not model:
-            raise ValueError("A model must be supplied during instantiation of a TimeSeriesForecastingPipeline")
-
         if "context_length" not in kwargs:
             kwargs["context_length"] = model.config.context_length
 
@@ -161,7 +159,7 @@ def __init__(
         else:
             kwargs["frequency_token"] = None
 
-        super().__init__(*args, **kwargs)
+        super().__init__(model, *args, **kwargs)
 
         if self.framework == "tf":
             raise ValueError(f"The {self.__class__} is only available in PyTorch.")