Create ts_forecast_tool.py

tools/ts_forecast_tool.py

@@ -0,0 +1,90 @@
+# space/tools/ts_forecast_tool.py
+import numpy as np
+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
+
+
+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.
+    """
+    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]:
+        """
+        Zero-shot forecast: no fine-tuning, just apply pre-trained model.
+        Returns dict with "forecast" (future horizon) and "context" (last context window).
+        """
+        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)
+        }