File size: 3,600 Bytes
852fd6f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
# 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)
}
|