Update tools/ts_forecast_tool.py

tools/ts_forecast_tool.py

@@ -1,30 +1,115 @@
 # space/tools/ts_forecast_tool.py
+import os
+from typing import Optional, Dict
+
 import torch
 import pandas as pd
-from transformers import AutoModelForTimeSeriesForecasting, AutoTokenizer
+
+from utils.tracing import Tracer
+from utils.config import AppConfig
+
+# We avoid unavailable task-specific heads.
+# Use a generic AutoModel and attempt capability-based calls.
+from transformers import AutoModel, AutoConfig
+
 
 class TimeseriesForecastTool:
     """
-    Lightweight wrapper around ibm-granite/granite-timeseries-ttm-r1
-    using the Transformers interface.
+    Lightweight wrapper around 'ibm-granite/granite-timeseries-ttm-r1' for zero-shot forecasting.
+
+    This wrapper:
+      - loads the model with `AutoModel.from_pretrained`
+      - checks for a `.predict(...)` method first
+      - else tries calling the model with `prediction_length=horizon`
+      - returns a Pandas DataFrame with a single 'forecast' column
+
+    Expected input:
+      - series: pd.Series with a DatetimeIndex (regular frequency recommended)
+      - horizon: int, number of future steps
+
+    NOTE:
+      Different library versions expose different APIs. If your environment/model
+      lacks a compatible inference method, we raise a clear RuntimeError with
+      guidance rather than failing at import time.
     """
-    def __init__(self,
-                 model_id="ibm-granite/granite-timeseries-ttm-r1",
-                 device=None):
+
+    def __init__(
+        self,
+        cfg: Optional[AppConfig],
+        tracer: Optional[Tracer],
+        model_id: str = "ibm-granite/granite-timeseries-ttm-r1",
+        device: Optional[str] = None,
+    ):
+        self.cfg = cfg
+        self.tracer = tracer
+        self.model_id = model_id
+
         self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
-        self.model = AutoModelForTimeSeriesForecasting.from_pretrained(model_id).to(self.device)
-        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
-
-    def zeroshot_forecast(self, series: pd.Series, horizon: int = 96):
-        """
-        series: pd.Series indexed by datetime
-        horizon: forecast steps
-        """
-        values = series.values.astype("float32")
-        inputs = torch.tensor(values, dtype=torch.float32).unsqueeze(0).to(self.device)
+        # Load config + model generically
+        self.config = AutoConfig.from_pretrained(self.model_id)
+        self.model = AutoModel.from_pretrained(self.model_id)
+        self.model.to(self.device)
+        self.model.eval()
+
+    def zeroshot_forecast(self, series: pd.Series, horizon: int = 96) -> pd.DataFrame:
+        if not isinstance(series, pd.Series):
+            raise ValueError("series must be a pandas Series")
+        if series.empty:
+            return pd.DataFrame(columns=["forecast"])
+
+        # Ensure numeric tensor
+        values = series.astype("float32").to_numpy()
+        x = torch.tensor(values, dtype=torch.float32, device=self.device).unsqueeze(0)
+
         with torch.no_grad():
-            preds = self.model(inputs, prediction_length=horizon).predictions
-        return pd.DataFrame(
-            preds.squeeze().cpu().numpy(),
-            columns=["forecast"]
+            # 1) Preferred: explicit .predict API
+            if hasattr(self.model, "predict"):
+                try:
+                    preds = self.model.predict(x, prediction_length=horizon)
+                    yhat = preds if isinstance(preds, torch.Tensor) else torch.tensor(preds)
+                    out = yhat.squeeze().detach().cpu().numpy()
+                    return pd.DataFrame({"forecast": out})
+                except Exception as e:
+                    raise RuntimeError(
+                        f"Granite model has a 'predict' method but it failed at runtime: {e}"
+                    )
+
+            # 2) Fallback: call forward with a 'prediction_length' kwarg if supported
+            try:
+                outputs = self.model(x, prediction_length=horizon)
+                # Try common attribute names
+                for k in ("predictions", "prediction", "logits", "output"):
+                    if hasattr(outputs, k):
+                        tensor = getattr(outputs, k)
+                        if isinstance(tensor, (tuple, list)):
+                            tensor = tensor[0]
+                        if not isinstance(tensor, torch.Tensor):
+                            tensor = torch.tensor(tensor)
+                        out = tensor.squeeze().detach().cpu().numpy()
+                        # If multi-dim, take last dimension as forecast
+                        if out.ndim > 1:
+                            out = out[-1] if out.shape[0] == horizon else out.reshape(-1)
+                        return pd.DataFrame({"forecast": out})
+                # If outputs is a raw tensor
+                if isinstance(outputs, torch.Tensor):
+                    out = outputs.squeeze().detach().cpu().numpy()
+                    if out.ndim > 1:
+                        out = out[-1] if out.shape[0] == horizon else out.reshape(-1)
+                    return pd.DataFrame({"forecast": out})
+            except TypeError:
+                # Some builds may not accept prediction_length at all
+                pass
+            except Exception as e:
+                raise RuntimeError(
+                    f"Calling the model forward for forecasting failed: {e}"
+                )
+
+        # If we get here, the installed combo doesn't expose an inference entrypoint we can use.
+        raise RuntimeError(
+            "The installed transformers/model combo does not expose a usable zero-shot "
+            "forecasting interface (no `.predict` and forward(...) didn't accept "
+            "`prediction_length`). Consider:\n"
+            "  • Upgrading transformers/torch versions\n"
+            "  • Using the 'granite-tsfm-public' PyPI if available in your region\n"
+            "  • Switching to a classic forecaster for now (e.g., ARIMA/XGBoost)\n"
         )