Update tools/ts_forecast_tool.py

tools/ts_forecast_tool.py

@@ -1,90 +1,30 @@
 # space/tools/ts_forecast_tool.py
-import numpy as np
+import torch
 import pandas as pd
-from typing import Dict, Optional
-
-from utils.tracing import Tracer
-from utils.config import AppConfig
-
-# Granite TTM imports (from tsfm_public)
-from tsfm_public.models.tinytimemixer.modeling_tinytimemixer import TinyTimeMixerForPrediction
-from tsfm_public.toolkit.config import TSPPConfig
-from tsfm_public.toolkit.dataset import TimeSeriesDataset
-from tsfm_public.toolkit.trainer import Trainer
-
+from transformers import AutoModelForTimeSeriesForecasting, AutoTokenizer
 
 class TimeseriesForecastTool:
     """
-    Wraps IBM Granite TinyTimeMixer (TTM) for multivariate forecasting.
-    Expects a wide dataframe with:
-      - 'timestamp' (datetime64[ns])
-      - one or more numeric series columns to forecast (targets)
-      - optional control/exogenous columns (known-in-future features)
-    You must provide context_length and forecast_length to match a TTM variant.
+    Lightweight wrapper around ibm-granite/granite-timeseries-ttm-r1
+    using the Transformers interface.
     """
-    def __init__(self, cfg: AppConfig, tracer: Tracer,
-                 hf_model_id: str = "ibm-granite/granite-timeseries-ttm-r1",
-                 revision: str = "main",
-                 context_length: int = 512,
-                 forecast_length: int = 96,
-                 target_cols: Optional[list] = None,
-                 control_cols: Optional[list] = None):
-        self.cfg = cfg
-        self.tracer = tracer
-
-        self.context_length = context_length
-        self.forecast_length = forecast_length
-        self.target_cols = target_cols or []
-        self.control_cols = control_cols or []
-
-        # Build TSPP config
-        self.tspp_config = TSPPConfig(
-            context_length=context_length,
-            prediction_length=forecast_length,
-            target_cols=self.target_cols,
-            known_cov_cols=self.control_cols,  # known-in-future exogenous
-            time_col="timestamp",
-            freq=None  # inferred; you can set "H" or "T" if you know it
-        )
-
-        # Load model from HF (r1; try r2 for newer variants if needed)
-        self.model = TinyTimeMixerForPrediction.from_pretrained(
-            hf_model_id, revision=revision
-        )
-
-    def _build_dataset(self, df: pd.DataFrame) -> TimeSeriesDataset:
-        # Minimal build: single item dataset from dataframe (you can batch multiple series)
-        item = df.sort_values("timestamp").reset_index(drop=True)
-        return TimeSeriesDataset.from_pandas(
-            item,
-            tspp_config=self.tspp_config
-        )
-
-    def zeroshot_forecast(self, df: pd.DataFrame) -> Dict[str, pd.DataFrame]:
+    def __init__(self,
+                 model_id="ibm-granite/granite-timeseries-ttm-r1",
+                 device=None):
+        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):
         """
-        Zero-shot forecast: no fine-tuning, just apply pre-trained model.
-        Returns dict with "forecast" (future horizon) and "context" (last context window).
+        series: pd.Series indexed by datetime
+        horizon: forecast steps
         """
-        dset = self._build_dataset(df)
-        trainer = Trainer(model=self.model)
-        out = trainer.evaluate(dset)  # Granite API uses 'evaluate' for zeroshot
-
-        # Convert to a user-friendly dataframe
-        # out contains predictions for target_cols for next forecast_length steps
-        preds = out["predictions"]  # dict: {col: np.ndarray [forecast_length]}
-        horizon_idx = pd.RangeIndex(self.forecast_length, name="step")
-        forecast_df = pd.DataFrame(preds, index=horizon_idx).reset_index(drop=True)
-
-        try:
-            self.tracer.trace_event("ts_forecast", {
-                "targets": self.target_cols,
-                "ctx": self.context_length,
-                "h": self.forecast_length
-            })
-        except Exception:
-            pass
-
-        return {
-            "forecast": forecast_df,
-            "context": df.tail(self.context_length).reset_index(drop=True)
-        }
+        values = series.values.astype("float32")
+        inputs = torch.tensor(values, dtype=torch.float32).unsqueeze(0).to(self.device)
+        with torch.no_grad():
+            preds = self.model(inputs, prediction_length=horizon).predictions
+        return pd.DataFrame(
+            preds.squeeze().cpu().numpy(),
+            columns=["forecast"]
+        )