diff --git a/README.md b/README.md index a75587c..f949d8b 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,84 @@ # TimesFM -PLACEHOLDER +TimesFM (Time Series Foundation Model) is a pretrained time-series foundation +model developed by Google Research for time-series forecasting. + +* Paper: + [A decoder-only foundation model for time-series forecasting](https://arxiv.org/abs/2310.10688), + ICML 2024. +* All checkpoints: + [TimesFM Hugging Face Collection](https://huggingface.co/collections/google/timesfm-release-66e4be5fdb56e960c1e482a6). +* [Google Research blog](https://research.google/blog/a-decoder-only-foundation-model-for-time-series-forecasting/). +* [TimesFM in BigQuery](https://cloud.google.com/bigquery/docs/timesfm-model): + an official Google product. + +This open version is not an officially supported Google product. + +**Latest Model Version:** TimesFM 2.5 + +**Archived Model Versions:** + +- 1.0 and 2.0: relevant code archived in the sub directory `v1`. You can `pip + install timesfm==1.3.0` to install an older version of this package to load + them. + +## Update - Sept. 15, 2025 + +TimesFM 2.5 is out! + +Comparing to TimesFM 2.0, this new 2.5 model: + +- uses 200M parameters, down from 500M. +- supports up to 16k context length, up from 2048. +- supports continuous quantile forecast up to 1k horizon via an optional 30M + quantile head. +- gets rid of the `frequency` indicator. +- has a couple of new forecasting flags. + +Along with the model upgrade we have also upgraded the inference API. This repo +will be under construction over the next few weeks to + +1. add support for an upcoming Flax version of the model (faster inference). +2. add back covariate support. +3. populate more docstrings, docs and notebook. + +### Install + +TODO(siriuz42): Package timesfm==2.0.0 and upload to PyPI . + +Run + +```shell +git clone https://github.com/google-research/timesfm.git +cd timesfm +pip install -e . +``` + +### Code Example + +```python +import numpy as np +import timesfm +model = timesfm.TimesFM_2p5_200M_torch() +model.load_checkpoint() +model.compile( + timesfm.ForecastConfig( + max_context=1024, + max_horizon=256, + normalize_inputs=True, + use_continuous_quantile_head=True, + force_flip_invariance=True, + infer_is_positive=True, + fix_quantile_crossing=True, + ) +) +point_forecast, quantile_forecast = model.forecast( + horizon=12, + inputs=[ + np.linspace(0, 1, 100), + np.sin(np.linspace(0, 20, 67)), + ], # Two dummy inputs +) +point_forecast.shape # (2, 12) +quantile_forecast.shape # (2, 12, 10): mean, then 10th to 90th quantiles. +``` diff --git a/src/timesfm/__init__.py b/src/timesfm/__init__.py index cf83c93..85177a2 100644 --- a/src/timesfm/__init__.py +++ b/src/timesfm/__init__.py @@ -14,6 +14,7 @@ """TimesFM API.""" +from .configs import ForecastConfig from .timesfm_2p5 import timesfm_2p5_torch TimesFM_2p5_200M_torch = timesfm_2p5_torch.TimesFM_2p5_200M_torch diff --git a/src/timesfm/configs.py b/src/timesfm/configs.py index e02672d..f5ba9c5 100644 --- a/src/timesfm/configs.py +++ b/src/timesfm/configs.py @@ -20,7 +20,33 @@ from typing import Literal @dataclasses.dataclass(frozen=False) class ForecastConfig: - """Options for forecasting.""" + """Options for forecasting. + + Attributes: + max_context: The maximum context length. This is used by the complied decode + function at inference time during batched inference. Any input time series + with length less than max_context will be padded with zeros, and with + length greater than max_context will be truncated. + max_horizon: The maximum horizon length. This is used by the complied decode + function at inference time during batched inference. The compiled cached + decoding function will by default forecast till max_horizon. + normalize_inputs: Whether to normalize the inputs. This is useful when the + raw inputs are of extremely large or small magnitudes which may result in + numerical issues. + window_size: The window size for decomposed forecasting. + TODO(siriuz42):implement it. + per_core_batch_size: The batch size per core. Used at inference time during + batched inference when multiple GPU / TPU devices are used. + use_continuous_quantile_head: Whether to use a separate continuous quantile + head to avoid quantile collapsing. + force_flip_invariance: Whether to force flip invariance. TimesFM guarantees + that TimesFM(aX + b) = a * TimesFM(x) + b for a >= 0 by default. This flag + extends it to a < 0 as well. + infer_is_positive: Whether to guarantee nonnegativity of the output if the + input is nonnegative. + fix_quantile_crossing: Whether to fix quantile crossing. + return_backcast: Whether to return backcast. + """ max_context: int = 0 max_horizon: int = 0 diff --git a/src/timesfm/timesfm_2p5/timesfm_2p5_base.py b/src/timesfm/timesfm_2p5/timesfm_2p5_base.py index 1eaea07..ba5902e 100644 --- a/src/timesfm/timesfm_2p5/timesfm_2p5_base.py +++ b/src/timesfm/timesfm_2p5/timesfm_2p5_base.py @@ -126,7 +126,13 @@ class TimesFM_2p5_200M_Definition: class TimesFM_2p5: - """Abstract base class for TimesFM models.""" + """Abstract base class for TimesFM models. + + Attributes: + forecast_config: Configuration for forecasting flags. + compiled_decode: Compiled decode function. + global_batch_size: Global batch size. + """ forecast_config: ForecastConfig | None = None compiled_decode: Callable[..., Any] | None = None diff --git a/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py b/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py index 6edecca..dfb21aa 100644 --- a/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py +++ b/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py @@ -229,8 +229,22 @@ class TimesFM_2p5_200M_torch_module(nn.Module): return renormed_outputs, renormed_quantile_spread, ar_renormed_outputs - def forecast_naive(self, horizon: int, inputs: Sequence[np.ndarray]): - """Forecasts the time series.""" + def forecast_naive( + self, horizon: int, inputs: Sequence[np.ndarray] + ) -> list[np.ndarray]: + """Forecasts the time series. + + This is a naive implementation for debugging purposes. No forecasting + flags are used here. Forecasting quality can be subpar. + + Args: + horizon: The number of time points to forecast. + inputs: A sequence of numpy arrays, each representing a time series to + query forecast for. + + Returns: + A list of numpy arrays of forecasts. + """ outputs = [] for each_input in inputs: input_t = torch.tensor(each_input, dtype=torch.float32) @@ -265,13 +279,14 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5): *, path: str | None = None, hf_repo_id: str | None = "google/timesfm-2.5-200m-pytorch", - ): + ) -> None: """Loads a PyTorch safetensors TimesFM model. Args: path: Path to a local checkpoint. If not provided, will try to download from the default Hugging Face repo. - hf_repo_id: Use another Hugging Face repo ID. + hf_repo_id: If provided, will download from the specified Hugging Face + repo instead. """ if path: pass @@ -287,7 +302,16 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5): raise ValueError("Either path or hf_repo_id must be provided.") self.model.load_checkpoint(path) - def compile(self, forecast_config: configs.ForecastConfig, **kwargs): + def compile(self, forecast_config: configs.ForecastConfig, **kwargs) -> None: + """Attempts to compile the model for fast decoding. + + See configs.ForecastConfig for more details on the supported flags. + + Args: + forecast_config: Configuration for forecasting flags. + **kwargs: Additional keyword arguments to pass to model.compile(). + """ + if kwargs.get("backend", None) is not None: self.model.compile(**kwargs) self.global_batch_size = ( @@ -354,13 +378,10 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5): pf_outputs, quantile_spreads, ar_outputs = self.model.decode( forecast_config.max_horizon, inputs, masks ) - full_forecast = torch.cat( - [ - pf_outputs[:, -1, ...], - ar_outputs.reshape(batch_size, -1, self.model.q), - ], - dim=1, - ) + to_cat = [pf_outputs[:, -1, ...]] + if ar_outputs is not None: + to_cat.append(ar_outputs.reshape(batch_size, -1, self.model.q)) + full_forecast = torch.cat(to_cat, dim=1) flip_quantile_fn = lambda x: torch.cat( [x[..., :1], torch.flip(x[..., 1:], dims=(-1,))], dim=-1 @@ -372,13 +393,12 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5): ) flipped_quantile_spreads = flip_quantile_fn(flipped_quantile_spreads) flipped_pf_outputs = flip_quantile_fn(flipped_pf_outputs) - flipped_full_forecast = torch.cat( - [ - flipped_pf_outputs[:, -1, ...], - flipped_ar_outputs.reshape(batch_size, -1, self.model.q), - ], - dim=1, - ) + to_cat = [flipped_pf_outputs[:, -1, ...]] + if flipped_ar_outputs is not None: + to_cat.append( + flipped_ar_outputs.reshape(batch_size, -1, self.model.q) + ) + flipped_full_forecast = torch.cat(to_cat, dim=1) quantile_spreads = (quantile_spreads - flipped_quantile_spreads) / 2 pf_outputs = (pf_outputs - flipped_pf_outputs) / 2 full_forecast = (full_forecast - flipped_full_forecast) / 2