Chronos 2

orion/pipelines/pretrained/chronos2/chronos2.json

@@ -0,0 +1,36 @@
+{
+    "primitives": [
+        "mlstars.custom.timeseries_preprocessing.time_segments_aggregate",
+        "sklearn.impute.SimpleImputer",
+        "mlstars.custom.timeseries_preprocessing.rolling_window_sequences",
+        "orion.primitives.chronos2.Chronos2",
+        "orion.primitives.timeseries_errors.regression_errors",
+        "orion.primitives.timeseries_anomalies.find_anomalies"
+    ],
+    "init_params": {
+        "mlstars.custom.timeseries_preprocessing.time_segments_aggregate#1": {
+            "time_column": "timestamp",
+            "interval": 21600,
+            "method": "mean"
+        },
+        "mlstars.custom.timeseries_preprocessing.rolling_window_sequences#1": {
+            "target_column": 0,
+            "window_size": 250
+        },
+        "orion.primitives.timeseries_anomalies.find_anomalies#1": {
+            "window_size_portion": 0.33,
+            "window_step_size_portion": 0.1,
+            "fixed_threshold": true
+        }
+    },
+    "input_names": {
+        "orion.primitives.timeseries_anomalies.find_anomalies#1": {
+            "index": "target_index"
+        }
+    },
+    "output_names": {
+        "orion.primitives.timeseries_anomalies.find_anomalies#1": {
+            "y": "anomalies"
+        }
+    }
+}

orion/pipelines/pretrained/chronos2/chronos2_artificialwithanomaly.json

@@ -0,0 +1,6 @@
+{
+    "mlstars.custom.timeseries_preprocessing.time_segments_aggregate#1": {
+        "time_column": "timestamp",
+        "interval": 600
+    }
+}

orion/pipelines/pretrained/chronos2/chronos2_msl.json

@@ -0,0 +1,6 @@
+{
+    "mlstars.custom.timeseries_preprocessing.time_segments_aggregate#1": {
+        "time_column": "timestamp",
+        "interval": 21600
+    }
+}

orion/pipelines/pretrained/chronos2/chronos2_realadexchange.json

@@ -0,0 +1,6 @@
+{
+    "mlstars.custom.timeseries_preprocessing.time_segments_aggregate#1": {
+        "time_column": "timestamp",
+        "interval": 3600
+    }
+}

orion/pipelines/pretrained/chronos2/chronos2_realawscloudwatch.json

@@ -0,0 +1,6 @@
+{
+    "mlstars.custom.timeseries_preprocessing.time_segments_aggregate#1": {
+        "time_column": "timestamp",
+        "interval": 600
+    }
+}

orion/pipelines/pretrained/chronos2/chronos2_realtraffic.json

@@ -0,0 +1,6 @@
+{
+    "mlstars.custom.timeseries_preprocessing.time_segments_aggregate#1": {
+        "time_column": "timestamp",
+        "interval": 600
+    }
+}

orion/pipelines/pretrained/chronos2/chronos2_realtweets.json

@@ -0,0 +1,6 @@
+{
+    "mlstars.custom.timeseries_preprocessing.time_segments_aggregate#1": {
+        "time_column": "timestamp",
+        "interval": 600
+    }
+}

orion/pipelines/pretrained/chronos2/chronos2_smap.json

@@ -0,0 +1,6 @@
+{
+    "mlstars.custom.timeseries_preprocessing.time_segments_aggregate#1": {
+        "time_column": "timestamp",
+        "interval": 21600
+    }
+}

orion/pipelines/pretrained/chronos2/chronos2_ucr.json

@@ -0,0 +1,6 @@
+{
+    "mlstars.custom.timeseries_preprocessing.time_segments_aggregate#1": {
+        "time_column": "timestamp",
+        "interval": 300
+    }
+}

orion/pipelines/pretrained/chronos2/chronos2_yahooa1.json

@@ -0,0 +1,6 @@
+{
+    "mlstars.custom.timeseries_preprocessing.time_segments_aggregate#1": {
+        "time_column": "timestamp",
+        "interval": 1
+    }
+}

orion/pipelines/pretrained/chronos2/chronos2_yahooa2.json

@@ -0,0 +1,6 @@
+{
+    "mlstars.custom.timeseries_preprocessing.time_segments_aggregate#1": {
+        "time_column": "timestamp",
+        "interval": 1
+    }
+}

orion/pipelines/pretrained/chronos2/chronos2_yahooa3.json

@@ -0,0 +1,6 @@
+{
+    "mlstars.custom.timeseries_preprocessing.time_segments_aggregate#1": {
+        "time_column": "timestamp",
+        "interval": 1
+    }
+}

orion/pipelines/pretrained/chronos2/chronos2_yahooa4.json

@@ -0,0 +1,6 @@
+{
+    "mlstars.custom.timeseries_preprocessing.time_segments_aggregate#1": {
+        "time_column": "timestamp",
+        "interval": 1
+    }
+}

orion/primitives/chronos2.py

@@ -0,0 +1,102 @@
+"""
+This primitive an implementation of Amazon's Chronos2 model for timeseries forecasting.
+
+The model implementation can be found at
+https://huggingface.co/amazon/chronos-2
+
+Note: This primitive assumes that Chronos2 doesn't care about specific timestamps
+of the data. We fill in the timestamps with a linear sequence of timestamps in order
+for the model to work.
+"""
+
+import numpy as np
+import pandas as pd
+import torch
+
+from chronos import Chronos2Pipeline
+
+
+class Chronos2:
+    """Chronos2 model for timeseries forecasting.
+
+    Args:
+        pred_len (int):
+            Prediction horizon length. Default to 1.
+        repo_id (str):
+            Directory of the model checkpoint. Default to "amazon/chronos-2"
+        batch_size(int):
+            Size of one batch. Default to 32.
+        target (int):
+            Index of target column in multivariate case. Default to 0.
+        start_time (datetime):
+            Start time of the timeseries. Default to Jan 1, 2020 00:00:00.
+        time_interval (int):
+            Time interval between two samples in seconds. Default to 600.
+    """
+
+    def __init__(self,
+                 pred_len=1,
+                 repo_id="amazon/chronos-2",
+                 batch_size=32,
+                 target=0,
+                 start_time=pd.to_datetime("2000-01-01 00:00:00"),
+                 time_interval=600):
+
+        self.pred_len = pred_len
+        self.batch_size = batch_size
+        self.target = f"{target}"
+        self.start_time = start_time
+        self.time_interval = pd.Timedelta(seconds=time_interval)
+
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model = Chronos2Pipeline.from_pretrained(repo_id, device_map=device)
+
+    def predict(self, X, force=False):
+        """Forecasting timeseries
+
+        Args:
+            X (ndarray):
+                input timeseries with shape (n_windows, window_size, n_features).
+        Return:
+            ndarray:
+                forecasted timeseries.
+        """
+        n_windows, window_size, n_features = X.shape
+
+        outs = []
+
+        for i in range(0, n_windows, self.batch_size):
+            x_batch = self.convert_to_df(X[i:i + self.batch_size], start_batch_at=i)
+            y_batch = self.model.predict_df(
+                df=x_batch,
+                prediction_length=self.pred_len,
+                quantile_levels=[0.5],
+                id_column="item_id",
+                timestamp_column="timestamp",
+                target=self.target,
+            )
+
+            y_batch = y_batch.sort_values(["item_id", "timestamp"])
+            preds = np.stack(
+                y_batch.groupby("item_id", sort=False)["predictions"]
+                .apply(lambda s: s.to_numpy())
+                .to_list()
+            )
+            outs.append(preds)
+
+        return np.concatenate(outs, axis=0)
+
+    def convert_to_df(self, x_batch, start_batch_at=0):
+        n_windows_in_batch, window_size, n_features = x_batch.shape
+
+        rows = []
+        for window in range(n_windows_in_batch):
+            for data_entry in range(window_size):
+                rows.append({
+                    "timestamp": self.start_time + self.time_interval * data_entry,
+                    "item_id": f"window_{start_batch_at + window}",
+                    **{f"{i}": x_batch[window, data_entry, i] for i in range(n_features)}
+                })
+
+        rows = pd.DataFrame(rows)
+        return rows

orion/primitives/jsons/orion.primitives.chronos2.Chronos2.json

@@ -0,0 +1,58 @@
+{
+    "name": "orion.primitives.chronos2.Chronos2",
+    "contributors": [
+        "Allen Baranov <baranov@mit.edu>"
+    ],
+    "documentation": "https://huggingface.co/amazon/chronos-2",
+    "description": "Amazon Chronos2 model for timeseries forecasting",
+    "classifiers": {
+        "type": "estimator",
+        "subtype": "regressor"
+    },
+    "modalities": [],
+    "primitive": "orion.primitives.chronos2.Chronos2",
+    "produce": {
+        "method": "predict",
+        "args": [
+            {
+                "name": "X",
+                "type": "ndarray"
+            },
+            {
+                "name": "force",
+                "type": "bool",
+                "default": false
+            }
+        ],
+        "output": [
+            {
+                "name": "y_hat",
+                "type": "ndarray"
+            }
+        ]
+    },
+    "hyperparameters": {
+        "fixed": {
+            "pred_len": {
+                "type": "int",
+                "default": 1
+            },
+            "repo_id": {
+                "type": "str",
+                "default": "amazon/chronos-2"
+            },
+            "batch_size": {
+                "type": "int",
+                "default": 32
+            },
+            "target": {
+                "type": "int",
+                "default": 0
+            },
+            "time_interval": {
+                "type": "int",
+                "default": 600
+            }
+        }
+    }
+}

setup.py

@@ -48,6 +48,9 @@
     "timesfm[torch]>=1.2.0,<1.5;python_version>='3.11'",
     "jax;python_version>='3.11'",
 
+    #chronos2
+    'chronos-forecasting>=2.2.0,<2.3.0',
+
     'wrapt>=1.14,<1.15',
 ]