Add Chronos2 to time-series benchmarks

benchmark/timeseries/fev_bench/README.md

@@ -12,6 +12,15 @@ pip install . fev==0.10.0
 
 ## Run
 
+- **`Chronos2`**:
+
+  ```bash
+  python -m benchmark.timeseries.fev_bench.main --model chronos2
+  ```
+
+  `Chronos2` follows official evaluation with cross-learning across batches, and contexts truncated to 8,192 time steps.
+  Unlike the official `float32` evaluation, we use `float16` mixed precision for improved efficiency with negligible differences in aggregate benchmark accuracy.
+
 - **`TimesFM3`**:
 
   ```bash

benchmark/timeseries/fev_bench/main.py

@@ -18,7 +18,7 @@
 parser.add_argument(
     "--model",
     choices=tuple(MODEL_CONFIGS),
-    default="timesfm3",
+    default="chronos2",
     help="SDM model to benchmark.",
 )
 parser.add_argument(

benchmark/timeseries/fev_bench/predictor.py

@@ -147,6 +147,7 @@ def _fit_predict(self, task: fev.Task) -> list[datasets.DatasetDict]:
                                 y_context=TableTensor.from_tensor(y_i),
                                 x_query=TableTensor.from_tensor(x_query),
                                 num_estimators=self.config.num_estimators,
+                                cross_learning=self.config.cross_learning,
                             )
                         pred = out.numerical
                         pred = pred.view(y.size(0), task.horizon, size, -1)
@@ -155,7 +156,7 @@ def _fit_predict(self, task: fev.Task) -> list[datasets.DatasetDict]:
                     pred = torch.cat(preds, dim=-2)[..., : y.size(-1), :]
                     pred = pred.cpu().numpy()  # [B, R, C, Q]
                     levels = out.numerical_columns[: pred.shape[-1]]
-                    levels = [float(f"0.{q.split('q')[-1]}") for q in levels]
+                    levels = [int(q.split("q")[-1]) / 100 for q in levels]
                     wanted = [0.5, *task.quantile_levels]
                     index = [levels.index(q) for q in wanted]
                     values.append(pred[..., index])

benchmark/timeseries/gift_eval/README.md

@@ -13,6 +13,15 @@ pip install --no-deps "git+https://github.com/SalesforceAIResearch/gift-eval.git
 
 ## Run
 
+- **`Chronos2`**:
+
+  ```bash
+  python -m benchmark.timeseries.gift_eval.main --model chronos2
+  ```
+
+  `Chronos2` follows official evaluation with cross-learning across batches, and contexts truncated to 8,192 time steps.
+  Unlike the official `float32` evaluation, we use `float16` mixed precision for improved efficiency with negligible differences in aggregate benchmark accuracy.
+
 - **`TimesFM3`**:
 
   ```bash

benchmark/timeseries/gift_eval/main.py

@@ -8,14 +8,15 @@
 import os
 import warnings
 from functools import partial
+from itertools import islice
 from pathlib import Path
 from unittest.mock import patch
 
 import pandas as pd
 import torch
 from gift_eval.data import Dataset
 from gluonts.ev import metrics
-from gluonts.model import evaluate_model
+from gluonts.model import evaluate_forecasts
 from gluonts.time_feature import get_seasonality
 from huggingface_hub import snapshot_download
 from tqdm.auto import tqdm
@@ -33,7 +34,7 @@
 parser.add_argument(
     "--model",
     choices=tuple(MODEL_CONFIGS),
-    default="timesfm3",
+    default="chronos2",
     help="SDM model to benchmark.",
 )
 parser.add_argument(
@@ -95,12 +96,38 @@
 
         dataset = Dataset(name=name, term=term, to_univariate=False)
         predictor = SDMPredictor(model, config, dataset.prediction_length)
+
+        if config.cross_learning:  # Share information only within same window:
+            outs = []
+            with tqdm(total=len(dataset.test_data), desc=dataset_id) as pbar:
+                for window in range(dataset.test_data.windows):
+                    inputs = islice(
+                        dataset.test_data.input,
+                        window,
+                        None,
+                        dataset.test_data.windows,
+                    )
+                    windows = []
+                    for out in predictor.predict(inputs):
+                        windows.append(out)
+                        pbar.update()
+                    outs.append(windows)
+            outs = [out for windows in zip(*outs) for out in windows]
+        else:
+            outs = list(
+                tqdm(
+                    predictor.predict(dataset.test_data.input),
+                    total=len(dataset.test_data),
+                    desc=dataset_id,
+                )
+            )
+
         with patch(
             "gluonts.model.evaluation.tqdm",
-            partial(tqdm, total=len(dataset.test_data), desc=dataset_id),
+            partial(tqdm, disable=True),
         ):
-            result = evaluate_model(
-                predictor,
+            result = evaluate_forecasts(
+                outs,
                 test_data=dataset.test_data,
                 metrics=METRICS,
                 seasonality=get_seasonality(dataset.freq),

benchmark/timeseries/model.py

@@ -17,9 +17,19 @@ class ModelConfig:
     num_estimators: int = 1
     max_context_size: int | None = None
     max_variates: int | None = None
+    cross_learning: bool = False
 
 
 MODEL_CONFIGS = {
+    "chronos2": ModelConfig(
+        name="Chronos2",
+        factory=lambda device: sdm.models.Chronos2(device=device),
+        autocast_dtype=torch.float16,
+        num_estimators=1,
+        max_context_size=8_192,
+        max_variates=None,
+        cross_learning=True,
+    ),
     "timesfm3": ModelConfig(
         name="TimesFM3",
         factory=lambda device: sdm.models.TimesFM3(

benchmark/timeseries/predictor.py

@@ -107,6 +107,7 @@ def predict(
                             device=self.device,
                         ),
                         num_estimators=self.config.num_estimators,
+                        cross_learning=self.config.cross_learning,
                     )
                 pred = out.numerical
                 pred = pred.view(y.size(0), self.prediction_length, size, -1)
@@ -115,7 +116,7 @@ def predict(
             pred = torch.cat(preds, dim=-2)[..., : y.size(-1), :]
             pred = pred.permute(0, 3, 1, 2).cpu().numpy()  # [B, Q, R, C]
             levels = out.numerical_columns[: pred.shape[1]]
-            levels = [str(float(f"0.{q.split('q')[-1]}")) for q in levels]
+            levels = [str(int(q.split("q")[-1]) / 100) for q in levels]
 
             for entry, pred_i in zip(batch, pred):
                 if entry["target"].ndim == 1:

benchmark/timeseries/time_bench/README.md

@@ -13,6 +13,15 @@ pip install --no-deps "git+https://github.com/zqiao11/TIME.git@3ca5c41d71c76f3c7
 
 ## Run
 
+- **`Chronos2`**:
+
+  ```bash
+  python -m benchmark.timeseries.time_bench.main --model chronos2
+  ```
+
+  `Chronos2` follows official evaluation with contexts truncated to 8,192 time steps.
+  Unlike the official `float32` evaluation, we use `float16` mixed precision for improved efficiency with negligible differences in aggregate benchmark accuracy.
+
 - **`TimesFM3`**:
 
   ```bash

benchmark/timeseries/time_bench/main.py

@@ -7,6 +7,7 @@
 import os
 import warnings
 from contextlib import redirect_stdout
+from dataclasses import replace
 from pathlib import Path
 
 import numpy as np
@@ -29,7 +30,7 @@
 parser.add_argument(
     "--model",
     choices=tuple(MODEL_CONFIGS),
-    default="timesfm3",
+    default="chronos2",
     help="SDM model to benchmark.",
 )
 parser.add_argument(
@@ -41,7 +42,7 @@
 
 warnings.filterwarnings("ignore", message=r"'.*' is deprecated.*")
 
-config = MODEL_CONFIGS[args.model]
+config = replace(MODEL_CONFIGS[args.model], cross_learning=False)
 dataset_config = load_dataset_config()
 path = Path(__file__).parent.parent / "time_bench_out" / config.name
 
@@ -79,7 +80,9 @@
             )
         )
 
-        quantiles = np.stack([out.forecast_array for out in outs])
+        levels = [i / 10 for i in range(1, 10)]
+        indices = [outs[0].forecast_keys.index(str(q)) for q in levels]
+        quantiles = np.stack([out.forecast_array for out in outs])[:, indices]
         if quantiles.ndim == 4:
             quantiles = quantiles.swapaxes(-1, -2)  # [B, Q, C, R]
 
@@ -90,7 +93,7 @@
                 ds_config=f"{name}/{term}",
                 output_base_dir=str(path),
                 seasonality=get_seasonality(dataset.freq),
-                quantile_levels=[float(q) for q in outs[0].forecast_keys],
+                quantile_levels=levels,
             )
         with np.load(path / name / term / "metrics.npz") as metrics:
             mase = np.nanmean(metrics["MASE"])