chronosx

2026-05-23 09:39:35 +00:00 · 2025-09-05 11:38:58 +00:00 · 2025-09-05 11:38:58 +00:00 · 563cb70f12
commit 563cb70f12
parent ccb656cc2b
100 changed files with 3243 additions and 10704 deletions
--- a/.github/ISSUE_TEMPLATE/bug_report.md
+++ b/.github/ISSUE_TEMPLATE/bug_report.md
@ -1,33 +0,0 @@
---
-name: Bug report
-about: Create a report to help us reproduce and fix a bug
-title: "[BUG]"
-labels: ['bug']
-assignees: ''
-
---
-
-**Bug report checklist**
-<!-- Please ensure at least one of the following to help the developers troubleshoot the problem:  -->
- [ ] I provided code that demonstrates a minimal reproducible example. <!-- Ideal, especially via source install --> 
- [ ] I confirmed bug exists on the latest mainline of Chronos via source install. <!-- Preferred --> 
-
-**Describe the bug**
-<!-- A clear and concise description of what the bug is. -->
-
-**Expected behavior**
-<!-- A clear and concise description of what you expected to happen. -->
-
-**To reproduce**
-<!-- A minimal script to reproduce the issue. Links to Colab notebooks or similar tools are encouraged.  
-If the code is too long, feel free to put it in a public gist and link it in the issue: https://gist.github.com.  
-In short, we are going to copy-paste your code to run it and we expect to get the same result as you. -->
-
-**Environment description**
-Operating system:
-Python version:
-CUDA version:
-PyTorch version: 
-HuggingFace transformers version:
-HuggingFace accelerate version:
-
--- a/.github/ISSUE_TEMPLATE/config.yml
+++ b/.github/ISSUE_TEMPLATE/config.yml
@ -1,8 +0,0 @@
-blank_issues_enabled: false
-contact_links:
-  - name: Frequently asked questions
-    url: https://github.com/amazon-science/chronos-forecasting/issues?q=is%3Aissue+label%3AFAQ
-    about: Check the frequently asked questions before opening a new one
-  - name: Discussions
-    url: https://github.com/amazon-science/chronos-forecasting/discussions/new
-    about: Use this to ask questions and start discussions
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@ -1,52 +0,0 @@
-name: CI
-
-on:
-  push:
-    branches: ["main"] # Run only on main branch
-  pull_request:
-    branches: ["**"] # Run on any branch
-  schedule:
-    - cron: "0 8 * * *" # Run at 8 AM UTC
-
-jobs:
-  type-check:
-    strategy:
-      max-parallel: 4
-      fail-fast: false
-      matrix:
-        python-version: ["3.11"]
-        platform: [ubuntu-latest]
-
-    runs-on: ${{ matrix.platform }}
-
-    steps:
-      - uses: actions/checkout@v4
-      - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v5
-        with:
-          python-version: ${{ matrix.python-version }}
-      - name: Install dependencies
-        run: pip install ".[typecheck]" -f https://download.pytorch.org/whl/cpu/torch_stable.html
-      - name: Type checks with mypy
-        run: mypy src test
-
-  test:
-    strategy:
-      max-parallel: 4
-      fail-fast: false
-      matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12"]
-        platform: [ubuntu-latest]
-
-    runs-on: ${{ matrix.platform }}
-
-    steps:
-      - uses: actions/checkout@v4
-      - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v5
-        with:
-          python-version: ${{ matrix.python-version }}
-      - name: Install dependencies
-        run: pip install ".[test]" -f https://download.pytorch.org/whl/cpu/torch_stable.html
-      - name: Test with pytest
-        run: pytest
--- a/.github/workflows/eval-model.yml
+++ b/.github/workflows/eval-model.yml
@ -1,37 +0,0 @@
-# Evaluates Chronos-Bolt (Small) model on selected datasets
-name: Evaluate
-
-on:
-  # Runs only with read privilages for the GITHUB_TOKEN
-  pull_request:
-    branches: ["main"] # Run on PRs to main branch
-    types:
-      - opened      # When a PR is created
-      - reopened    # When a closed PR is reopened
-      - synchronize # When new commits are pushed to the PR
-      - labeled     # When a label is added to the PR
-
-jobs:
-  evaluate-and-print:
-    if: contains(github.event.pull_request.labels.*.name, 'run-eval')  # Only run if 'run-eval' label is added
-    runs-on: ubuntu-latest
-    env:
-      RESULTS_CSV: "eval-ci-metrics-${{ github.event.pull_request.number }}.csv"
-
-    steps:
-      - name: Checkout Repository
-        uses: actions/checkout@v4
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: '3.11'
-
-      - name: Install Dependencies
-        run: pip install ".[evaluation]" -f https://download.pytorch.org/whl/cpu/torch_stable.html
-
-      - name: Run Eval Script
-        run: python scripts/evaluation/evaluate.py ci/evaluate/backtest_config.yaml $RESULTS_CSV --chronos-model-id=amazon/chronos-bolt-small --device=cpu --torch-dtype=float32
-      
-      - name: Print CSV
-        run: cat $RESULTS_CSV
--- a/.github/workflows/publish-to-pypi.yml
+++ b/.github/workflows/publish-to-pypi.yml
@ -1,27 +0,0 @@
-name: Publish Python Package to PyPi
-
-on:
-  release:
-    types: [published]
-
-jobs:
-  deploy-to-pypi:
-    runs-on: ubuntu-latest
-    environment: release
-    permissions:
-      id-token: write
-    steps:
-    - uses: actions/checkout@v4
-    - name: Set up Python
-      uses: actions/setup-python@v5
-      with:
-        python-version: '3.11'
-    - name: Install dependencies
-      run: |
-        python -m pip install -U pip
-        python -m pip install setuptools wheel build
-    - name: Build package
-      run: |
-        python -m build
-    - name: Publish to PyPi
-      uses: pypa/gh-action-pypi-publish@release/v1
--- a/CITATION.cff
+++ b/CITATION.cff
@ -1,78 +1,62 @@
 cff-version: 1.2.0
-title: "Chronos: Learning the Language of Time Series"
+title: "ChronosX: Adapting Pretrained Time Series Models with Exogenous Variables"
 message: "If you find Chronos models useful for your research, please consider citing the associated paper."
 authors:
+  - family-names: Arango
+    given-names: Sebastian Pineda
+  - family-names: Mercado
+    given-names: Pedro
+  - family-names: Kapoor
+    given-names: Shubham
  - family-names: Ansari
    given-names: Abdul Fatir
  - family-names: Stella
    given-names: Lorenzo
-  - family-names: Turkmen
-    given-names: Caner
-  - family-names: Zhang
-    given-names: Xiyuan
-  - family-names: Mercado
-    given-names: Pedro
  - family-names: Shen
    given-names: Huibin
+  - family-names: Senetaire
+    given-names: Hugo
+  - family-names: Turkmen
+    given-names: Caner
  - family-names: Shchur
    given-names: Oleksandr
-  - family-names: Rangapuram
-    given-names: Syama Syndar
-  - family-names: Arango
-    given-names: Sebastian Pineda
-  - family-names: Kapoor
-    given-names: Shubham
-  - family-names: Zschiegner
-    given-names: Jasper
  - family-names: Maddix
    given-names: Danielle C.
-  - family-names: Mahoney
-    given-names: Michael W.
-  - family-names: Torkkola
-    given-names: Kari
-  - family-names: Wilson
-    given-names: Andrew Gordon
  - family-names: Bohlke-Schneider
    given-names: Michael
  - family-names: Wang
    given-names: Yuyang
+  - family-names: Rangapuram
+    given-names: Syama Syndar
 preferred-citation:
  type: article
  authors:
-    - family-names: Ansari
-      given-names: Abdul Fatir
-    - family-names: Stella
-      given-names: Lorenzo
-    - family-names: Turkmen
-      given-names: Caner
-    - family-names: Zhang
-      given-names: Xiyuan
-    - family-names: Mercado
-      given-names: Pedro
-    - family-names: Shen
-      given-names: Huibin
-    - family-names: Shchur
-      given-names: Oleksandr
-    - family-names: Rangapuram
-      given-names: Syama Syndar
-    - family-names: Arango
-      given-names: Sebastian Pineda
-    - family-names: Kapoor
-      given-names: Shubham
-    - family-names: Zschiegner
-      given-names: Jasper
-    - family-names: Maddix
-      given-names: Danielle C.
-    - family-names: Mahoney
-      given-names: Michael W.
-    - family-names: Torkkola
-      given-names: Kari
-    - family-names: Wilson
-      given-names: Andrew Gordon
-    - family-names: Bohlke-Schneider
-      given-names: Michael
-    - family-names: Wang
-      given-names: Yuyang
-  title: "Chronos: Learning the Language of Time Series"
-  journal: "arXiv preprint arXiv:2403.07815"
-  year: 2024
+  - family-names: Arango
+    given-names: Sebastian Pineda
+  - family-names: Mercado
+    given-names: Pedro
+  - family-names: Kapoor
+    given-names: Shubham
+  - family-names: Ansari
+    given-names: Abdul Fatir
+  - family-names: Stella
+    given-names: Lorenzo
+  - family-names: Shen
+    given-names: Huibin
+  - family-names: Senetaire
+    given-names: Hugo
+  - family-names: Turkmen
+    given-names: Caner
+  - family-names: Shchur
+    given-names: Oleksandr
+  - family-names: Maddix
+    given-names: Danielle C.
+  - family-names: Bohlke-Schneider
+    given-names: Michael
+  - family-names: Wang
+    given-names: Yuyang
+  - family-names: Rangapuram
+    given-names: Syama Syndar
+  title: "ChronosX: Adapting Pretrained Time Series Models with Exogenous Variables"
+  journal: "arXiv preprint arXiv:2503.12107"
+  year: 2025
--- a/CODE_OF_CONDUCT.md
+++ b/CODE_OF_CONDUCT.md
@ -1,4 +1,4 @@
 ## Code of Conduct
 This project has adopted the [Amazon Open Source Code of Conduct](https://aws.github.io/code-of-conduct).
 For more information see the [Code of Conduct FAQ](https://aws.github.io/code-of-conduct-faq) or contact
-opensource-codeofconduct@amazon.com with any additional questions or comments.
+opensource-codeofconduct@amazon.com with any additional questions or comments.
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@ -56,4 +56,4 @@ If you discover a potential security issue in this project we ask that you notif

 ## Licensing

-See the [LICENSE](LICENSE) file for our project's licensing. We will ask you to confirm the licensing of your contribution.
+See the [LICENSE](LICENSE) file for our project's licensing. We will ask you to confirm the licensing of your contribution.
--- a/2
+++ b/2
@ -172,4 +172,4 @@
      of any other Contributor, and only if You agree to indemnify,
      defend, and hold each Contributor harmless for any liability
      incurred by, or claims asserted against, such Contributor by reason
-      of your accepting any such warranty or additional liability.
+      of your accepting any such warranty or additional liability.
--- a/2
+++ b/2
@ -1 +1 @@
-Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
--- a/README.md
+++ b/README.md
@ -1,6 +1,124 @@
-<div align="center">
-
 # ChronosX: Adapting Pretrained Time Series Models with Exogenous Variables
-## 🚀 Code soon to be released 🚀 

-</div>
+[![preprint](https://img.shields.io/static/v1?label=arXiv&message=2503.12107&color=B31B1B&logo=arXiv)](https://arxiv.org/abs/2503.12107)
+[![License: MIT](https://img.shields.io/badge/License-Apache--2.0-green.svg)](https://opensource.org/licenses/Apache-2.0)
+
+This repository provides the forecasting model ChronosX introduced in the paper
+[ChronosX: Adapting Pretrained Time Series Models with Exogenous Variables](https://arxiv.org/abs/2503.12107).
+
+## ✨ Introduction
+ChronosX introduces a new method to incorporate covariates into pretrained time series forecasting models. ChronosX incorporates covariate information into pretrained forecasting models through modular blocks that inject past and future covariate information, without necessarily modifying the pretrained model in consideration.
+
+For further details on ChronosX please refer to the paper [ChronosX: Adapting Pretrained Time Series Models with Exogenous Variables](https://arxiv.org/abs/2503.12107).
+
+## 📈 Usage
+
+To perform inference with Chronos or Chronos-Bolt models, the easiest way is to install this package through `pip`:
+
+```sh
+pip install git+https://github.com/amazon-science/chronos-forecasting.git@chronosx
+```
+
+### Forecasting
+
+A minimal example showing how to perform forecasting using ChronosX:
+
+```python
+import numpy as np
+import yaml
+
+from chronosx.chronosx import ChronosXPipeline
+from chronosx.utils import ChronosDataset, load_and_split_dataset
+from gluonts.ev.metrics import MASE, MeanWeightedSumQuantileLoss
+from gluonts.model.evaluation import evaluate_forecasts
+from pathlib import Path
+
+
+config_path = "./scripts/experiments/configs/datasets.yaml"
+output_dir = Path(f"./output/finetune")
+output_dir.mkdir(exist_ok=True, parents=True)
+
+with open(config_path) as fp:
+    backtest_configs = yaml.safe_load(fp)
+
+dataset_config = backtest_configs[0]
+prediction_length = dataset_config["prediction_length"]
+num_covariates = 2 * len(dataset_config["covariates_fields"])
+
+
+# Load Chronos
+pipeline = ChronosXPipeline(
+    pretrained_model_name_or_path="amazon/chronos-t5-small",
+    prediction_length=prediction_length,
+    num_covariates=num_covariates,
+)
+
+# load Dataset
+train_dataset, test_dataset = load_and_split_dataset(backtest_config=dataset_config)
+quantized_train_dataset = ChronosDataset(
+    datasets=[train_dataset],
+    probabilities=[1.0],
+    tokenizer=pipeline.tokenizer,
+    prediction_length=prediction_length,
+    mode="training",
+).shuffle()
+
+# fine tune model
+_, save_model_path = pipeline.finetune(
+    output_dir,
+    quantized_train_dataset,
+    skip_pretrained_validation=True,
+)
+
+# Evaluate fine tuned model
+pipeline = ChronosXPipeline(
+    prediction_length=prediction_length,
+    num_covariates=num_covariates,
+    pretrained_model_name_or_path=output_dir / "final-checkpoint",
+)
+
+pipeline.chronosx.eval()
+forecasts = pipeline.generate_forecasts(test_dataset.input)
+
+metrics = (
+    evaluate_forecasts(
+        forecasts,
+        test_data=test_dataset,
+        metrics=[
+            MASE(),
+            MeanWeightedSumQuantileLoss(np.arange(0.05, 1, 0.05).round(2).tolist()),
+        ],
+    )
+    .reset_index(drop=True)
+    .to_dict(orient="records")
+)
+
+print(metrics)
+```
+
+### Experiments
+
+Scripts for experiments can be found in [this folder](./scripts/).
+
+## 📝 Citation
+
+If you find ChronosX useful for your research, please consider citing the associated [paper](https://arxiv.org/abs/2503.12107):
+
+```
+@inproceedings{
+arango2025chronosx,
+title={ChronosX: Adapting Pretrained Time Series Models with Exogenous Variables},
+author={Sebastian Pineda Arango and Pedro Mercado and Shubham Kapoor and Abdul Fatir Ansari and Lorenzo Stella and Huibin Shen and Hugo Henri Joseph Senetaire and Ali Caner Turkmen and Oleksandr Shchur and Danielle C. Maddix and Bernie Wang and Michael Bohlke-Schneider and Syama Sundar Rangapuram},
+booktitle={The 28th International Conference on Artificial Intelligence and Statistics},
+year={2025},
+url={https://openreview.net/forum?id=f4nWNn0RjV}
+}
+```
+
+## 🛡️ Security
+
+See [CONTRIBUTING](CONTRIBUTING.md#security-issue-notifications) for more information.
+
+## 📃 License
+
+This project is licensed under the Apache-2.0 License.
--- a/ci/evaluate/backtest_config.yaml
+++ b/ci/evaluate/backtest_config.yaml
@ -1,37 +0,0 @@
-# From In-domain
- name: taxi_30min # 30 min
-  hf_repo: autogluon/chronos_datasets
-  offset: -48
-  prediction_length: 48
-  num_rolls: 1
-# From Zero-shot
- name: ETTh # Hourly
-  hf_repo: autogluon/chronos_datasets_extra
-  offset: -24
-  prediction_length: 24
-  num_rolls: 1
- name: monash_covid_deaths # Daily
-  hf_repo: autogluon/chronos_datasets
-  offset: -30
-  prediction_length: 30
-  num_rolls: 1
- name: monash_nn5_weekly # Weekly
-  hf_repo: autogluon/chronos_datasets
-  offset: -8
-  prediction_length: 8
-  num_rolls: 1
- name: monash_fred_md # Monthly
-  hf_repo: autogluon/chronos_datasets
-  offset: -12
-  prediction_length: 12
-  num_rolls: 1
- name: monash_m3_quarterly # Quarterly
-  hf_repo: autogluon/chronos_datasets
-  offset: -8
-  prediction_length: 8
-  num_rolls: 1
- name: monash_tourism_yearly # Yearly
-  hf_repo: autogluon/chronos_datasets
-  offset: -4
-  prediction_length: 4
-  num_rolls: 1
--- a/figures/chronos-logo.png
+++ b/figures/chronos-logo.png
--- a/figures/main-figure.png
+++ b/figures/main-figure.png
--- a/figures/zero_shot-agg_scaled_score.svg
+++ b/figures/zero_shot-agg_scaled_score.svg
--- a/notebooks/deploy-chronos-bolt-to-amazon-sagemaker.ipynb
+++ b/notebooks/deploy-chronos-bolt-to-amazon-sagemaker.ipynb
--- a/pyproject.toml
+++ b/pyproject.toml
@ -1,19 +1,19 @@
 [project]
-name = "chronos-forecasting"
-version = "1.5.0"
+name = "chronosx"
+version = "1.0.0"
 authors = [
-  { name="Abdul Fatir Ansari", email="ansarnd@amazon.com" },
-  { name="Lorenzo Stella", email="stellalo@amazon.com" },
-  { name="Caner Turkmen", email="atturkm@amazon.com" },
+  { name="Pedro Mercado", email="pedroml@amazon.com" }
 ]
-description = "Chronos: Pretrained models for time series forecasting"
+description = "ChronosX: Adapting Pretrained Time Series Models with Exogenous Variables"
 readme = "README.md"
 license = { file = "LICENSE" }
-requires-python = ">=3.9"
+requires-python = ">=3.10"
 dependencies = [
-    "torch>=2.0,<3",         # package was tested on 2.2
-    "transformers>=4.48,<5",
-    "accelerate>=0.32,<1",
+    "chronos-forecasting[training]==1.5.0",
+    "transformers==4.47.1",
+    "accelerate==0.34.2",
+    "datasetsforecast==1.0.0",
+    "fev==0.5.0",
 ]
 classifiers = [
    "Programming Language :: Python :: 3",
@ -27,18 +27,12 @@ requires = ["hatchling"]
 build-backend = "hatchling.build"

 [tool.hatch.build.targets.wheel]
-packages = ["src/chronos"]
-
-[project.optional-dependencies]
-test = ["pytest~=8.0", "numpy~=1.21"]
-typecheck = ["mypy~=1.9"]
-training = ["gluonts[pro]~=0.15", "numpy~=1.21", "datasets~=2.18", "typer", "typer-config", "joblib", "scikit-learn", "tensorboard"]
-evaluation = ["gluonts[pro]~=0.15", "numpy~=1.21", "datasets~=2.18", "typer"]
+packages = ["src/chronosx"]

 [project.urls]
-Homepage = "https://github.com/amazon-science/chronos-forecasting"
+Homepage = "https://www.amazon.science/publications/chronosx-adapting-pretrained-time-series-models-with-exogenous-variables"
 Issues = "https://github.com/amazon-science/chronos-forecasting/issues"
-Paper = "https://arxiv.org/abs/2403.07815"
+Paper = "https://arxiv.org/pdf/2503.12107"

 [tool.mypy]
-ignore_missing_imports = true
+ignore_missing_imports = true
--- a/scripts/.DS_Store
+++ b/scripts/.DS_Store
--- a/scripts/README.md
+++ b/scripts/README.md
@ -1,153 +0,0 @@
-# Usage Examples
-
-## Generating Synthetic Time Series (KernelSynth)
-
- Install this package with with the `training` extra:
-    ```
-    pip install "chronos-forecasting[training] @ git+https://github.com/amazon-science/chronos-forecasting.git"
-    ```
- Run `kernel-synth.py`:
-    ```sh
-    # With defaults used in the paper (1M time series and 5 max_kernels)
-    python kernel-synth.py
-
-    # You may optionally specify num-series and max-kernels
-    python kernel-synth.py \
-        --num-series <num of series to generate> \
-        --max-kernels <max number of kernels to use per series>
-    ```
-    The generated time series will be saved in a [GluonTS](https://github.com/awslabs/gluonts)-comptabile arrow file `kernelsynth-data.arrow`.
-
-## Pretraining (and fine-tuning) Chronos models
- Install this package with with the `training` extra:
-    ```
-    pip install "chronos-forecasting[training] @ git+https://github.com/amazon-science/chronos-forecasting.git"
-    ```
- Convert your time series dataset into a GluonTS-compatible file dataset. We recommend using the arrow format. You may use the `convert_to_arrow` function from the following snippet for that. Optionally, you may use [synthetic data from KernelSynth](#generating-synthetic-time-series-kernelsynth) to follow along.
-    ```py
-    from pathlib import Path
-    from typing import List, Union
-
-    import numpy as np
-    from gluonts.dataset.arrow import ArrowWriter
-
-
-    def convert_to_arrow(
-        path: Union[str, Path],
-        time_series: Union[List[np.ndarray], np.ndarray],
-        compression: str = "lz4",
-    ):
-        """
-        Store a given set of series into Arrow format at the specified path.
-
-        Input data can be either a list of 1D numpy arrays, or a single 2D
-        numpy array of shape (num_series, time_length).
-        """
-        assert isinstance(time_series, list) or (
-            isinstance(time_series, np.ndarray) and
-            time_series.ndim == 2
-        )
-
-        # Set an arbitrary start time
-        start = np.datetime64("2000-01-01 00:00", "s")
-
-        dataset = [
-            {"start": start, "target": ts} for ts in time_series
-        ]
-
-        ArrowWriter(compression=compression).write_to_file(
-            dataset,
-            path=path,
-        )
-
-
-    if __name__ == "__main__":
-        # Generate 20 random time series of length 1024
-        time_series = [np.random.randn(1024) for i in range(20)]
-
-        # Convert to GluonTS arrow format
-        convert_to_arrow("./noise-data.arrow", time_series=time_series)
-    ```
- Modify the [training configs](training/configs) to use your data. Let's use the KernelSynth data as an example.
-    ```yaml
-    # List of training data files
-    training_data_paths:
-    - "/path/to/kernelsynth-data.arrow"
-    # Mixing probability of each dataset file
-    probability:
-    - 1.0
-    ```
-    You may optionally change other parameters of the config file, as required. For instance, if you're interested in fine-tuning the model from a pretrained Chronos checkpoint, you should change the `model_id`, set `random_init: false`, and (optionally) change other parameters such as `max_steps` and `learning_rate`.
- Start the training (or fine-tuning) job:
-    ```sh
-    # On single GPU
-    CUDA_VISIBLE_DEVICES=0 python training/train.py --config /path/to/modified/config.yaml
-
-    # On multiple GPUs (example with 8 GPUs)
-    torchrun --nproc-per-node=8 training/train.py --config /path/to/modified/config.yaml
-
-    # Fine-tune `amazon/chronos-t5-small` for 1000 steps with initial learning rate of 1e-3
-    CUDA_VISIBLE_DEVICES=0 python training/train.py --config /path/to/modified/config.yaml \
-        --model-id amazon/chronos-t5-small \
-        --no-random-init \
-        --max-steps 1000 \
-        --learning-rate 0.001
-    ```
-    The output and checkpoints will be saved in `output/run-{id}/`.
-> [!TIP]  
-> If the initial training step is too slow, you might want to change the `shuffle_buffer_length` and/or set `torch_compile` to `false`.
-
-> [!IMPORTANT]  
-> When pretraining causal models (such as GPT2), the training script does [`LastValueImputation`](https://github.com/awslabs/gluonts/blob/f0f2266d520cb980f4c1ce18c28b003ad5cd2599/src/gluonts/transform/feature.py#L103) for missing values by default. If you pretrain causal models, please ensure that missing values are imputed similarly before passing the context tensor to `ChronosPipeline.predict()` for accurate results.
- (Optional) Once trained, you can easily push your fine-tuned model to HuggingFace🤗 Hub. Before that, do not forget to [create an access token](https://huggingface.co/settings/tokens) with **write permissions** and put it in `~/.cache/huggingface/token`. Here's a snippet that will push a fine-tuned model to HuggingFace🤗 Hub at `<your_hf_username>/chronos-t5-small-fine-tuned`.
-    ```py
-    from chronos import ChronosPipeline
-
-    pipeline = ChronosPipeline.from_pretrained("/path/to/fine-tuned/model/ckpt/dir/")
-    pipeline.model.model.push_to_hub("chronos-t5-small-fine-tuned")
-    ```
-
-## Evaluating Chronos models
-
-Follow these steps to compute the WQL and MASE values for the in-domain and zero-shot benchmarks in our paper.
-
- Install this package with with the `evaluation` extra:
-    ```
-    pip install "chronos-forecasting[evaluation] @ git+https://github.com/amazon-science/chronos-forecasting.git"
-    ```
- Run the evaluation script:
-    ```sh
-    # In-domain evaluation
-    # Results will be saved in: evaluation/results/chronos-t5-small-in-domain.csv
-    python evaluation/evaluate.py evaluation/configs/in-domain.yaml evaluation/results/chronos-t5-small-in-domain.csv \
-        --chronos-model-id "amazon/chronos-t5-small" \
-        --batch-size=32 \
-        --device=cuda:0 \
-        --num-samples 20
-
-    # Zero-shot evaluation
-    # Results will be saved in: evaluation/results/chronos-t5-small-zero-shot.csv
-    python evaluation/evaluate.py evaluation/configs/zero-shot.yaml evaluation/results/chronos-t5-small-zero-shot.csv \
-        --chronos-model-id "amazon/chronos-t5-small" \
-        --batch-size=32 \
-        --device=cuda:0 \
-        --num-samples 20
-    ```
- Use the following snippet to compute the aggregated relative WQL and MASE scores:
-    ```py
-    import pandas as pd
-    from scipy.stats import gmean  # requires: pip install scipy
-
-
-    def agg_relative_score(model_df: pd.DataFrame, baseline_df: pd.DataFrame):
-        relative_score = model_df.drop("model", axis="columns") / baseline_df.drop(
-            "model", axis="columns"
-        )
-        return relative_score.agg(gmean)
-
-
-    result_df = pd.read_csv("evaluation/results/chronos-t5-small-in-domain.csv").set_index("dataset")
-    baseline_df = pd.read_csv("evaluation/results/seasonal-naive-in-domain.csv").set_index("dataset")
-
-    agg_score_df = agg_relative_score(result_df, baseline_df)
-    ```
--- a/scripts/evaluation/agg-relative-score.py
+++ b/scripts/evaluation/agg-relative-score.py
@ -1,60 +0,0 @@
-import pandas as pd
-import typer
-from scipy.stats import gmean
-from pathlib import Path
-
-app = typer.Typer(pretty_exceptions_enable=False)
-DEFAULT_RESULTS_DIR = Path(__file__).parent / "results"
-
-
-def agg_relative_score(model_csv: Path, baseline_csv: Path):
-    model_df = pd.read_csv(model_csv).set_index("dataset")
-    baseline_df = pd.read_csv(baseline_csv).set_index("dataset")
-    relative_score = model_df.drop("model", axis="columns") / baseline_df.drop(
-        "model", axis="columns"
-    )
-    return relative_score.agg(gmean)
-
-
-@app.command()
-def main(
-    model_name: str,
-    baseline_name: str = "seasonal-naive",
-    results_dir: Path = DEFAULT_RESULTS_DIR,
-):
-    """
-    Compute the aggregated relative score as reported in the Chronos paper.
-    Results will be saved to {results_dir}/{model_name}-agg-rel-scores.csv
-
-    Parameters
-    ----------
-    model_name : str
-        Name of the model used in the CSV files. The in-domain and zero-shot CSVs
-        are expected to be named {model_name}-in-domain.csv and {model_name}-zero-shot.csv.
-    results_dir : Path, optional, default = results/
-        Directory where results CSVs generated by evaluate.py are stored
-    """
-
-    in_domain_agg_score_df = agg_relative_score(
-        results_dir / f"{model_name}-in-domain.csv",
-        results_dir / f"{baseline_name}-in-domain.csv",
-    )
-    in_domain_agg_score_df.name = "value"
-    in_domain_agg_score_df.index.name = "metric"
-
-    zero_shot_agg_score_df = agg_relative_score(
-        results_dir / f"{model_name}-zero-shot.csv",
-        results_dir / f"{baseline_name}-zero-shot.csv",
-    )
-    zero_shot_agg_score_df.name = "value"
-    zero_shot_agg_score_df.index.name = "metric"
-
-    agg_score_df = pd.concat(
-        {"in-domain": in_domain_agg_score_df, "zero-shot": zero_shot_agg_score_df},
-        names=["benchmark"],
-    )
-    agg_score_df.to_csv(f"{results_dir}/{model_name}-agg-rel-scores.csv")
-
-
-if __name__ == "__main__":
-    app()
--- a/scripts/evaluation/configs/in-domain.yaml
+++ b/scripts/evaluation/configs/in-domain.yaml
@ -1,78 +0,0 @@
-# Backtest configs for the 15 "in-domain" datasets.
-# The training portion of these datasets was part of the 
-# training corpus for Chronos models.
- name: electricity_15min
-  hf_repo: autogluon/chronos_datasets
-  offset: -5376
-  prediction_length: 24
-  num_rolls: 1
- name: monash_electricity_hourly
-  hf_repo: autogluon/chronos_datasets
-  offset: -24
-  prediction_length: 24
-  num_rolls: 1
- name: monash_electricity_weekly
-  hf_repo: autogluon/chronos_datasets
-  offset: -8
-  prediction_length: 8
-  num_rolls: 1
- name: monash_kdd_cup_2018
-  hf_repo: autogluon/chronos_datasets
-  offset: -48
-  prediction_length: 48
-  num_rolls: 1
- name: m4_daily
-  hf_repo: autogluon/chronos_datasets
-  offset: -14
-  prediction_length: 14
-  num_rolls: 1
- name: m4_hourly
-  hf_repo: autogluon/chronos_datasets
-  offset: -48
-  prediction_length: 48
-  num_rolls: 1
- name: m4_monthly
-  hf_repo: autogluon/chronos_datasets
-  offset: -18
-  prediction_length: 18
-  num_rolls: 1
- name: m4_weekly
-  hf_repo: autogluon/chronos_datasets
-  offset: -13
-  prediction_length: 13
-  num_rolls: 1
- name: monash_pedestrian_counts
-  hf_repo: autogluon/chronos_datasets
-  offset: -48
-  prediction_length: 48
-  num_rolls: 1
- name: taxi_30min
-  hf_repo: autogluon/chronos_datasets
-  offset: -48
-  prediction_length: 48
-  num_rolls: 1
- name: uber_tlc_hourly
-  hf_repo: autogluon/chronos_datasets
-  offset: -24
-  prediction_length: 24
-  num_rolls: 1
- name: uber_tlc_daily
-  hf_repo: autogluon/chronos_datasets
-  offset: -7
-  prediction_length: 7
-  num_rolls: 1
- name: monash_rideshare
-  hf_repo: autogluon/chronos_datasets
-  offset: -24
-  prediction_length: 24
-  num_rolls: 1
- name: monash_temperature_rain
-  hf_repo: autogluon/chronos_datasets
-  offset: -30
-  prediction_length: 30
-  num_rolls: 1
- name: monash_london_smart_meters
-  hf_repo: autogluon/chronos_datasets
-  offset: -48
-  prediction_length: 48
-  num_rolls: 1
--- a/scripts/evaluation/configs/zero-shot.yaml
+++ b/scripts/evaluation/configs/zero-shot.yaml
@ -1,137 +0,0 @@
-# Backtest configs for the 27 "zero-shot" datasets.
-# These datasets were not seen by Chronos models during training.
- name: monash_traffic
-  hf_repo: autogluon/chronos_datasets
-  offset: -24
-  prediction_length: 24
-  num_rolls: 1
- name: monash_australian_electricity
-  hf_repo: autogluon/chronos_datasets
-  offset: -48
-  prediction_length: 48
-  num_rolls: 1
- name: ercot
-  hf_repo: autogluon/chronos_datasets
-  offset: -24
-  prediction_length: 24
-  num_rolls: 1
- name: ETTm
-  hf_repo: autogluon/chronos_datasets_extra
-  offset: -96
-  prediction_length: 24
-  num_rolls: 1
- name: ETTh
-  hf_repo: autogluon/chronos_datasets_extra
-  offset: -24
-  prediction_length: 24
-  num_rolls: 1
- name: exchange_rate
-  hf_repo: autogluon/chronos_datasets
-  offset: -30
-  prediction_length: 30
-  num_rolls: 1
- name: nn5
-  hf_repo: autogluon/chronos_datasets
-  offset: -56
-  prediction_length: 56
-  num_rolls: 1
- name: monash_nn5_weekly
-  hf_repo: autogluon/chronos_datasets
-  offset: -8
-  prediction_length: 8
-  num_rolls: 1
- name: monash_weather
-  hf_repo: autogluon/chronos_datasets
-  offset: -30
-  prediction_length: 30
-  num_rolls: 1
- name: monash_covid_deaths
-  hf_repo: autogluon/chronos_datasets
-  offset: -30
-  prediction_length: 30
-  num_rolls: 1
- name: monash_fred_md
-  hf_repo: autogluon/chronos_datasets
-  offset: -12
-  prediction_length: 12
-  num_rolls: 1
- name: m4_quarterly
-  hf_repo: autogluon/chronos_datasets
-  offset: -8
-  prediction_length: 8
-  num_rolls: 1
- name: m4_yearly
-  hf_repo: autogluon/chronos_datasets
-  offset: -6
-  prediction_length: 6
-  num_rolls: 1
- name: dominick
-  hf_repo: autogluon/chronos_datasets
-  offset: -8
-  prediction_length: 8
-  num_rolls: 1
- name: m5
-  hf_repo: autogluon/chronos_datasets
-  offset: -28
-  prediction_length: 28
-  num_rolls: 1
- name: monash_tourism_monthly
-  hf_repo: autogluon/chronos_datasets
-  offset: -24
-  prediction_length: 24
-  num_rolls: 1
- name: monash_tourism_quarterly
-  hf_repo: autogluon/chronos_datasets
-  offset: -8
-  prediction_length: 8
-  num_rolls: 1
- name: monash_tourism_yearly
-  hf_repo: autogluon/chronos_datasets
-  offset: -4
-  prediction_length: 4
-  num_rolls: 1
- name: monash_car_parts
-  hf_repo: autogluon/chronos_datasets
-  offset: -12
-  prediction_length: 12
-  num_rolls: 1
- name: monash_hospital
-  hf_repo: autogluon/chronos_datasets
-  offset: -12
-  prediction_length: 12
-  num_rolls: 1
- name: monash_cif_2016
-  hf_repo: autogluon/chronos_datasets
-  offset: -12
-  prediction_length: 12
-  num_rolls: 1
- name: monash_m1_yearly
-  hf_repo: autogluon/chronos_datasets
-  offset: -6
-  prediction_length: 6
-  num_rolls: 1
- name: monash_m1_quarterly
-  hf_repo: autogluon/chronos_datasets
-  offset: -8
-  prediction_length: 8
-  num_rolls: 1
- name: monash_m1_monthly
-  hf_repo: autogluon/chronos_datasets
-  offset: -18
-  prediction_length: 18
-  num_rolls: 1
- name: monash_m3_monthly
-  hf_repo: autogluon/chronos_datasets
-  offset: -18
-  prediction_length: 18
-  num_rolls: 1
- name: monash_m3_yearly
-  hf_repo: autogluon/chronos_datasets
-  offset: -6
-  prediction_length: 6
-  num_rolls: 1
- name: monash_m3_quarterly
-  hf_repo: autogluon/chronos_datasets
-  offset: -8
-  prediction_length: 8
-  num_rolls: 1
--- a/scripts/evaluation/evaluate.py
+++ b/scripts/evaluation/evaluate.py
@ -1,253 +0,0 @@
-import logging
-from pathlib import Path
-from typing import Iterable, Optional
-
-import datasets
-import numpy as np
-import pandas as pd
-import torch
-import typer
-import yaml
-from gluonts.dataset.split import split
-from gluonts.ev.metrics import MASE, MeanWeightedSumQuantileLoss
-from gluonts.itertools import batcher
-from gluonts.model.evaluation import evaluate_forecasts
-from gluonts.model.forecast import QuantileForecast, SampleForecast
-from tqdm.auto import tqdm
-
-from chronos import (
-    BaseChronosPipeline,
-    ChronosBoltPipeline,
-    ChronosPipeline,
-    ForecastType,
-)
-
-app = typer.Typer(pretty_exceptions_enable=False)
-
-
-def to_gluonts_univariate(hf_dataset: datasets.Dataset):
-    series_fields = [
-        col
-        for col in hf_dataset.features
-        if isinstance(hf_dataset.features[col], datasets.Sequence)
-    ]
-    series_fields.remove("timestamp")
-    dataset_length = hf_dataset.info.splits["train"].num_examples * len(series_fields)
-
-    # Assumes that all time series in the dataset have the same frequency
-    dataset_freq = pd.DatetimeIndex(hf_dataset[0]["timestamp"]).to_period()[0].freqstr
-
-    gts_dataset = []
-    for hf_entry in hf_dataset:
-        for field in series_fields:
-            gts_dataset.append(
-                {
-                    "start": pd.Period(
-                        hf_entry["timestamp"][0],
-                        freq=dataset_freq,
-                    ),
-                    "target": hf_entry[field],
-                }
-            )
-    assert len(gts_dataset) == dataset_length
-
-    return gts_dataset
-
-
-def load_and_split_dataset(backtest_config: dict):
-    hf_repo = backtest_config["hf_repo"]
-    dataset_name = backtest_config["name"]
-    offset = backtest_config["offset"]
-    prediction_length = backtest_config["prediction_length"]
-    num_rolls = backtest_config["num_rolls"]
-
-    # This is needed because the datasets in autogluon/chronos_datasets_extra cannot
-    # be distribued due to license restrictions and must be generated on the fly
-    trust_remote_code = True if hf_repo == "autogluon/chronos_datasets_extra" else False
-
-    ds = datasets.load_dataset(
-        hf_repo, dataset_name, split="train", trust_remote_code=trust_remote_code
-    )
-    ds.set_format("numpy")
-
-    gts_dataset = to_gluonts_univariate(ds)
-
-    # Split dataset for evaluation
-    _, test_template = split(gts_dataset, offset=offset)
-    test_data = test_template.generate_instances(prediction_length, windows=num_rolls)
-
-    return test_data
-
-
-def generate_forecasts(
-    test_data_input: Iterable,
-    pipeline: BaseChronosPipeline,
-    prediction_length: int,
-    batch_size: int,
-    **predict_kwargs,
-):
-    # Generate forecasts
-    forecast_outputs = []
-    for batch in tqdm(batcher(test_data_input, batch_size=batch_size)):
-        context = [torch.tensor(entry["target"]) for entry in batch]
-        forecast_outputs.append(
-            pipeline.predict(
-                context,
-                prediction_length=prediction_length,
-                **predict_kwargs,
-            ).numpy()
-        )
-    forecast_outputs = np.concatenate(forecast_outputs)
-
-    # Convert forecast samples into gluonts Forecast objects
-    forecasts = []
-    for item, ts in zip(forecast_outputs, test_data_input):
-        forecast_start_date = ts["start"] + len(ts["target"])
-
-        if pipeline.forecast_type == ForecastType.SAMPLES:
-            forecasts.append(
-                SampleForecast(samples=item, start_date=forecast_start_date)
-            )
-        elif pipeline.forecast_type == ForecastType.QUANTILES:
-            forecasts.append(
-                QuantileForecast(
-                    forecast_arrays=item,
-                    forecast_keys=list(map(str, pipeline.quantiles)),
-                    start_date=forecast_start_date,
-                )
-            )
-
-    return forecasts
-
-
-@app.command()
-def main(
-    config_path: Path,
-    metrics_path: Path,
-    chronos_model_id: str = "amazon/chronos-t5-small",
-    device: str = "cuda",
-    torch_dtype: str = "bfloat16",
-    batch_size: int = 32,
-    num_samples: int = 20,
-    temperature: Optional[float] = None,
-    top_k: Optional[int] = None,
-    top_p: Optional[float] = None,
-):
-    """Evaluate Chronos models.
-
-    Parameters
-    ----------
-    config_path : Path
-        Path to the evaluation config. See ./configs/.
-    metrics_path : Path
-        Path to the CSV file where metrics will be saved.
-    chronos_model_id : str, optional, default = "amazon/chronos-t5-small"
-        HuggingFace ID of the Chronos model or local path
-        Available models on HuggingFace:
-        Chronos:
-            - amazon/chronos-t5-tiny
-            - amazon/chronos-t5-mini
-            - amazon/chronos-t5-small
-            - amazon/chronos-t5-base
-            - amazon/chronos-t5-large
-        Chronos-Bolt:
-            - amazon/chronos-bolt-tiny
-            - amazon/chronos-bolt-mini
-            - amazon/chronos-bolt-small
-            - amazon/chronos-bolt-base
-    device : str, optional, default = "cuda"
-        Device on which inference will be performed
-    torch_dtype : str, optional
-        Model's dtype, by default "bfloat16"
-    batch_size : int, optional, default = 32
-        Batch size for inference. For Chronos-Bolt models, significantly larger
-        batch sizes can be used
-    num_samples : int, optional, default = 20
-        Number of samples to draw when using the original Chronos models
-    temperature : Optional[float], optional, default = 1.0
-        Softmax temperature to used for the original Chronos models
-    top_k : Optional[int], optional, default = 50
-        Top-K sampling, by default None
-    top_p : Optional[float], optional, default = 1.0
-        Top-p sampling, by default None
-    """
-    if isinstance(torch_dtype, str):
-        torch_dtype = getattr(torch, torch_dtype)
-    assert isinstance(torch_dtype, torch.dtype)
-
-    # Load Chronos
-    pipeline = BaseChronosPipeline.from_pretrained(
-        chronos_model_id,
-        device_map=device,
-        torch_dtype=torch_dtype,
-    )
-
-    if isinstance(pipeline, ChronosPipeline):
-        predict_kwargs = dict(
-            num_samples=num_samples,
-            temperature=temperature,
-            top_k=top_k,
-            top_p=top_p,
-        )
-    elif isinstance(pipeline, ChronosBoltPipeline):
-        predict_kwargs = {}
-
-    # Load backtest configs
-    with open(config_path) as fp:
-        backtest_configs = yaml.safe_load(fp)
-
-    result_rows = []
-    for config in backtest_configs:
-        dataset_name = config["name"]
-        prediction_length = config["prediction_length"]
-
-        logger.info(f"Loading {dataset_name}")
-        test_data = load_and_split_dataset(backtest_config=config)
-
-        logger.info(
-            f"Generating forecasts for {dataset_name} "
-            f"({len(test_data.input)} time series)"
-        )
-        forecasts = generate_forecasts(
-            test_data.input,
-            pipeline=pipeline,
-            prediction_length=prediction_length,
-            batch_size=batch_size,
-            **predict_kwargs,
-        )
-
-        logger.info(f"Evaluating forecasts for {dataset_name}")
-        metrics = (
-            evaluate_forecasts(
-                forecasts,
-                test_data=test_data,
-                metrics=[
-                    MASE(),
-                    MeanWeightedSumQuantileLoss(np.arange(0.1, 1.0, 0.1)),
-                ],
-                batch_size=5000,
-            )
-            .reset_index(drop=True)
-            .to_dict(orient="records")
-        )
-        result_rows.append(
-            {"dataset": dataset_name, "model": chronos_model_id, **metrics[0]}
-        )
-
-    # Save results to a CSV file
-    results_df = (
-        pd.DataFrame(result_rows)
-        .rename(
-            {"MASE[0.5]": "MASE", "mean_weighted_sum_quantile_loss": "WQL"},
-            axis="columns",
-        )
-        .sort_values(by="dataset")
-    )
-    results_df.to_csv(metrics_path, index=False)
-
-
-if __name__ == "__main__":
-    logging.basicConfig(format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
-    logger = logging.getLogger("Chronos Evaluation")
-    logger.setLevel(logging.INFO)
-    app()
--- a/scripts/evaluation/results/chronos-bolt-base-agg-rel-scores.csv
+++ b/scripts/evaluation/results/chronos-bolt-base-agg-rel-scores.csv
@ -1,5 +0,0 @@
-benchmark,metric,value
-in-domain,MASE,0.6800133628315155
-in-domain,WQL,0.5339263811489279
-zero-shot,MASE,0.7914551113353537
-zero-shot,WQL,0.6241424984163773
--- a/scripts/evaluation/results/chronos-bolt-base-in-domain.csv
+++ b/scripts/evaluation/results/chronos-bolt-base-in-domain.csv
@ -1,16 +0,0 @@
-dataset,model,MASE,WQL
-electricity_15min,amazon/chronos-bolt-base,0.41069374835605243,0.0703533790998506
-m4_daily,amazon/chronos-bolt-base,3.205192517121196,0.02110308498174413
-m4_hourly,amazon/chronos-bolt-base,0.8350129849014075,0.025353803894164
-m4_monthly,amazon/chronos-bolt-base,0.9491758928362231,0.09382496106659234
-m4_weekly,amazon/chronos-bolt-base,2.0847827409162742,0.03816605075768161
-monash_electricity_hourly,amazon/chronos-bolt-base,1.254966217685461,0.09442192616975713
-monash_electricity_weekly,amazon/chronos-bolt-base,1.8391546050108039,0.06410971963960499
-monash_kdd_cup_2018,amazon/chronos-bolt-base,0.6405985809360102,0.2509172188706336
-monash_london_smart_meters,amazon/chronos-bolt-base,0.701398572604996,0.3218915088923906
-monash_pedestrian_counts,amazon/chronos-bolt-base,0.2646412642278343,0.18789459806066328
-monash_rideshare,amazon/chronos-bolt-base,0.7695376426829713,0.11637119433040358
-monash_temperature_rain,amazon/chronos-bolt-base,0.8983612698773724,0.6050555216496304
-taxi_30min,amazon/chronos-bolt-base,0.7688908266765317,0.2363178601205094
-uber_tlc_daily,amazon/chronos-bolt-base,0.8231767493519677,0.0926036406916842
-uber_tlc_hourly,amazon/chronos-bolt-base,0.6632193728217927,0.14987786887626975
--- a/scripts/evaluation/results/chronos-bolt-base-zero-shot.csv
+++ b/scripts/evaluation/results/chronos-bolt-base-zero-shot.csv
@ -1,28 +0,0 @@
-dataset,model,MASE,WQL
-ETTh,amazon/chronos-bolt-base,0.7479154031956647,0.07062173821055001
-ETTm,amazon/chronos-bolt-base,0.6334357237512225,0.052261607745858835
-dominick,amazon/chronos-bolt-base,0.8560272479913918,0.3453573743726445
-ercot,amazon/chronos-bolt-base,0.6933217425507392,0.02142183038021456
-exchange_rate,amazon/chronos-bolt-base,1.7095176257412634,0.01200682136751536
-m4_quarterly,amazon/chronos-bolt-base,1.2244670010522907,0.0771066518089854
-m4_yearly,amazon/chronos-bolt-base,3.513752058541554,0.12142798053483984
-m5,amazon/chronos-bolt-base,0.9152230096463854,0.561999688057527
-monash_australian_electricity,amazon/chronos-bolt-base,0.7403239930185613,0.03584034231329335
-monash_car_parts,amazon/chronos-bolt-base,0.8550263912438314,0.9945122291263591
-monash_cif_2016,amazon/chronos-bolt-base,0.9988541862779904,0.016456104842296485
-monash_covid_deaths,amazon/chronos-bolt-base,38.901749109066415,0.047410971217640714
-monash_fred_md,amazon/chronos-bolt-base,0.6468787708795645,0.04185083716355386
-monash_hospital,amazon/chronos-bolt-base,0.6883138394434054,0.057032869931903894
-monash_m1_monthly,amazon/chronos-bolt-base,1.0997677446267855,0.1392311148066238
-monash_m1_quarterly,amazon/chronos-bolt-base,1.7737851980875563,0.1007118219350403
-monash_m1_yearly,amazon/chronos-bolt-base,4.404672537832342,0.1504617654430952
-monash_m3_monthly,amazon/chronos-bolt-base,0.8510696834878182,0.09269673913736748
-monash_m3_quarterly,amazon/chronos-bolt-base,1.2890908822598466,0.07615133571216029
-monash_m3_yearly,amazon/chronos-bolt-base,2.9067097980770082,0.12934285625258413
-monash_nn5_weekly,amazon/chronos-bolt-base,0.9158766337957451,0.08352114810139548
-monash_tourism_monthly,amazon/chronos-bolt-base,1.5283388458731357,0.09026425492612797
-monash_tourism_quarterly,amazon/chronos-bolt-base,1.756127005530011,0.06448060953595125
-monash_tourism_yearly,amazon/chronos-bolt-base,3.691545772463519,0.16548820700844424
-monash_traffic,amazon/chronos-bolt-base,0.7843310867739336,0.23148632068725078
-monash_weather,amazon/chronos-bolt-base,0.8115247139672316,0.13350830777170594
-nn5,amazon/chronos-bolt-base,0.5764084996361287,0.1500519584148468
--- a/scripts/evaluation/results/chronos-bolt-mini-agg-rel-scores.csv
+++ b/scripts/evaluation/results/chronos-bolt-mini-agg-rel-scores.csv
@ -1,5 +0,0 @@
-benchmark,metric,value
-in-domain,MASE,0.7268373301543752
-in-domain,WQL,0.565140251955324
-zero-shot,MASE,0.8221798917822493
-zero-shot,WQL,0.6441645845380903
--- a/scripts/evaluation/results/chronos-bolt-mini-in-domain.csv
+++ b/scripts/evaluation/results/chronos-bolt-mini-in-domain.csv
@ -1,16 +0,0 @@
-dataset,model,MASE,WQL
-electricity_15min,amazon/chronos-bolt-mini,0.44185193304080733,0.0731477927531107
-m4_daily,amazon/chronos-bolt-mini,3.1342608828747456,0.0206872246743766
-m4_hourly,amazon/chronos-bolt-mini,0.9218285923038745,0.024383114886067574
-m4_monthly,amazon/chronos-bolt-mini,0.9628339921394529,0.09502498697494888
-m4_weekly,amazon/chronos-bolt-mini,2.2330452369879255,0.039393515325238534
-monash_electricity_hourly,amazon/chronos-bolt-mini,1.6195944363428718,0.11468972600782207
-monash_electricity_weekly,amazon/chronos-bolt-mini,1.866105365159433,0.06019900031840434
-monash_kdd_cup_2018,amazon/chronos-bolt-mini,0.74790954883436,0.3012661161484388
-monash_london_smart_meters,amazon/chronos-bolt-mini,0.7187830347765344,0.32984510693830227
-monash_pedestrian_counts,amazon/chronos-bolt-mini,0.308633944815819,0.23331301029432483
-monash_rideshare,amazon/chronos-bolt-mini,0.818948044410056,0.1297966960374544
-monash_temperature_rain,amazon/chronos-bolt-mini,0.9035244443682741,0.605031064086567
-taxi_30min,amazon/chronos-bolt-mini,0.812010120941363,0.25232294549917317
-uber_tlc_daily,amazon/chronos-bolt-mini,0.8507256206478295,0.10101757743084538
-uber_tlc_hourly,amazon/chronos-bolt-mini,0.6685484898085609,0.1515245941548974
--- a/scripts/evaluation/results/chronos-bolt-mini-zero-shot.csv
+++ b/scripts/evaluation/results/chronos-bolt-mini-zero-shot.csv
@ -1,28 +0,0 @@
-dataset,model,MASE,WQL
-ETTh,amazon/chronos-bolt-mini,0.8057126710113404,0.07740387596411452
-ETTm,amazon/chronos-bolt-mini,0.6100793941108849,0.05129333450944573
-dominick,amazon/chronos-bolt-mini,0.8664152477208024,0.3499696999160997
-ercot,amazon/chronos-bolt-mini,0.6871250728215426,0.02448804863744021
-exchange_rate,amazon/chronos-bolt-mini,1.3520551553333662,0.00934663373172766
-m4_quarterly,amazon/chronos-bolt-mini,1.2569644266281508,0.07833787023275976
-m4_yearly,amazon/chronos-bolt-mini,3.7611003052413796,0.12931927951165456
-m5,amazon/chronos-bolt-mini,0.9188876472137485,0.5661303206519673
-monash_australian_electricity,amazon/chronos-bolt-mini,0.8823559450287066,0.04493688824488474
-monash_car_parts,amazon/chronos-bolt-mini,0.8604081423647779,1.0041876404811494
-monash_cif_2016,amazon/chronos-bolt-mini,1.0762361363763873,0.017641893717784202
-monash_covid_deaths,amazon/chronos-bolt-mini,38.83915011538576,0.06098317835750057
-monash_fred_md,amazon/chronos-bolt-mini,0.6169859211923081,0.03256236965040934
-monash_hospital,amazon/chronos-bolt-mini,0.6924431064606051,0.05766349075348645
-monash_m1_monthly,amazon/chronos-bolt-mini,1.147893030263777,0.13270222658510553
-monash_m1_quarterly,amazon/chronos-bolt-mini,1.8662100001165818,0.09846363409254102
-monash_m1_yearly,amazon/chronos-bolt-mini,5.319154632748303,0.16167328827180308
-monash_m3_monthly,amazon/chronos-bolt-mini,0.8758452776118432,0.09493431248614057
-monash_m3_quarterly,amazon/chronos-bolt-mini,1.3555175243802005,0.07808062465932723
-monash_m3_yearly,amazon/chronos-bolt-mini,3.605769430055575,0.15711010456482008
-monash_nn5_weekly,amazon/chronos-bolt-mini,0.9347141924977239,0.08522899825844342
-monash_tourism_monthly,amazon/chronos-bolt-mini,1.649587479665881,0.0979648261309891
-monash_tourism_quarterly,amazon/chronos-bolt-mini,1.8471553663088986,0.06501077791766902
-monash_tourism_yearly,amazon/chronos-bolt-mini,3.9932920493826245,0.1743539122097316
-monash_traffic,amazon/chronos-bolt-mini,0.8355442361271347,0.24351051123330386
-monash_weather,amazon/chronos-bolt-mini,0.800013628350165,0.13041050756802045
-nn5,amazon/chronos-bolt-mini,0.611917632501032,0.1570111102680171
--- a/scripts/evaluation/results/chronos-bolt-small-agg-rel-scores.csv
+++ b/scripts/evaluation/results/chronos-bolt-small-agg-rel-scores.csv
@ -1,5 +0,0 @@
-benchmark,metric,value
-in-domain,MASE,0.7030801652116672
-in-domain,WQL,0.5443547623341555
-zero-shot,MASE,0.8192127745093378
-zero-shot,WQL,0.6356097843099521
--- a/scripts/evaluation/results/chronos-bolt-small-in-domain.csv
+++ b/scripts/evaluation/results/chronos-bolt-small-in-domain.csv
@ -1,16 +0,0 @@
-dataset,model,MASE,WQL
-electricity_15min,amazon/chronos-bolt-small,0.44920089250026723,0.08115291306964295
-m4_daily,amazon/chronos-bolt-small,3.201966619014735,0.02143368277732494
-m4_hourly,amazon/chronos-bolt-small,0.8686298207618999,0.020368729287465817
-m4_monthly,amazon/chronos-bolt-small,0.9537717737278778,0.0939247807527992
-m4_weekly,amazon/chronos-bolt-small,2.1236755094789177,0.03785184715517262
-monash_electricity_hourly,amazon/chronos-bolt-small,1.3728906161330452,0.09452411472431674
-monash_electricity_weekly,amazon/chronos-bolt-small,1.8703239487242378,0.06648479071326366
-monash_kdd_cup_2018,amazon/chronos-bolt-small,0.6458631909979771,0.25148489931571666
-monash_london_smart_meters,amazon/chronos-bolt-small,0.7126939688565166,0.326874529903459
-monash_pedestrian_counts,amazon/chronos-bolt-small,0.3015070035798365,0.2285590441093863
-monash_rideshare,amazon/chronos-bolt-small,0.823726965741684,0.12409769473500927
-monash_temperature_rain,amazon/chronos-bolt-small,0.8980348827836525,0.5984819599873311
-taxi_30min,amazon/chronos-bolt-small,0.7597818149895785,0.2348569752311862
-uber_tlc_daily,amazon/chronos-bolt-small,0.8460854328036702,0.09666483354735897
-uber_tlc_hourly,amazon/chronos-bolt-small,0.6662547495017634,0.1524256346268063
--- a/scripts/evaluation/results/chronos-bolt-small-zero-shot.csv
+++ b/scripts/evaluation/results/chronos-bolt-small-zero-shot.csv
@ -1,28 +0,0 @@
-dataset,model,MASE,WQL
-ETTh,amazon/chronos-bolt-small,0.792521748651108,0.07590654063011319
-ETTm,amazon/chronos-bolt-small,0.6209623928936988,0.05056189722606397
-dominick,amazon/chronos-bolt-small,0.8706134610400587,0.34811141409475416
-ercot,amazon/chronos-bolt-small,0.7562857616685997,0.02596064260343696
-exchange_rate,amazon/chronos-bolt-small,1.774835301692689,0.011363548847621512
-m4_quarterly,amazon/chronos-bolt-small,1.2478142413437487,0.07808795122806232
-m4_yearly,amazon/chronos-bolt-small,3.6925595655002574,0.12772564181388502
-m5,amazon/chronos-bolt-small,0.9195435643571084,0.5668430814831332
-monash_australian_electricity,amazon/chronos-bolt-small,0.8128424798841111,0.041509852162861564
-monash_car_parts,amazon/chronos-bolt-small,0.8584574663781737,1.0074689402521324
-monash_cif_2016,amazon/chronos-bolt-small,1.0182471909074982,0.01581964877692293
-monash_covid_deaths,amazon/chronos-bolt-small,36.467595559655145,0.0427382859406882
-monash_fred_md,amazon/chronos-bolt-small,0.6132863794635253,0.03730410577241995
-monash_hospital,amazon/chronos-bolt-small,0.6954489513780618,0.058119864671526154
-monash_m1_monthly,amazon/chronos-bolt-small,1.1277621848099244,0.1335656174632902
-monash_m1_quarterly,amazon/chronos-bolt-small,1.8356144904231688,0.09363028483838018
-monash_m1_yearly,amazon/chronos-bolt-small,5.098146069746402,0.15669928873371905
-monash_m3_monthly,amazon/chronos-bolt-small,0.8685125121306435,0.09396568468255145
-monash_m3_quarterly,amazon/chronos-bolt-small,1.3269103591066727,0.07691022995374203
-monash_m3_yearly,amazon/chronos-bolt-small,3.40993282700627,0.1547639821304127
-monash_nn5_weekly,amazon/chronos-bolt-small,0.9266513350636507,0.08452821221908001
-monash_tourism_monthly,amazon/chronos-bolt-small,1.6106732721197876,0.09362336754317802
-monash_tourism_quarterly,amazon/chronos-bolt-small,1.8357819365308639,0.06734337535269994
-monash_tourism_yearly,amazon/chronos-bolt-small,3.8963100495394194,0.16766064312072784
-monash_traffic,amazon/chronos-bolt-small,0.8598507749866499,0.25173786112983054
-monash_weather,amazon/chronos-bolt-small,0.8020408743877911,0.13258563963844888
-nn5,amazon/chronos-bolt-small,0.5833047644729239,0.15066847836762787
--- a/scripts/evaluation/results/chronos-bolt-tiny-agg-rel-scores.csv
+++ b/scripts/evaluation/results/chronos-bolt-tiny-agg-rel-scores.csv
@ -1,5 +0,0 @@
-benchmark,metric,value
-in-domain,MASE,0.7403252781013574
-in-domain,WQL,0.5733728165523524
-zero-shot,MASE,0.8445407343705457
-zero-shot,WQL,0.6678781905023173
--- a/scripts/evaluation/results/chronos-bolt-tiny-in-domain.csv
+++ b/scripts/evaluation/results/chronos-bolt-tiny-in-domain.csv
@ -1,16 +0,0 @@
-dataset,model,MASE,WQL
-electricity_15min,amazon/chronos-bolt-tiny,0.4676384089765091,0.0861229808117837
-m4_daily,amazon/chronos-bolt-tiny,3.1789994761356795,0.020961883512815756
-m4_hourly,amazon/chronos-bolt-tiny,0.9348005698736752,0.021087527284114574
-m4_monthly,amazon/chronos-bolt-tiny,0.965298729632761,0.0950380483243082
-m4_weekly,amazon/chronos-bolt-tiny,2.261575511029903,0.04093653263178429
-monash_electricity_hourly,amazon/chronos-bolt-tiny,1.5739346351263623,0.10808418398945202
-monash_electricity_weekly,amazon/chronos-bolt-tiny,1.8628689103722829,0.05773335283584782
-monash_kdd_cup_2018,amazon/chronos-bolt-tiny,0.6869549985391232,0.28012801092758166
-monash_london_smart_meters,amazon/chronos-bolt-tiny,0.7284234905933779,0.33496438244693033
-monash_pedestrian_counts,amazon/chronos-bolt-tiny,0.32338947321773864,0.2530637833749087
-monash_rideshare,amazon/chronos-bolt-tiny,0.8562780835002918,0.1304317657933891
-monash_temperature_rain,amazon/chronos-bolt-tiny,0.9030707620825977,0.6064087080755548
-taxi_30min,amazon/chronos-bolt-tiny,0.9122159603256838,0.28002194370731626
-uber_tlc_daily,amazon/chronos-bolt-tiny,0.9087055420190513,0.11193388685815164
-uber_tlc_hourly,amazon/chronos-bolt-tiny,0.6716569179590032,0.15310845458208555
--- a/scripts/evaluation/results/chronos-bolt-tiny-zero-shot.csv
+++ b/scripts/evaluation/results/chronos-bolt-tiny-zero-shot.csv
@ -1,28 +0,0 @@
-dataset,model,MASE,WQL
-ETTh,amazon/chronos-bolt-tiny,0.7941225847155844,0.07480860969990633
-ETTm,amazon/chronos-bolt-tiny,0.6508270995240056,0.05440068825429993
-dominick,amazon/chronos-bolt-tiny,0.876060127216559,0.35175949052933253
-ercot,amazon/chronos-bolt-tiny,0.7309134980173839,0.02468604544464515
-exchange_rate,amazon/chronos-bolt-tiny,1.6857262567077134,0.011477224264784112
-m4_quarterly,amazon/chronos-bolt-tiny,1.2605908919338378,0.0789049420017836
-m4_yearly,amazon/chronos-bolt-tiny,3.7118394116161757,0.1286932555969197
-m5,amazon/chronos-bolt-tiny,0.9195469670062033,0.5634881835998845
-monash_australian_electricity,amazon/chronos-bolt-tiny,0.8419304693259403,0.042040993880313904
-monash_car_parts,amazon/chronos-bolt-tiny,0.8625579150452282,1.0009987800801836
-monash_cif_2016,amazon/chronos-bolt-tiny,1.095219642027011,0.017550336784241796
-monash_covid_deaths,amazon/chronos-bolt-tiny,40.674057986280744,0.06723714516685976
-monash_fred_md,amazon/chronos-bolt-tiny,0.6127387450520702,0.04747523852271518
-monash_hospital,amazon/chronos-bolt-tiny,0.6980246281225624,0.05864223243167421
-monash_m1_monthly,amazon/chronos-bolt-tiny,1.1625495971731141,0.13142237467151166
-monash_m1_quarterly,amazon/chronos-bolt-tiny,1.8941765599193754,0.09972207844232561
-monash_m1_yearly,amazon/chronos-bolt-tiny,5.136332694531757,0.160331813128038
-monash_m3_monthly,amazon/chronos-bolt-tiny,0.8744553726704598,0.09435519378597752
-monash_m3_quarterly,amazon/chronos-bolt-tiny,1.364563776692303,0.07875066385737857
-monash_m3_yearly,amazon/chronos-bolt-tiny,3.3685961410254928,0.15158076519486274
-monash_nn5_weekly,amazon/chronos-bolt-tiny,0.9324436794013877,0.0847385189968909
-monash_tourism_monthly,amazon/chronos-bolt-tiny,1.7895936775088157,0.1058167042693116
-monash_tourism_quarterly,amazon/chronos-bolt-tiny,2.095262637810499,0.0710732570354461
-monash_tourism_yearly,amazon/chronos-bolt-tiny,4.042821441327848,0.172613367251472
-monash_traffic,amazon/chronos-bolt-tiny,0.8836032533767518,0.2574297134210491
-monash_weather,amazon/chronos-bolt-tiny,0.8005348255663177,0.13111355494466076
-nn5,amazon/chronos-bolt-tiny,0.7228248498869763,0.1816913098894226
--- a/scripts/evaluation/results/chronos-t5-base-agg-rel-scores.csv
+++ b/scripts/evaluation/results/chronos-t5-base-agg-rel-scores.csv
@ -1,5 +0,0 @@
-benchmark,metric,value
-in-domain,MASE,0.7007558507277635
-in-domain,WQL,0.5786300105297922
-zero-shot,MASE,0.8155209321160994
-zero-shot,WQL,0.6424634919486323
--- a/scripts/evaluation/results/chronos-t5-base-in-domain.csv
+++ b/scripts/evaluation/results/chronos-t5-base-in-domain.csv
@ -1,16 +0,0 @@
-dataset,model,MASE,WQL
-electricity_15min,amazon/chronos-t5-base,0.39879754957261204,0.07738953262286181
-m4_daily,amazon/chronos-t5-base,3.160575865614404,0.02194256368254537
-m4_hourly,amazon/chronos-t5-base,0.6938747745332102,0.026354948301302205
-m4_monthly,amazon/chronos-t5-base,0.971951848755026,0.10355213196432872
-m4_weekly,amazon/chronos-t5-base,2.0143841267657945,0.03639741235815474
-monash_electricity_hourly,amazon/chronos-t5-base,1.5717251971297332,0.1078882125804548
-monash_electricity_weekly,amazon/chronos-t5-base,1.7862927210886668,0.06255982783148449
-monash_kdd_cup_2018,amazon/chronos-t5-base,0.6335225775496138,0.2684272353843692
-monash_london_smart_meters,amazon/chronos-t5-base,0.8362014889190201,0.4265549499082726
-monash_pedestrian_counts,amazon/chronos-t5-base,0.2817708325561419,0.20810108090665583
-monash_rideshare,amazon/chronos-t5-base,0.8614480533175364,0.1356591190888703
-monash_temperature_rain,amazon/chronos-t5-base,0.9692405156151607,0.660155448791624
-taxi_30min,amazon/chronos-t5-base,0.8186287575356217,0.26236060366367003
-uber_tlc_daily,amazon/chronos-t5-base,0.8338648311528079,0.0970875577681834
-uber_tlc_hourly,amazon/chronos-t5-base,0.6647193438331641,0.15436646659512512
--- a/scripts/evaluation/results/chronos-t5-base-zero-shot.csv
+++ b/scripts/evaluation/results/chronos-t5-base-zero-shot.csv
@ -1,28 +0,0 @@
-dataset,model,MASE,WQL
-ETTh,amazon/chronos-t5-base,0.7653491494991778,0.08087267701042929
-ETTm,amazon/chronos-t5-base,0.7737006634032871,0.07008650633028274
-dominick,amazon/chronos-t5-base,0.8194044957573132,0.33201307438298133
-ercot,amazon/chronos-t5-base,0.5014399265038706,0.013589435745554596
-exchange_rate,amazon/chronos-t5-base,2.055616906406159,0.011066070028466317
-m4_quarterly,amazon/chronos-t5-base,1.2253036947743137,0.08327936201395683
-m4_yearly,amazon/chronos-t5-base,3.639991540990927,0.13539258375263963
-m5,amazon/chronos-t5-base,0.9391874615167101,0.5867234116216755
-monash_australian_electricity,amazon/chronos-t5-base,1.2944069383163321,0.07070604202031877
-monash_car_parts,amazon/chronos-t5-base,0.9071940271035218,1.077797124337994
-monash_cif_2016,amazon/chronos-t5-base,0.9840747802099565,0.011825556826558836
-monash_covid_deaths,amazon/chronos-t5-base,42.68503365359237,0.042229910495746356
-monash_fred_md,amazon/chronos-t5-base,0.4857773806790164,0.021204829049512715
-monash_hospital,amazon/chronos-t5-base,0.7053005021431749,0.05630687524507516
-monash_m1_monthly,amazon/chronos-t5-base,1.1153039466137842,0.12724419775326076
-monash_m1_quarterly,amazon/chronos-t5-base,1.746093728928804,0.1123583549291933
-monash_m1_yearly,amazon/chronos-t5-base,4.401291522370069,0.18541586641719554
-monash_m3_monthly,amazon/chronos-t5-base,0.8627172231908679,0.09640536232169555
-monash_m3_quarterly,amazon/chronos-t5-base,1.1696030904401578,0.07392876900131434
-monash_m3_yearly,amazon/chronos-t5-base,3.1298600218573775,0.1486674447940158
-monash_nn5_weekly,amazon/chronos-t5-base,0.9334860602210187,0.08972736821598823
-monash_tourism_monthly,amazon/chronos-t5-base,1.7937702435879332,0.10260220444264027
-monash_tourism_quarterly,amazon/chronos-t5-base,1.7791494997972261,0.06852507950474919
-monash_tourism_yearly,amazon/chronos-t5-base,3.8359926053603197,0.20722699382964643
-monash_traffic,amazon/chronos-t5-base,0.8015262383138622,0.25565153982140926
-monash_weather,amazon/chronos-t5-base,0.8159511190589147,0.13802320967454584
-nn5,amazon/chronos-t5-base,0.5927076179914024,0.1630476065585159
--- a/scripts/evaluation/results/chronos-t5-large-agg-rel-scores.csv
+++ b/scripts/evaluation/results/chronos-t5-large-agg-rel-scores.csv
@ -1,5 +0,0 @@
-benchmark,metric,value
-in-domain,MASE,0.6944869734691035
-in-domain,WQL,0.5596857927462495
-zero-shot,MASE,0.8213682201405101
-zero-shot,WQL,0.6504834081319559
--- a/scripts/evaluation/results/chronos-t5-large-in-domain.csv
+++ b/scripts/evaluation/results/chronos-t5-large-in-domain.csv
@ -1,16 +0,0 @@
-dataset,model,MASE,WQL
-electricity_15min,amazon/chronos-t5-large,0.3866310906621673,0.07759528615667297
-m4_daily,amazon/chronos-t5-large,3.134560968849699,0.02158279722410466
-m4_hourly,amazon/chronos-t5-large,0.6975930649233378,0.02086427219957674
-m4_monthly,amazon/chronos-t5-large,0.9585550091429409,0.10091221432814867
-m4_weekly,amazon/chronos-t5-large,2.0191422600104425,0.036912838355537186
-monash_electricity_hourly,amazon/chronos-t5-large,1.4069912853901292,0.09642382339452431
-monash_electricity_weekly,amazon/chronos-t5-large,1.7501880036182798,0.05765306465830232
-monash_kdd_cup_2018,amazon/chronos-t5-large,0.6788042816175427,0.2853553329804835
-monash_london_smart_meters,amazon/chronos-t5-large,0.8290300790418726,0.4235436387853963
-monash_pedestrian_counts,amazon/chronos-t5-large,0.2764118100521592,0.18692234491663473
-monash_rideshare,amazon/chronos-t5-large,0.8758058784466208,0.140260325368757
-monash_temperature_rain,amazon/chronos-t5-large,0.9738403865035117,0.6604571928063249
-taxi_30min,amazon/chronos-t5-large,0.8245662397270109,0.2653520120326771
-uber_tlc_daily,amazon/chronos-t5-large,0.8044165990021739,0.09499035584302248
-uber_tlc_hourly,amazon/chronos-t5-large,0.6700665937164474,0.15190288476653066
--- a/scripts/evaluation/results/chronos-t5-large-zero-shot.csv
+++ b/scripts/evaluation/results/chronos-t5-large-zero-shot.csv
@ -1,28 +0,0 @@
-dataset,model,MASE,WQL
-ETTh,amazon/chronos-t5-large,0.78160443631164,0.07884375667736107
-ETTm,amazon/chronos-t5-large,0.7325919639389967,0.06656858270921162
-dominick,amazon/chronos-t5-large,0.8200108271155829,0.3311575649734524
-ercot,amazon/chronos-t5-large,0.6050812633742764,0.01822996942395577
-exchange_rate,amazon/chronos-t5-large,2.3439287001928744,0.014841231672174684
-m4_quarterly,amazon/chronos-t5-large,1.2169666607868148,0.08235162400898562
-m4_yearly,amazon/chronos-t5-large,3.5524979814018947,0.1325675848907479
-m5,amazon/chronos-t5-large,0.9422990989146737,0.585615077637479
-monash_australian_electricity,amazon/chronos-t5-large,1.480849838497958,0.07973968848149568
-monash_car_parts,amazon/chronos-t5-large,0.901547374873302,1.0467398096496576
-monash_cif_2016,amazon/chronos-t5-large,0.9906388185665337,0.011966178555329998
-monash_covid_deaths,amazon/chronos-t5-large,44.07354193681227,0.06108999981222163
-monash_fred_md,amazon/chronos-t5-large,0.5184400880318044,0.01675533888399231
-monash_hospital,amazon/chronos-t5-large,0.7055308474630898,0.0552450850258613
-monash_m1_monthly,amazon/chronos-t5-large,1.0888995301234758,0.12729911122909737
-monash_m1_quarterly,amazon/chronos-t5-large,1.7477134564031453,0.10618253695380094
-monash_m1_yearly,amazon/chronos-t5-large,4.250667049416348,0.17128879333643188
-monash_m3_monthly,amazon/chronos-t5-large,0.8559326975903808,0.09572577431396007
-monash_m3_quarterly,amazon/chronos-t5-large,1.1867267751420676,0.07449254281607631
-monash_m3_yearly,amazon/chronos-t5-large,3.0239493021840635,0.14814710375646464
-monash_nn5_weekly,amazon/chronos-t5-large,0.9228721852437364,0.08948447200571868
-monash_tourism_monthly,amazon/chronos-t5-large,1.7304427846580348,0.09983169221760163
-monash_tourism_quarterly,amazon/chronos-t5-large,1.6437184365114073,0.0690906057781915
-monash_tourism_yearly,amazon/chronos-t5-large,3.6268503118928535,0.17732007043832695
-monash_traffic,amazon/chronos-t5-large,0.7985975530866148,0.25313515740581755
-monash_weather,amazon/chronos-t5-large,0.8187388457436171,0.1387756772600068
-nn5,amazon/chronos-t5-large,0.5755260854173723,0.15733693855465292
--- a/scripts/evaluation/results/chronos-t5-mini-agg-rel-scores.csv
+++ b/scripts/evaluation/results/chronos-t5-mini-agg-rel-scores.csv
@ -1,5 +0,0 @@
-benchmark,metric,value
-in-domain,MASE,0.7249816823595568
-in-domain,WQL,0.5965372489622094
-zero-shot,MASE,0.8411995116926901
-zero-shot,WQL,0.6888397962259065
--- a/scripts/evaluation/results/chronos-t5-mini-in-domain.csv
+++ b/scripts/evaluation/results/chronos-t5-mini-in-domain.csv
@ -1,16 +0,0 @@
-dataset,model,MASE,WQL
-electricity_15min,amazon/chronos-t5-mini,0.4446629660227641,0.08114657599239496
-m4_daily,amazon/chronos-t5-mini,3.1533349226194005,0.022000507013584743
-m4_hourly,amazon/chronos-t5-mini,0.7616830292996938,0.024630575107847653
-m4_monthly,amazon/chronos-t5-mini,0.9934074425853089,0.10402168689068064
-m4_weekly,amazon/chronos-t5-mini,2.1407189608104416,0.04138058102434373
-monash_electricity_hourly,amazon/chronos-t5-mini,1.3698378948313894,0.09189698159081384
-monash_electricity_weekly,amazon/chronos-t5-mini,1.9238345295706893,0.07015383787479901
-monash_kdd_cup_2018,amazon/chronos-t5-mini,0.6027861468526459,0.25493489598663444
-monash_london_smart_meters,amazon/chronos-t5-mini,0.8570035850603943,0.4356582737588471
-monash_pedestrian_counts,amazon/chronos-t5-mini,0.30374539593979855,0.2374083216051065
-monash_rideshare,amazon/chronos-t5-mini,0.8157349455509949,0.12963515638823117
-monash_temperature_rain,amazon/chronos-t5-mini,1.010161905102516,0.6919171702485583
-taxi_30min,amazon/chronos-t5-mini,0.9318379552979712,0.31229508015999674
-uber_tlc_daily,amazon/chronos-t5-mini,0.9213437323817685,0.10475291429149586
-uber_tlc_hourly,amazon/chronos-t5-mini,0.6812621470377416,0.15982192635434303
--- a/scripts/evaluation/results/chronos-t5-mini-zero-shot.csv
+++ b/scripts/evaluation/results/chronos-t5-mini-zero-shot.csv
@ -1,28 +0,0 @@
-dataset,model,MASE,WQL
-ETTh,amazon/chronos-t5-mini,0.789678971785092,0.08068969536800001
-ETTm,amazon/chronos-t5-mini,0.7521219674190734,0.06791782942706617
-dominick,amazon/chronos-t5-mini,0.8207116999488602,0.34004499734299765
-ercot,amazon/chronos-t5-mini,0.5462749489237783,0.015035001020343136
-exchange_rate,amazon/chronos-t5-mini,2.1326718165798657,0.015073846769933199
-m4_quarterly,amazon/chronos-t5-mini,1.271761811062081,0.08575942238385105
-m4_yearly,amazon/chronos-t5-mini,3.7340853642679126,0.13938781939783162
-m5,amazon/chronos-t5-mini,0.9421556321929742,0.5961689098871504
-monash_australian_electricity,amazon/chronos-t5-mini,1.046297291920238,0.05424453772723559
-monash_car_parts,amazon/chronos-t5-mini,0.8913523483805221,1.0174797526818506
-monash_cif_2016,amazon/chronos-t5-mini,1.0674111822055679,0.016800831829085764
-monash_covid_deaths,amazon/chronos-t5-mini,43.69727825485175,0.08788117644141617
-monash_fred_md,amazon/chronos-t5-mini,0.46227452519609524,0.01871860604459728
-monash_hospital,amazon/chronos-t5-mini,0.7112593459108532,0.05831005112661489
-monash_m1_monthly,amazon/chronos-t5-mini,1.1756557848450433,0.14192178371159841
-monash_m1_quarterly,amazon/chronos-t5-mini,1.795009199698074,0.11760148522768847
-monash_m1_yearly,amazon/chronos-t5-mini,5.078889706085604,0.1882823108615221
-monash_m3_monthly,amazon/chronos-t5-mini,0.900404391663476,0.09935931092075681
-monash_m3_quarterly,amazon/chronos-t5-mini,1.2604342624229292,0.07807204797138119
-monash_m3_yearly,amazon/chronos-t5-mini,3.4395976709464255,0.16085249526114198
-monash_nn5_weekly,amazon/chronos-t5-mini,0.9459117943913629,0.09042762527674755
-monash_tourism_monthly,amazon/chronos-t5-mini,1.920865545569713,0.10791754513335952
-monash_tourism_quarterly,amazon/chronos-t5-mini,1.7957439111869486,0.07514539225156464
-monash_tourism_yearly,amazon/chronos-t5-mini,4.134958090482728,0.2202036957350168
-monash_traffic,amazon/chronos-t5-mini,0.8546792774237857,0.2668831661775284
-monash_weather,amazon/chronos-t5-mini,0.8607748244159247,0.15031866806333247
-nn5,amazon/chronos-t5-mini,0.6497211196906223,0.17352254058241523
--- a/scripts/evaluation/results/chronos-t5-small-agg-rel-scores.csv
+++ b/scripts/evaluation/results/chronos-t5-small-agg-rel-scores.csv
@ -1,5 +0,0 @@
-benchmark,metric,value
-in-domain,MASE,0.7296140269944743
-in-domain,WQL,0.6086958548874499
-zero-shot,MASE,0.8303721909132112
-zero-shot,WQL,0.6649587072099045
--- a/scripts/evaluation/results/chronos-t5-small-in-domain.csv
+++ b/scripts/evaluation/results/chronos-t5-small-in-domain.csv
@ -1,16 +0,0 @@
-dataset,model,MASE,WQL
-electricity_15min,amazon/chronos-t5-small,0.4115559557750193,0.08085148902238105
-m4_daily,amazon/chronos-t5-small,3.1384304946608896,0.02129901023419818
-m4_hourly,amazon/chronos-t5-small,0.7300874075370588,0.024686127211237932
-m4_monthly,amazon/chronos-t5-small,0.9797264456494642,0.10297069145186107
-m4_weekly,amazon/chronos-t5-small,2.0802214537692607,0.03959222330783002
-monash_electricity_hourly,amazon/chronos-t5-small,1.530308399040219,0.10765947926209926
-monash_electricity_weekly,amazon/chronos-t5-small,1.9249616494404531,0.07593976499899265
-monash_kdd_cup_2018,amazon/chronos-t5-small,0.6911172359201715,0.2863722811236367
-monash_london_smart_meters,amazon/chronos-t5-small,0.8405756252443325,0.4300875548402115
-monash_pedestrian_counts,amazon/chronos-t5-small,0.30836963006151696,0.2442543970311678
-monash_rideshare,amazon/chronos-t5-small,0.8436277753840817,0.1363421932158997
-monash_temperature_rain,amazon/chronos-t5-small,1.0176003932416664,0.6847726381172435
-taxi_30min,amazon/chronos-t5-small,0.976277213614167,0.32770172988517626
-uber_tlc_daily,amazon/chronos-t5-small,0.8694727058784919,0.0994889223610958
-uber_tlc_hourly,amazon/chronos-t5-small,0.6738672444888639,0.1573990617753753
--- a/scripts/evaluation/results/chronos-t5-small-zero-shot.csv
+++ b/scripts/evaluation/results/chronos-t5-small-zero-shot.csv
@ -1,28 +0,0 @@
-dataset,model,MASE,WQL
-ETTh,amazon/chronos-t5-small,0.8516754221042285,0.08667817580712385
-ETTm,amazon/chronos-t5-small,0.6825432730635727,0.06076472147001207
-dominick,amazon/chronos-t5-small,0.8108766032127683,0.3368104617474581
-ercot,amazon/chronos-t5-small,0.564879593858422,0.015547628920969682
-exchange_rate,amazon/chronos-t5-small,1.8143459139100264,0.014492477372711763
-m4_quarterly,amazon/chronos-t5-small,1.2415331521819728,0.08383826063189778
-m4_yearly,amazon/chronos-t5-small,3.738749650935195,0.1384514201649314
-m5,amazon/chronos-t5-small,0.9368713240675598,0.5896066252181699
-monash_australian_electricity,amazon/chronos-t5-small,1.2241146217392032,0.06951399165882449
-monash_car_parts,amazon/chronos-t5-small,0.8917508090523597,1.0314986717260015
-monash_cif_2016,amazon/chronos-t5-small,1.0187937383419037,0.014633240218233142
-monash_covid_deaths,amazon/chronos-t5-small,42.298997211368935,0.06339512778191682
-monash_fred_md,amazon/chronos-t5-small,0.4742159923922472,0.01486734736993978
-monash_hospital,amazon/chronos-t5-small,0.709814741753487,0.05704674270057172
-monash_m1_monthly,amazon/chronos-t5-small,1.1723041163998773,0.13799049510465802
-monash_m1_quarterly,amazon/chronos-t5-small,1.8077827825737092,0.11323432989795904
-monash_m1_yearly,amazon/chronos-t5-small,4.739967673537301,0.1730738338876877
-monash_m3_monthly,amazon/chronos-t5-small,0.8856577322724943,0.09985251429658573
-monash_m3_quarterly,amazon/chronos-t5-small,1.278907982396775,0.08094041554590593
-monash_m3_yearly,amazon/chronos-t5-small,3.382470310192457,0.157363937435307
-monash_nn5_weekly,amazon/chronos-t5-small,0.9277396908126303,0.08963913763368506
-monash_tourism_monthly,amazon/chronos-t5-small,1.9251180766131313,0.10943962474253494
-monash_tourism_quarterly,amazon/chronos-t5-small,1.7623454951333655,0.06862432764377493
-monash_tourism_yearly,amazon/chronos-t5-small,3.987690476709746,0.19960492460202509
-monash_traffic,amazon/chronos-t5-small,0.8204223927835267,0.2571189517024486
-monash_weather,amazon/chronos-t5-small,0.8550633590487968,0.1479701971025123
-nn5,amazon/chronos-t5-small,0.6130789183153671,0.16771392719859998
--- a/scripts/evaluation/results/chronos-t5-tiny-agg-rel-scores.csv
+++ b/scripts/evaluation/results/chronos-t5-tiny-agg-rel-scores.csv
@ -1,5 +0,0 @@
-benchmark,metric,value
-in-domain,MASE,0.7649019745781727
-in-domain,WQL,0.6288613368129368
-zero-shot,MASE,0.8704764463925718
-zero-shot,WQL,0.7108912052035352
--- a/scripts/evaluation/results/chronos-t5-tiny-in-domain.csv
+++ b/scripts/evaluation/results/chronos-t5-tiny-in-domain.csv
@ -1,16 +0,0 @@
-dataset,model,MASE,WQL
-electricity_15min,amazon/chronos-t5-tiny,0.5091784254243783,0.08236334376190152
-m4_daily,amazon/chronos-t5-tiny,3.203164895930929,0.022152192084951595
-m4_hourly,amazon/chronos-t5-tiny,0.8171321441164723,0.027490760558343874
-m4_monthly,amazon/chronos-t5-tiny,1.005839207921131,0.10388015368939435
-m4_weekly,amazon/chronos-t5-tiny,2.2148332313370735,0.043429655561156084
-monash_electricity_hourly,amazon/chronos-t5-tiny,1.6190021089002615,0.10967453530956882
-monash_electricity_weekly,amazon/chronos-t5-tiny,2.0774597917676734,0.08159998975612164
-monash_kdd_cup_2018,amazon/chronos-t5-tiny,0.6730886827096076,0.2616610603634618
-monash_london_smart_meters,amazon/chronos-t5-tiny,0.8830447519225436,0.4499607073491794
-monash_pedestrian_counts,amazon/chronos-t5-tiny,0.3042105240185045,0.23387631681117071
-monash_rideshare,amazon/chronos-t5-tiny,0.8431350112476247,0.1378817076926394
-monash_temperature_rain,amazon/chronos-t5-tiny,0.9887398447367799,0.6957797286648015
-taxi_30min,amazon/chronos-t5-tiny,1.035544060665179,0.3450476958104713
-uber_tlc_daily,amazon/chronos-t5-tiny,0.93025919000775,0.1105323649942084
-uber_tlc_hourly,amazon/chronos-t5-tiny,0.697558054147913,0.16320255844336232
--- a/scripts/evaluation/results/chronos-t5-tiny-zero-shot.csv
+++ b/scripts/evaluation/results/chronos-t5-tiny-zero-shot.csv
@ -1,28 +0,0 @@
-dataset,model,MASE,WQL
-ETTh,amazon/chronos-t5-tiny,0.8184074113571701,0.08578203438707048
-ETTm,amazon/chronos-t5-tiny,0.9103621000781905,0.07975361086322658
-dominick,amazon/chronos-t5-tiny,0.8538295532466194,0.3597090770361857
-ercot,amazon/chronos-t5-tiny,0.7273437589773705,0.020843170924006626
-exchange_rate,amazon/chronos-t5-tiny,1.6621128608546154,0.01085145980896454
-m4_quarterly,amazon/chronos-t5-tiny,1.2696259955861924,0.0861404188925996
-m4_yearly,amazon/chronos-t5-tiny,3.5293881164900527,0.13281575565500411
-m5,amazon/chronos-t5-tiny,0.9394059505709506,0.5981531758388589
-monash_australian_electricity,amazon/chronos-t5-tiny,1.4558820561269024,0.07673567331332948
-monash_car_parts,amazon/chronos-t5-tiny,0.9058206654011024,1.0236307963149358
-monash_cif_2016,amazon/chronos-t5-tiny,1.09349564130852,0.014066593076202984
-monash_covid_deaths,amazon/chronos-t5-tiny,46.53079664940016,0.09201919385053775
-monash_fred_md,amazon/chronos-t5-tiny,0.48008374212956456,0.03219550761153211
-monash_hospital,amazon/chronos-t5-tiny,0.7062562198194838,0.05790409320432609
-monash_m1_monthly,amazon/chronos-t5-tiny,1.214892145549996,0.14723095246308077
-monash_m1_quarterly,amazon/chronos-t5-tiny,1.8968576926613199,0.11026972972622998
-monash_m1_yearly,amazon/chronos-t5-tiny,4.829453202075546,0.17286063726000958
-monash_m3_monthly,amazon/chronos-t5-tiny,0.9095746605884618,0.10117875324490073
-monash_m3_quarterly,amazon/chronos-t5-tiny,1.3234957548639883,0.08209032993637215
-monash_m3_yearly,amazon/chronos-t5-tiny,3.1489371074890093,0.1492445630072877
-monash_nn5_weekly,amazon/chronos-t5-tiny,0.9637480731663901,0.09205994784693056
-monash_tourism_monthly,amazon/chronos-t5-tiny,2.151677532807024,0.11356761694754255
-monash_tourism_quarterly,amazon/chronos-t5-tiny,1.9116538900950555,0.07191734222366106
-monash_tourism_yearly,amazon/chronos-t5-tiny,3.820615532600914,0.19709256337364625
-monash_traffic,amazon/chronos-t5-tiny,0.878709088458116,0.2632101606272236
-monash_weather,amazon/chronos-t5-tiny,0.8504899606521996,0.14787595319625085
-nn5,amazon/chronos-t5-tiny,0.7021735456568664,0.19071330483289695
--- a/scripts/evaluation/results/seasonal-naive-in-domain.csv
+++ b/scripts/evaluation/results/seasonal-naive-in-domain.csv
@ -1,16 +0,0 @@
-dataset,model,MASE,WQL
-electricity_15min,seasonal-naive,0.4978697476132387,0.1169378163151378
-m4_daily,seasonal-naive,3.278424323759728,0.0279332664832445
-m4_hourly,seasonal-naive,1.1932105781333862,0.0483091941403194
-m4_monthly,seasonal-naive,1.2597170386001693,0.1455332906092934
-m4_weekly,seasonal-naive,2.777295109814942,0.0633986476090776
-monash_electricity_hourly,seasonal-naive,1.839634785956572,0.1468968206229902
-monash_electricity_weekly,seasonal-naive,3.0371656285424,0.1979332504059267
-monash_kdd_cup_2018,seasonal-naive,0.9943785889052376,0.5555856702439576
-monash_london_smart_meters,seasonal-naive,0.9661872287141056,0.5413187715028914
-monash_pedestrian_counts,seasonal-naive,0.3691951941442247,0.3185271550430794
-monash_rideshare,seasonal-naive,1.2495987545425715,0.1860080644135506
-monash_temperature_rain,seasonal-naive,2.243384627173123,1.4244854980220072
-taxi_30min,seasonal-naive,1.160268631066241,0.4711417890926274
-uber_tlc_daily,seasonal-naive,1.37803447078482,0.2313550175912078
-uber_tlc_hourly,seasonal-naive,0.930916273455971,0.298849044501192
--- a/scripts/evaluation/results/seasonal-naive-zero-shot.csv
+++ b/scripts/evaluation/results/seasonal-naive-zero-shot.csv
@ -1,28 +0,0 @@
-dataset,model,MASE,WQL
-ETTh,seasonal-naive,0.9316203114697056,0.1220896585205886
-ETTm,seasonal-naive,1.1693053852270578,0.1413480385734046
-dominick,seasonal-naive,0.8706150115348875,0.4529164093744346
-ercot,seasonal-naive,0.7613354813741452,0.0366036447606282
-exchange_rate,seasonal-naive,1.7401824286954128,0.0129841406759913
-m4_quarterly,seasonal-naive,1.6022471766126911,0.1186484661559648
-m4_yearly,seasonal-naive,3.974360261259571,0.1614389663357925
-m5,seasonal-naive,1.399206213076729,1.0240883478068443
-monash_australian_electricity,seasonal-naive,1.2533189641227642,0.0836951323308387
-monash_car_parts,seasonal-naive,1.2014638390969912,1.5999522140809177
-monash_cif_2016,seasonal-naive,1.289290577415544,0.0150830409089921
-monash_covid_deaths,seasonal-naive,46.91239825526407,0.1330848762571827
-monash_fred_md,seasonal-naive,1.1008000463101226,0.1222237702571737
-monash_hospital,seasonal-naive,0.9205278266364826,0.0726263373268254
-monash_m1_monthly,seasonal-naive,1.3144614957646543,0.1914632595030148
-monash_m1_quarterly,seasonal-naive,2.077536550805995,0.1495022062865622
-monash_m1_yearly,seasonal-naive,4.894322225232431,0.2092955931101782
-monash_m3_monthly,seasonal-naive,1.1462045758327934,0.1485446007554992
-monash_m3_quarterly,seasonal-naive,1.425343793700714,0.1012520529806161
-monash_m3_yearly,seasonal-naive,3.1717102364409517,0.1665329650420048
-monash_nn5_weekly,seasonal-naive,1.0628482559107015,0.1226908962169196
-monash_tourism_monthly,seasonal-naive,1.630939994944413,0.1041824322151567
-monash_tourism_quarterly,seasonal-naive,1.6989892627474672,0.1193750169177449
-monash_tourism_yearly,seasonal-naive,3.5520097206480883,0.2091826587673241
-monash_traffic,seasonal-naive,1.0767397173107436,0.3618532196990004
-monash_weather,seasonal-naive,1.0038475713182748,0.2165947349654047
-nn5,seasonal-naive,1.2917285866431214,0.4246208074843067
--- a/scripts/experiments/README.md
+++ b/scripts/experiments/README.md
@ -0,0 +1,8 @@
+# Experiments in Real Datasets
+
+An example on how to run experiments with real datasets is shown in [this script](./experiment_with_real_datasets.py).
+
+The configurations per dataset are located in [this yaml file](./configs/datasets.yaml).
+
+### On electricity price forecasting datasets
+At the time of writing our paper we took the electricity price forecasting datasets from [here](https://github.com/Nixtla/transfer-learning-time-series/blob/main/datasets/exogenous-vars-electricity.csv) and [here](https://github.com/Nixtla/transfer-learning-time-series/blob/main/datasets/electricity.csv). As one can see from the commit history [here](https://github.com/Nixtla/transfer-learning-time-series/pull/17), it seems that there were some fixes on how the data was preprocessed. Since the data has been taken from [zenodo](https://zenodo.org/records/4624805), we have provided both versions of these datasets so that anyone interested can run our model with both the previous and corrected dataset versions.
--- a/scripts/experiments/configs/datasets.yaml
+++ b/scripts/experiments/configs/datasets.yaml
@ -0,0 +1,196 @@
+- name: ETTh
+  hf_repo: autogluon/chronos_datasets_extra
+  offset: -24
+  prediction_length: 24
+  num_rolls: 1
+  covariates_fields: ["HULL", "MUFL", "MULL", "LUFL", "LULL"]
+  series_fields: "OT"
+- name: ETTm
+  hf_repo: autogluon/chronos_datasets_extra
+  offset: -24
+  prediction_length: 24
+  num_rolls: 1
+  covariates_fields: ["HULL", "MUFL", "MULL", "LUFL", "LULL"]
+  series_fields: "OT"
+- name: epf_electricity_be_paper
+  hf_repo: autogluon/fev_datasets
+  offset: -24
+  prediction_length: 24
+  num_rolls: 1
+  covariates_fields: ['exogenous1', 'exogenous2']
+  series_fields: "target"
+- name: epf_electricity_de_paper
+  hf_repo: autogluon/fev_datasets
+  offset: -24
+  prediction_length: 24
+  num_rolls: 1
+  covariates_fields: ['exogenous1', 'exogenous2']
+  series_fields: "target"
+- name: epf_electricity_fr_paper
+  hf_repo: autogluon/fev_datasets
+  offset: -24
+  prediction_length: 24
+  num_rolls: 1
+  covariates_fields: ['exogenous1', 'exogenous2']
+  series_fields: "target"
+- name: epf_electricity_np_paper
+  hf_repo: autogluon/fev_datasets
+  offset: -24
+  prediction_length: 24
+  num_rolls: 1
+  covariates_fields: ['exogenous1', 'exogenous2']
+  series_fields: "target"
+- name: epf_electricity_pjm_paper
+  hf_repo: autogluon/fev_datasets
+  offset: -24
+  prediction_length: 24
+  num_rolls: 1
+  covariates_fields: ['exogenous1', 'exogenous2']
+  series_fields: "target"
+- name: proenfo_bull
+  hf_repo: autogluon/fev_datasets
+  offset: -24
+  prediction_length: 24
+  num_rolls: 1
+  covariates_fields: ['airtemperature', 'dewtemperature', 'sealvlpressure']
+  series_fields: "target"
+- name: proenfo_cockatoo
+  hf_repo: autogluon/fev_datasets
+  offset: -24
+  prediction_length: 24
+  num_rolls: 1
+  covariates_fields: ['airtemperature', 'dewtemperature', 'sealvlpressure', 'winddirection', 'windspeed']
+  series_fields: "target"
+- name: proenfo_covid19
+  hf_repo: autogluon/fev_datasets
+  offset: -24
+  prediction_length: 24
+  num_rolls: 1
+  covariates_fields: ['pressure_kpa', 'cloud_cover_perc', 'humidity_perc', 'airtemperature', 'wind_direction_deg', 'wind_speed_kmh']
+  series_fields: "target"
+- name: proenfo_gfc12_load
+  hf_repo: autogluon/fev_datasets
+  offset: -24
+  prediction_length: 24
+  num_rolls: 1
+  covariates_fields: ['airtemperature']
+  series_fields: "target"
+- name: proenfo_gfc14_load
+  hf_repo: autogluon/fev_datasets
+  offset: -24
+  prediction_length: 24
+  num_rolls: 1
+  covariates_fields: ['airtemperature']
+  series_fields: "target"
+- name: proenfo_gfc17_load
+  hf_repo: autogluon/fev_datasets
+  offset: -24
+  prediction_length: 24
+  num_rolls: 1
+  covariates_fields: ['airtemperature']
+  series_fields: "target"
+- name: proenfo_hog
+  hf_repo: autogluon/fev_datasets
+  offset: -24
+  prediction_length: 24
+  num_rolls: 1
+  covariates_fields: ['airtemperature', 'dewtemperature', 'sealvlpressure', 'winddirection', 'windspeed']
+  series_fields: "target"
+- name: proenfo_pdb
+  hf_repo: autogluon/fev_datasets
+  offset: -24
+  prediction_length: 24
+  num_rolls: 1
+  covariates_fields: ['airtemperature']
+  series_fields: "target"
+- name: proenfo_spain
+  hf_repo: autogluon/fev_datasets
+  offset: -24
+  prediction_length: 24
+  num_rolls: 1
+  covariates_fields: [
+    'generation_biomass', 
+    'generation_fossil_brown_coal_lignite', 
+    'generation_fossil_coal_derived_gas', 
+    'generation_fossil_gas', 
+    'generation_fossil_hard_coal', 
+    'generation_fossil_oil', 
+    'generation_fossil_oil_shale',
+    'generation_fossil_peat',
+    'generation_geothermal',
+    'generation_hydro_pumped_storage_consumption',
+    'generation_hydro_run_of_river_and_poundage',
+    'generation_hydro_water_reservoir',
+    'generation_marine',
+    'generation_nuclear',
+    'generation_other',
+    'generation_other_renewable',
+    'generation_solar',
+    'generation_waste',
+    'generation_wind_offshore',
+    'generation_wind_onshore',
+  ]
+  series_fields: "target"
+- name: monash_rideshare
+  hf_repo: autogluon/chronos_datasets
+  offset: -24
+  prediction_length: 24
+  num_rolls: 1
+  covariates_fields: ["price_mean", "distance_mean", "surge_mean", "temp", "rain", "humidity", "clouds", "wind"]
+  series_fields: "api_calls"
+- name: m5_with_covariates
+  hf_repo: autogluon/fev_datasets
+  offset: -28
+  prediction_length: 28
+  num_rolls: 1
+  covariates_fields: [
+    'snap_CA', 
+    'snap_TX', 
+    'snap_WI', 
+    'sell_price', 
+    'event_type_1_Cultural',
+    'event_type_1_National', 
+    'event_type_1_Religious', 
+    'event_type_1_Sporting', 
+    'event_type_1_nan',
+    'event_type_2_Cultural', 
+    'event_type_2_Religious', 
+    'event_type_2_nan',
+   ]
+  series_fields: "target"
+- name: epf_electricity_be
+  hf_repo: autogluon/fev_datasets
+  offset: -24
+  prediction_length: 24
+  num_rolls: 1
+  covariates_fields: ["Generation forecast", "System load forecast"]
+  series_fields: "target"
+- name: epf_electricity_de
+  hf_repo: autogluon/fev_datasets
+  offset: -24
+  prediction_length: 24
+  num_rolls: 1
+  covariates_fields: ["Ampirion Load Forecast", "PV+Wind Forecast"]
+  series_fields: "target"
+- name: epf_electricity_fr
+  hf_repo: autogluon/fev_datasets
+  offset: -24
+  prediction_length: 24
+  num_rolls: 1
+  covariates_fields: ["Generation forecast", "System load forecast"]
+  series_fields: "target"
+- name: epf_electricity_np
+  hf_repo: autogluon/fev_datasets
+  offset: -24
+  prediction_length: 24
+  num_rolls: 1
+  covariates_fields: ["Grid load forecast", "Wind power forecast"]
+  series_fields: "target"
+- name: epf_electricity_pjm
+  hf_repo: autogluon/fev_datasets
+  offset: -24
+  prediction_length: 24
+  num_rolls: 1
+  covariates_fields: ["System load forecast", "Zonal COMED load foecast"]
+  series_fields: "target"
+
--- a/scripts/experiments/configs/synthetic_datasets.yaml
+++ b/scripts/experiments/configs/synthetic_datasets.yaml
@ -0,0 +1,256 @@
+- name: synthetic_datasets/simple_spikes_mult
+  filename: datasets/simple_spikes_mult
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/simple_spikes_add
+  filename: datasets/simple_spikes_add
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/diverse_spikes_mult
+  filename: datasets/diverse_spikes_mult
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/diverse_spikes_add
+  filename: datasets/diverse_spikes_add
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/noisy_spikes_mult
+  filename: datasets/noisy_spikes_mult
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/noisy_spikes_add
+  filename: datasets/noisy_spikes_add
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/simple_steps_mult
+  filename: datasets/simple_steps_mult
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/simple_steps_add
+  filename: datasets/simple_steps_add
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/diverse_steps_mult
+  filename: datasets/diverse_steps_mult
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/diverse_steps_add
+  filename: datasets/diverse_steps_add
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/noisy_steps_mult
+  filename: datasets/noisy_steps_mult
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/noisy_steps_add
+  filename: datasets/noisy_steps_add
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/simple_bells_mult
+  filename: datasets/simple_bells_mult
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/simple_bells_add
+  filename: datasets/simple_bells_add
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/diverse_bells_mult
+  filename: datasets/diverse_bells_mult
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/diverse_bells_add
+  filename: datasets/diverse_bells_add
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/noisy_bells_mult
+  filename: datasets/noisy_bells_mult
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/noisy_bells_add
+  filename: datasets/noisy_bells_add
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/single_bells_mult
+  filename: datasets/single_bells_mult
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/single_bells_add
+  filename: datasets/single_bells_add
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/single_steps_mult
+  filename: datasets/single_steps_mult
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/single_steps_add
+  filename: datasets/single_steps_add
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/single_spikes_mult
+  filename: datasets/single_spikes_mult
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/single_spikes_add
+  filename: datasets/single_spikes_add
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/single_arp_mult
+  filename: datasets/single_arp_mult
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/single_arp_add
+  filename: datasets/single_arp_add
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/simple_arp_mult
+  filename: datasets/simple_arp_mult
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/simple_arp_add
+  filename: datasets/simple_arp_add
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/diverse_arp_mult
+  filename: datasets/diverse_arp_mult
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/diverse_arp_add
+  filename: datasets/diverse_arp_add
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/noisy_arp_mult
+  filename: datasets/noisy_arp_mult
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
+- name: synthetic_datasets/noisy_arp_add
+  filename: datasets/noisy_arp_add
+  offset: -30
+  prediction_length: 30
+  num_rolls: 1
+  covariates_fields:
+  - covariate
+  series_fields: target
--- a/scripts/experiments/example.py
+++ b/scripts/experiments/example.py
@ -0,0 +1,69 @@
+import numpy as np
+import yaml
+
+from chronosx.chronosx import ChronosXPipeline
+from chronosx.utils import ChronosDataset, load_and_split_dataset
+from gluonts.ev.metrics import MASE, MeanWeightedSumQuantileLoss
+from gluonts.model.evaluation import evaluate_forecasts
+from pathlib import Path
+
+
+config_path = "./configs/datasets.yaml"
+output_dir = Path(f"../../output/finetune")
+output_dir.mkdir(exist_ok=True, parents=True)
+
+with open(config_path) as fp:
+    backtest_configs = yaml.safe_load(fp)
+
+dataset_config = backtest_configs[0]
+prediction_length = dataset_config["prediction_length"]
+num_covariates = 2 * len(dataset_config["covariates_fields"])
+
+
+# Load Chronos
+pipeline = ChronosXPipeline(
+    prediction_length=prediction_length,
+    num_covariates=num_covariates,
+)
+
+# load Dataset
+train_dataset, test_dataset = load_and_split_dataset(backtest_config=dataset_config)
+quantized_train_dataset = ChronosDataset(
+    datasets=[train_dataset],
+    probabilities=[1.0],
+    tokenizer=pipeline.tokenizer,
+    prediction_length=prediction_length,
+    mode="training",
+).shuffle()
+
+# fine tune model
+_, save_model_path = pipeline.finetune(
+    output_dir,
+    quantized_train_dataset,
+    skip_pretrained_validation=True,
+)
+
+# Evaluate fine tuned model
+pipeline = ChronosXPipeline(
+    prediction_length=prediction_length,
+    num_covariates=num_covariates,
+    pretrained_model_name_or_path=output_dir / "final-checkpoint",
+)
+
+pipeline.chronosx.eval()
+forecasts = pipeline.generate_forecasts(test_dataset.input)
+
+metrics = (
+    evaluate_forecasts(
+        forecasts,
+        test_data=test_dataset,
+        metrics=[
+            MASE(),
+            MeanWeightedSumQuantileLoss(np.arange(0.05, 1, 0.05).round(2).tolist()),
+        ],
+    )
+    .reset_index(drop=True)
+    .to_dict(orient="records")
+)
+
+print(metrics)
--- a/scripts/experiments/experiment_with_real_datasets.py
+++ b/scripts/experiments/experiment_with_real_datasets.py
@ -0,0 +1,204 @@
+import numpy as np
+import pandas as pd
+import random
+import time
+import yaml
+
+from chronosx.chronosx import ChronosXPipeline
+from chronosx.utils.chronos_dataset import ChronosDataset
+from chronosx.utils.hf_data_loader import load_and_split_dataset
+from chronosx.utils.utils import has_enough_observations
+from functools import partial
+from gluonts.dataset.split import split
+from gluonts.ev.metrics import MASE, MeanWeightedSumQuantileLoss
+from gluonts.itertools import Filter
+from gluonts.model.evaluation import evaluate_forecasts
+from pathlib import Path
+from pprint import pprint
+
+
+covariate_injection = "IIB+OIB"
+chronos_model_id = "amazon/chronos-t5-small"
+config_path = "./configs/datasets.yaml"
+output_dir = Path("../../output/finetune")
+output_metrics_dir = Path("../../output/metrics")
+output_metrics_dir.mkdir(exist_ok=True, parents=True)
+
+min_past = 1
+shuffle_buffer_length = 100
+num_runs = 3
+max_steps = 5000
+skip_pretrained_validation = False
+learning_rate_list = [0.01, 0.001, 0.0001]
+
+with open(config_path) as fp:
+    backtest_configs = yaml.safe_load(fp)
+
+dataset_config = backtest_configs[0]
+dataset_name = dataset_config["name"]
+prediction_length = dataset_config["prediction_length"]
+num_covariates = 2 * len(dataset_config["covariates_fields"])
+train_dataset, test_dataset = load_and_split_dataset(backtest_config=dataset_config)
+
+
+# Load Chronos
+pipeline = ChronosXPipeline(
+    prediction_length=prediction_length,
+    num_covariates=num_covariates,
+    covariate_injection=covariate_injection,
+    pretrained_model_name_or_path=chronos_model_id,
+)
+
+tokenizer = pipeline.tokenizer
+
+
+train_dataset = Filter(
+    partial(
+        has_enough_observations,
+        min_length=min_past + 2 * prediction_length,  # for validation and for train
+        max_missing_prop=0.5,
+    ),
+    train_dataset,
+)
+
+quantized_val_dataset = ChronosDataset(
+    datasets=[train_dataset],
+    probabilities=[1.0],
+    tokenizer=tokenizer,
+    prediction_length=prediction_length,
+    min_past=min_past,
+    mode="validation",
+)
+a = list(quantized_val_dataset)
+
+train_dataset, _ = split(train_dataset, offset=-prediction_length)
+quantized_train_dataset = ChronosDataset(
+    datasets=[train_dataset],
+    probabilities=[1.0],
+    tokenizer=tokenizer,
+    prediction_length=prediction_length,
+    min_past=min_past,
+    mode="training",
+).shuffle(shuffle_buffer_length=shuffle_buffer_length)
+
+# zero-shot evaluation on validation set
+val_of_zero_shot_pretrained_model = pipeline.evaluate_model_on_validation_set(
+    covariate_injection=None,
+    quantized_val_dataset=quantized_val_dataset,
+    output_dir=output_dir / dataset_name,
+)
+
+
+random.seed(int(time.time()))
+seed = random.randint(0, 2**32)  
+val_loss_per_lr = {}
+mean_val_loss_per_lr = {}
+
+for lr in learning_rate_list:
+    lr_str = str(lr)
+    val_loss_per_run = []
+    model_paths = []
+    for run_id in range(num_runs):
+        run_id_str = str(run_id)
+        output_dir_lr_run_id = Path(
+            output_dir / dataset_name / f"lr={lr}" / f"run_id={run_id}"
+        )
+        output_dir_lr_run_id.mkdir(exist_ok=True, parents=True)
+
+        quantized_train_dataset = ChronosDataset(
+            datasets=[train_dataset],
+            probabilities=[1.0],
+            tokenizer=tokenizer,
+            prediction_length=prediction_length,
+            min_past=min_past,
+            mode="training",
+        ).shuffle(
+            shuffle_buffer_length=shuffle_buffer_length,
+            random_seed=run_id + seed,
+        )
+
+        pipeline = ChronosXPipeline(
+            prediction_length=prediction_length,
+            num_covariates=num_covariates,
+            covariate_injection=covariate_injection,
+            pretrained_model_name_or_path=chronos_model_id,
+        )
+
+        val_loss, save_model_path = pipeline.finetune(
+            output_dir_lr_run_id,
+            quantized_train_dataset,
+            lr=lr,
+            quantized_val_dataset=quantized_val_dataset,
+            skip_pretrained_validation=skip_pretrained_validation,
+            max_steps=max_steps,
+        )
+
+        val_loss_per_run.append(val_loss)
+        model_paths.append(save_model_path)
+
+    mean_val_loss = np.mean(val_loss_per_run)
+    val_loss_per_lr[f"{lr}"] = val_loss_per_run
+    mean_val_loss_per_lr[f"{lr}"] = mean_val_loss
+    print(
+        f"lr: {lr} - val_loss_per_run: {val_loss_per_run} - mean_val_loss: {mean_val_loss}"
+    )
+
+
+print("val_loss_per_lr")
+pprint(val_loss_per_lr)
+
+print("mean_val_loss_per_lr")
+pprint(mean_val_loss_per_lr)
+
+best_val_loss = np.inf
+for lr_str, val_loss in mean_val_loss_per_lr.items():
+    if val_loss < best_val_loss:
+        best_val_loss = val_loss
+        best_lr_str = lr_str
+
+
+# evaluate model with best mean validation loss
+lr = float(best_lr_str)
+result_rows = []
+for run_id in range(num_runs):
+    pretrained_model_path = (
+        output_dir / dataset_name / f"lr={lr}" / f"run_id={run_id}" / "final-checkpoint"
+    )
+    pipeline = ChronosXPipeline(
+        prediction_length=prediction_length,
+        num_covariates=num_covariates,
+        covariate_injection=covariate_injection,
+        pretrained_model_name_or_path=pretrained_model_path,
+    )
+
+    pipeline.chronosx.eval()
+    forecasts = pipeline.generate_forecasts(
+        test_dataset.input,
+    )
+
+    metrics = (
+        evaluate_forecasts(
+            forecasts,
+            test_data=test_dataset,
+            metrics=[
+                MASE(),
+                MeanWeightedSumQuantileLoss(np.arange(0.05, 1, 0.05).round(2).tolist()),
+            ],
+            batch_size=5000,
+        )
+        .reset_index(drop=True)
+        .to_dict(orient="records")
+    )
+
+    result_rows.append(
+        {
+            "dataset": dataset_config["name"],
+            "covariate_injection": covariate_injection,
+            **metrics[0],
+        }
+    )
+
+df = pd.DataFrame(result_rows)
+df = df.set_index(["dataset", "covariate_injection"])
+pprint(df.mean())
+df.to_csv(output_metrics_dir / f"{dataset_name}_max_steps={max_steps}.csv")
--- a/scripts/experiments/synthetic_datasets/README.md
+++ b/scripts/experiments/synthetic_datasets/README.md
@ -0,0 +1,6 @@
+# Experiments with Synthetic Datasets
+
+Synthetic Datasets must be generated by executing this [script](./generate_synthetic_datasets.py)
+
+Once the synthetic datasets are generated, experiments can be executed with this [script](./experiment_with_synthetic_datasets.py)
+
--- a/scripts/experiments/synthetic_datasets/constants.py
+++ b/scripts/experiments/synthetic_datasets/constants.py
@ -0,0 +1,222 @@
+NUM_YEARS = 5
+START_DATE = "01/01/2024"
+END_DATE = "12/31/2028"
+NUM_EVENTS_PER_YEAR = 25
+DATASETS_INFO = {
+    "simple_spikes_mult": {
+        "function_covariates": "steps",
+        "op": "mult",
+        "function_target": "simple",
+        "delta": 2,
+        "num_events": 4 * NUM_EVENTS_PER_YEAR * NUM_YEARS,
+        "fixed_event_pos": 15,
+    },
+    "simple_spikes_add": {
+        "function_covariates": "steps",
+        "op": "add",
+        "delta": 2,
+        "function_target": "simple",
+        "num_events": 4 * NUM_EVENTS_PER_YEAR * NUM_YEARS,
+        "fixed_event_pos": 15,
+    },
+    "diverse_spikes_mult": {
+        "function_covariates": "steps",
+        "op": "mult",
+        "delta": 2,
+        "function_target": "diverse",
+        "num_events": 4 * NUM_EVENTS_PER_YEAR * NUM_YEARS,
+        "fixed_event_pos": 15,
+    },
+    "diverse_spikes_add": {
+        "function_covariates": "steps",
+        "op": "add",
+        "delta": 2,
+        "function_target": "diverse",
+        "num_events": 4 * NUM_EVENTS_PER_YEAR * NUM_YEARS,
+        "fixed_event_pos": 15,
+    },
+    "noisy_spikes_mult": {
+        "function_covariates": "steps",
+        "op": "mult",
+        "delta": 2,
+        "function_target": "noisy",
+        "num_events": 4 * NUM_EVENTS_PER_YEAR * NUM_YEARS,
+        "fixed_event_pos": 15,
+    },
+    "noisy_spikes_add": {
+        "function_covariates": "steps",
+        "op": "add",
+        "delta": 2,
+        "function_target": "noisy",
+        "num_events": 4 * NUM_EVENTS_PER_YEAR * NUM_YEARS,
+        "fixed_event_pos": 15,
+    },
+    "simple_steps_mult": {
+        "function_covariates": "steps",
+        "op": "mult",
+        "function_target": "simple",
+        "delta": 15,
+        "num_events": NUM_EVENTS_PER_YEAR * NUM_YEARS,
+    },
+    "simple_steps_add": {
+        "function_covariates": "steps",
+        "op": "add",
+        "function_target": "simple",
+        "delta": 15,
+        "num_events": NUM_EVENTS_PER_YEAR * NUM_YEARS,
+    },
+    "diverse_steps_mult": {
+        "function_covariates": "steps",
+        "op": "mult",
+        "function_target": "diverse",
+        "delta": 15,
+        "num_events": NUM_EVENTS_PER_YEAR * NUM_YEARS,
+    },
+    "diverse_steps_add": {
+        "function_covariates": "steps",
+        "op": "add",
+        "function_target": "diverse",
+        "delta": 15,
+        "num_events": NUM_EVENTS_PER_YEAR * NUM_YEARS,
+    },
+    "noisy_steps_mult": {
+        "function_covariates": "steps",
+        "op": "mult",
+        "function_target": "noisy",
+        "delta": 15,
+        "num_events": NUM_EVENTS_PER_YEAR * NUM_YEARS,
+    },
+    "noisy_steps_add": {
+        "function_covariates": "steps",
+        "op": "add",
+        "function_target": "noisy",
+        "delta": 15,
+        "num_events": NUM_EVENTS_PER_YEAR * NUM_YEARS,
+    },
+    "simple_bells_mult": {
+        "function_covariates": "bells",
+        "op": "mult",
+        "function_target": "simple",
+        "num_events": NUM_EVENTS_PER_YEAR * NUM_YEARS,
+        "fixed_event_pos": 15,
+    },
+    "simple_bells_add": {
+        "function_covariates": "bells",
+        "op": "add",
+        "function_target": "simple",
+        "num_events": NUM_EVENTS_PER_YEAR * NUM_YEARS,
+        "fixed_event_pos": 15,
+    },
+    "diverse_bells_mult": {
+        "function_covariates": "bells",
+        "op": "mult",
+        "function_target": "diverse",
+        "num_events": NUM_EVENTS_PER_YEAR * NUM_YEARS,
+        "fixed_event_pos": 15,
+    },
+    "diverse_bells_add": {
+        "function_covariates": "bells",
+        "op": "add",
+        "function_target": "diverse",
+        "num_events": NUM_EVENTS_PER_YEAR * NUM_YEARS,
+        "fixed_event_pos": 15,
+    },
+    "noisy_bells_mult": {
+        "function_covariates": "bells",
+        "op": "mult",
+        "function_target": "noisy",
+        "num_events": NUM_EVENTS_PER_YEAR * NUM_YEARS,
+        "fixed_event_pos": 15,
+    },
+    "noisy_bells_add": {
+        "function_covariates": "bells",
+        "op": "add",
+        "function_target": "noisy",
+        "num_events": NUM_EVENTS_PER_YEAR * NUM_YEARS,
+        "fixed_event_pos": 15,
+    },
+    "single_bells_mult": {
+        "function_covariates": "bells",
+        "op": "mult",
+        "function_target": "single",
+        "num_events": NUM_EVENTS_PER_YEAR * NUM_YEARS,
+        "fixed_event_pos": 15,
+    },
+    "single_bells_add": {
+        "function_covariates": "bells",
+        "op": "add",
+        "function_target": "single",
+        "num_events": NUM_EVENTS_PER_YEAR * NUM_YEARS,
+        "fixed_event_pos": 15,
+    },
+    "single_steps_mult": {
+        "function_covariates": "steps",
+        "op": "mult",
+        "function_target": "single",
+        "delta": 15,
+        "num_events": NUM_EVENTS_PER_YEAR * NUM_YEARS,
+    },
+    "single_steps_add": {
+        "function_covariates": "steps",
+        "op": "add",
+        "function_target": "single",
+        "delta": 15,
+        "num_events": NUM_EVENTS_PER_YEAR * NUM_YEARS,
+    },
+    "single_spikes_mult": {
+        "function_covariates": "steps",
+        "op": "mult",
+        "function_target": "single",
+        "delta": 2,
+        "num_events": 4 * NUM_EVENTS_PER_YEAR * NUM_YEARS,
+        "fixed_event_pos": 15,
+    },
+    "single_spikes_add": {
+        "function_covariates": "steps",
+        "op": "add",
+        "function_target": "single",
+        "delta": 2,
+        "num_events": 4 * NUM_EVENTS_PER_YEAR * NUM_YEARS,
+        "fixed_event_pos": 15,
+    },
+    "single_arp_mult": {
+        "function_covariates": "arp",
+        "op": "mult",
+        "function_target": "single",
+    },
+    "single_arp_add": {
+        "function_covariates": "arp",
+        "op": "add",
+        "function_target": "single",
+    },
+    "simple_arp_mult": {
+        "function_covariates": "arp",
+        "op": "mult",
+        "function_target": "simple",
+    },
+    "simple_arp_add": {
+        "function_covariates": "arp",
+        "op": "add",
+        "function_target": "simple",
+    },
+    "diverse_arp_mult": {
+        "function_covariates": "arp",
+        "op": "mult",
+        "function_target": "diverse",
+    },
+    "diverse_arp_add": {
+        "function_covariates": "arp",
+        "op": "add",
+        "function_target": "diverse",
+    },
+    "noisy_arp_mult": {
+        "function_covariates": "arp",
+        "op": "mult",
+        "function_target": "noisy",
+    },
+    "noisy_arp_add": {
+        "function_covariates": "arp",
+        "op": "add",
+        "function_target": "noisy",
+    },
+}
--- a/scripts/experiments/synthetic_datasets/covariates.py
+++ b/scripts/experiments/synthetic_datasets/covariates.py
@ -0,0 +1,101 @@
+import numpy as np
+from gluonts.dataset.artificial import recipe as rcp
+
+
+def make_covariate_gen(covariate_func):
+    def covariate_gen(series, op, function_target, **kwargs):
+        covariates = covariate_func(
+            series=series, op=op, function_target=function_target, **kwargs
+        )
+
+        if op == "add":
+            return covariates + series, covariates
+        elif op == "mult":
+            return np.multiply(covariates, series), covariates
+        else:
+            raise NameError()
+
+    return covariate_gen
+
+
+@make_covariate_gen
+def steps(
+    series: np.ndarray,
+    op: str,
+    function_target: str,
+    num_events: int,
+    delta: int,
+    fixed_event_pos: int = None,
+) -> np.ndarray:
+
+    length = len(series)
+    step_pos = np.random.randint(0, length, num_events)
+
+    if fixed_event_pos is not None:
+        step_pos[-1] = length - fixed_event_pos
+
+    step_delta = np.random.randint(1, delta, num_events)
+    steps_series = np.ones(length)
+    gamma = get_gamma(series, function_target, op)
+
+    for p, delta in zip(step_pos, step_delta):
+        steps_series[p : p + delta] = gamma
+
+    return steps_series
+
+
+@make_covariate_gen
+def bells(
+    series: np.ndarray,
+    op: str,
+    function_target: str,
+    num_events: int,
+    max_sigma: int = 15,
+    fixed_event_pos: int = None,
+):
+
+    length = len(series)
+    change_pos = np.random.randint(0, length, num_events)
+    sigma_change = np.random.uniform(1, max_sigma, num_events).reshape(-1, 1)
+
+    if fixed_event_pos is not None:
+        change_pos[-1] = length - fixed_event_pos
+        sigma_change[-1] = 7
+
+    gamma = get_gamma(series, function_target, op)
+    t = np.tile(np.arange(length).reshape(1, -1), (num_events, 1))
+    change_pos = np.tile(change_pos.reshape(-1, 1), (1, length))
+
+    bells_series = gamma * (
+        np.exp(-np.divide(t - change_pos, sigma_change) ** 2).sum(axis=0) + 1
+    )
+
+    return bells_series
+
+
+@make_covariate_gen
+def arp(series: np.ndarray, op: str, function_target: str):
+    phi1 = np.random.uniform(0, 1)
+    phi2 = 1 - phi1
+
+    # AutoRegressive Process
+    target = rcp.ARp(phi=[phi1, phi2], sigma=1)
+    x = rcp.evaluate(dict(target=target), length=len(series))["target"]
+    x = (x - x.min()) / (x.max() - x.min())
+
+    gamma = get_gamma(series, function_target, op)
+    arp_series = gamma * x
+
+    return arp_series
+
+
+def get_gamma(series: np.ndarray, function_target: str, op: str, level: int = 5):
+
+    gamma = (
+        level
+        if function_target in ["simple", "single"]
+        else np.random.uniform(1, level)
+    )
+    scale = 1 if op == "mult" else np.mean(np.abs(series))
+    gamma *= scale
+    return gamma
--- a/scripts/experiments/synthetic_datasets/experiment_with_synthetic_datasets.py
+++ b/scripts/experiments/synthetic_datasets/experiment_with_synthetic_datasets.py
@ -0,0 +1,213 @@
+import gc
+import numpy as np
+import pandas as pd
+import random
+import shutil
+import time
+import torch
+import yaml
+
+
+from chronosx.chronosx import ChronosXPipeline
+from chronosx.utils.chronos_dataset import ChronosDataset
+from chronosx.utils.hf_data_loader import load_and_split_dataset
+from chronosx.utils.utils import has_enough_observations
+from functools import partial
+from gluonts.dataset.split import split
+from gluonts.ev.metrics import MASE, MeanWeightedSumQuantileLoss
+from gluonts.itertools import Filter
+from gluonts.model.evaluation import evaluate_forecasts
+from pathlib import Path
+from pprint import pprint
+
+
+covariate_injection = "IIB+OIB"
+chronos_model_id = "amazon/chronos-t5-small"
+config_path = "../configs/synthetic_datasets.yaml"
+output_dir = Path("../../../output_synthetic_experiments/finetune")
+output_metrics_dir = Path("../../../output_synthetic_experiments/metrics")
+output_metrics_dir.mkdir(exist_ok=True, parents=True)
+
+min_past = 1
+shuffle_buffer_length = 100
+num_runs = 3
+max_steps = 5000
+skip_pretrained_validation = False
+learning_rate_list = [0.01, 0.001, 0.0001]
+
+with open(config_path) as fp:
+    backtest_configs = yaml.safe_load(fp)
+
+for dataset_config in backtest_configs:
+    dataset_name = dataset_config["name"]
+    prediction_length = dataset_config["prediction_length"]
+    num_covariates = 2 * len(dataset_config["covariates_fields"])
+    train_dataset, test_dataset = load_and_split_dataset(backtest_config=dataset_config)
+
+    # Load Chronos
+    pipeline = ChronosXPipeline(
+        prediction_length=prediction_length,
+        num_covariates=num_covariates,
+        covariate_injection=covariate_injection,
+        pretrained_model_name_or_path=chronos_model_id,
+    )
+
+    tokenizer = pipeline.tokenizer
+
+    train_dataset = Filter(
+        partial(
+            has_enough_observations,
+            min_length=min_past + 2 * prediction_length,  # for validation and for train
+            max_missing_prop=0.5,
+        ),
+        train_dataset,
+    )
+
+    quantized_val_dataset = ChronosDataset(
+        datasets=[train_dataset],
+        probabilities=[1.0],
+        tokenizer=tokenizer,
+        prediction_length=prediction_length,
+        min_past=min_past,
+        mode="validation",
+    )
+
+    train_dataset, _ = split(train_dataset, offset=-prediction_length)
+    quantized_train_dataset = ChronosDataset(
+        datasets=[train_dataset],
+        probabilities=[1.0],
+        tokenizer=tokenizer,
+        prediction_length=prediction_length,
+        min_past=min_past,
+        mode="training",
+    ).shuffle(shuffle_buffer_length=shuffle_buffer_length)
+
+    # zero-shot evaluation on validation set
+    val_of_zero_shot_pretrained_model = pipeline.evaluate_model_on_validation_set(
+        covariate_injection=None,
+        quantized_val_dataset=quantized_val_dataset,
+        output_dir=output_dir / dataset_name,
+    )
+
+    random.seed(int(time.time()))
+    seed = random.randint(0, 2**32)
+    val_loss_per_lr = {}
+    mean_val_loss_per_lr = {}
+
+    for lr in learning_rate_list:
+        lr_str = str(lr)
+        val_loss_per_run = []
+        model_paths = []
+        for run_id in range(num_runs):
+            run_id_str = str(run_id)
+            output_dir_lr_run_id = Path(
+                output_dir / dataset_name / f"lr={lr}" / f"run_id={run_id}"
+            )
+            output_dir_lr_run_id.mkdir(exist_ok=True, parents=True)
+
+            quantized_train_dataset = ChronosDataset(
+                datasets=[train_dataset],
+                probabilities=[1.0],
+                tokenizer=tokenizer,
+                prediction_length=prediction_length,
+                min_past=min_past,
+                mode="training",
+            ).shuffle(
+                shuffle_buffer_length=shuffle_buffer_length,
+                random_seed=run_id + seed,
+            )
+
+            pipeline = ChronosXPipeline(
+                prediction_length=prediction_length,
+                num_covariates=num_covariates,
+                covariate_injection=covariate_injection,
+                pretrained_model_name_or_path=chronos_model_id,
+            )
+
+            val_loss, save_model_path = pipeline.finetune(
+                output_dir_lr_run_id,
+                quantized_train_dataset,
+                lr=lr,
+                quantized_val_dataset=quantized_val_dataset,
+                skip_pretrained_validation=skip_pretrained_validation,
+                max_steps=max_steps,
+            )
+
+            val_loss_per_run.append(val_loss)
+            model_paths.append(save_model_path)
+
+        mean_val_loss = np.mean(val_loss_per_run)
+        val_loss_per_lr[f"{lr}"] = val_loss_per_run
+        mean_val_loss_per_lr[f"{lr}"] = mean_val_loss
+        print(
+            f"lr: {lr} - val_loss_per_run: {val_loss_per_run} - mean_val_loss: {mean_val_loss}"
+        )
+
+    print("val_loss_per_lr")
+    pprint(val_loss_per_lr)
+
+    print("mean_val_loss_per_lr")
+    pprint(mean_val_loss_per_lr)
+
+    best_val_loss = np.inf
+    for lr_str, val_loss in mean_val_loss_per_lr.items():
+        if val_loss < best_val_loss:
+            best_val_loss = val_loss
+            best_lr_str = lr_str
+
+    # evaluate model with best mean validation loss
+    lr = float(best_lr_str)
+    result_rows = []
+    for run_id in range(num_runs):
+        pretrained_model_path = (
+            output_dir
+            / dataset_name
+            / f"lr={lr}"
+            / f"run_id={run_id}"
+            / "final-checkpoint"
+        )
+        pipeline = ChronosXPipeline(
+            prediction_length=prediction_length,
+            num_covariates=num_covariates,
+            covariate_injection=covariate_injection,
+            pretrained_model_name_or_path=pretrained_model_path,
+        )
+
+        pipeline.chronosx.eval()
+        forecasts = pipeline.generate_forecasts(
+            test_dataset.input,
+        )
+
+        metrics = (
+            evaluate_forecasts(
+                forecasts,
+                test_data=test_dataset,
+                metrics=[
+                    MASE(),
+                    MeanWeightedSumQuantileLoss(
+                        np.arange(0.05, 1, 0.05).round(2).tolist()
+                    ),
+                ],
+                batch_size=5000,
+            )
+            .reset_index(drop=True)
+            .to_dict(orient="records")
+        )
+
+        result_rows.append(
+            {
+                "dataset": dataset_config["name"],
+                "covariate_injection": covariate_injection,
+                **metrics[0],
+            }
+        )
+
+    df = pd.DataFrame(result_rows)
+    df = df.set_index(["dataset", "covariate_injection"])
+    pprint(df.mean())
+    df.to_csv(output_metrics_dir / f"{dataset_name.split('/')[-1]}_max_steps={max_steps}.csv")
+
+    shutil.rmtree(output_dir / dataset_name)
+    del pipeline, quantized_train_dataset, quantized_val_dataset, forecasts
+    gc.collect()
+    torch.cuda.empty_cache()
--- a/scripts/experiments/synthetic_datasets/generate_synthetic_datasets.py
+++ b/scripts/experiments/synthetic_datasets/generate_synthetic_datasets.py
@ -0,0 +1,62 @@
+from pathlib import Path
+import datasets
+import numpy as np
+import os
+import pandas as pd
+
+from constants import START_DATE, END_DATE, DATASETS_INFO
+from covariates import steps, bells, arp
+from target import single, simple, diverse, noisy
+
+
+def make_positive(x):
+    return x - x.min()
+
+
+num_entries = 100
+
+
+def generate():
+    output_dir = Path("./datasets")
+    output_dir.mkdir(exist_ok=True, parents=True)
+
+    timestamp = np.array(
+        pd.date_range(
+            start=pd.to_datetime(START_DATE), end=pd.to_datetime(END_DATE), freq="D"
+        )
+    )
+
+    for dataset_name, dataset_info in DATASETS_INFO.items():
+
+        print(f"dataset_name:{dataset_name}")
+        function_target = dataset_info.pop("function_target")
+        function_covariates = dataset_info.pop("function_covariates")
+        operator = dataset_info.pop("op")
+
+        target_generator = eval(function_target)
+        covariate_generator = eval(function_covariates)
+
+        entries = []
+        for _ in range(num_entries):
+            target_before_covariates = target_generator(L=len(timestamp))
+            target_after_covariates, covariate = covariate_generator(
+                series=target_before_covariates,
+                op=operator,
+                function_target=function_target,
+                **dataset_info,
+            )
+            target_after_covariates = make_positive(target_after_covariates)
+            entry = {
+                "timestamp": timestamp,
+                "target": target_after_covariates,
+                "covariate": covariate,
+            }
+            entries.append(entry)
+
+        dest_file = os.path.join(output_dir, dataset_name)
+        ds = datasets.Dataset.from_list(entries)
+        ds.save_to_disk(dest_file)
+
+
+if __name__ == "__main__":
+    generate()
--- a/scripts/experiments/synthetic_datasets/target.py
+++ b/scripts/experiments/synthetic_datasets/target.py
@ -0,0 +1,53 @@
+import numpy as np
+
+
+def seasonal(t):
+    a1, a2, a3 = np.random.uniform(-5, 5, 3)  # amplitude
+    b1, b2, b3 = np.random.uniform(0, 1, 3)  # phase
+    series = (
+        a1 * np.sin(2 * np.pi * (t / 7 + b1))
+        + a2 * np.sin(2 * np.pi * (t / 30 + b2))
+        + a3 * np.sin(2 * np.pi * (t / 365 + b3))
+    )
+    return series
+
+
+def trend(t):
+    a4, a5 = np.random.uniform(-1, 1, 2)
+    series = a4 + a5 * t / 365
+    return series
+
+
+def noise(y):
+    return np.random.normal(0, 0.25 * np.mean(np.abs(y)), len(y))
+
+
+def simple_seasonal(t):
+    a1 = np.random.uniform(-5, 5, 1)
+    series = a1 * np.sin(2 * np.pi * t / 7)
+    return series
+
+
+def simple(L):
+    t = np.arange(L)
+    y = simple_seasonal(t)
+    return y
+
+
+def single(L):
+    t = np.arange(L)
+    y = np.sin(2 * np.pi * t / 7)
+    return y
+
+
+def diverse(L):
+    t = np.arange(L)
+    y = seasonal(t) + trend(t)
+    return y
+
+
+def noisy(L):
+    t = np.arange(L)
+    y = seasonal(t) + trend(t)
+    y += noise(y)
+    return y
--- a/scripts/kernel-synth.py
+++ b/scripts/kernel-synth.py
@ -1,200 +0,0 @@
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-import argparse
-import functools
-from pathlib import Path
-from typing import Optional
-
-import numpy as np
-from gluonts.dataset.arrow import ArrowWriter
-from joblib import Parallel, delayed
-from sklearn.gaussian_process import GaussianProcessRegressor
-from sklearn.gaussian_process.kernels import (
-    RBF,
-    ConstantKernel,
-    DotProduct,
-    ExpSineSquared,
-    Kernel,
-    RationalQuadratic,
-    WhiteKernel,
-)
-from tqdm.auto import tqdm
-
-LENGTH = 1024
-KERNEL_BANK = [
-    ExpSineSquared(periodicity=24 / LENGTH),  # H
-    ExpSineSquared(periodicity=48 / LENGTH),  # 0.5H
-    ExpSineSquared(periodicity=96 / LENGTH),  # 0.25H
-    ExpSineSquared(periodicity=24 * 7 / LENGTH),  # H
-    ExpSineSquared(periodicity=48 * 7 / LENGTH),  # 0.5H
-    ExpSineSquared(periodicity=96 * 7 / LENGTH),  # 0.25H
-    ExpSineSquared(periodicity=7 / LENGTH),  # D
-    ExpSineSquared(periodicity=14 / LENGTH),  # 0.5D
-    ExpSineSquared(periodicity=30 / LENGTH),  # D
-    ExpSineSquared(periodicity=60 / LENGTH),  # 0.5D
-    ExpSineSquared(periodicity=365 / LENGTH),  # D
-    ExpSineSquared(periodicity=365 * 2 / LENGTH),  # 0.5D
-    ExpSineSquared(periodicity=4 / LENGTH),  # W
-    ExpSineSquared(periodicity=26 / LENGTH),  # W
-    ExpSineSquared(periodicity=52 / LENGTH),  # W
-    ExpSineSquared(periodicity=4 / LENGTH),  # M
-    ExpSineSquared(periodicity=6 / LENGTH),  # M
-    ExpSineSquared(periodicity=12 / LENGTH),  # M
-    ExpSineSquared(periodicity=4 / LENGTH),  # Q
-    ExpSineSquared(periodicity=4 * 10 / LENGTH),  # Q
-    ExpSineSquared(periodicity=10 / LENGTH),  # Y
-    DotProduct(sigma_0=0.0),
-    DotProduct(sigma_0=1.0),
-    DotProduct(sigma_0=10.0),
-    RBF(length_scale=0.1),
-    RBF(length_scale=1.0),
-    RBF(length_scale=10.0),
-    RationalQuadratic(alpha=0.1),
-    RationalQuadratic(alpha=1.0),
-    RationalQuadratic(alpha=10.0),
-    WhiteKernel(noise_level=0.1),
-    WhiteKernel(noise_level=1.0),
-    ConstantKernel(),
-]
-
-
-def random_binary_map(a: Kernel, b: Kernel):
-    """
-    Applies a random binary operator (+ or *) with equal probability
-    on kernels ``a`` and ``b``.
-
-    Parameters
-    ----------
-    a
-        A GP kernel.
-    b
-        A GP kernel.
-
-    Returns
-    -------
-        The composite kernel `a + b` or `a * b`.
-    """
-    binary_maps = [lambda x, y: x + y, lambda x, y: x * y]
-    return np.random.choice(binary_maps)(a, b)
-
-
-def sample_from_gp_prior(
-    kernel: Kernel, X: np.ndarray, random_seed: Optional[int] = None
-):
-    """
-    Draw a sample from a GP prior.
-
-    Parameters
-    ----------
-    kernel
-        The GP covaraince kernel.
-    X
-        The input "time" points.
-    random_seed, optional
-        The random seed for sampling, by default None.
-
-    Returns
-    -------
-        A time series sampled from the GP prior.
-    """
-    if X.ndim == 1:
-        X = X[:, None]
-
-    assert X.ndim == 2
-    gpr = GaussianProcessRegressor(kernel=kernel)
-    ts = gpr.sample_y(X, n_samples=1, random_state=random_seed)
-
-    return ts
-
-
-def sample_from_gp_prior_efficient(
-    kernel: Kernel,
-    X: np.ndarray,
-    random_seed: Optional[int] = None,
-    method: str = "eigh",
-):
-    """
-    Draw a sample from a GP prior. An efficient version that allows specification
-    of the sampling method. The default sampling method used in GaussianProcessRegressor
-    is based on SVD which is significantly slower that alternatives such as `eigh` and
-    `cholesky`.
-
-    Parameters
-    ----------
-    kernel
-        The GP covaraince kernel.
-    X
-        The input "time" points.
-    random_seed, optional
-        The random seed for sampling, by default None.
-    method, optional
-        The sampling method for multivariate_normal, by default `eigh`.
-
-    Returns
-    -------
-        A time series sampled from the GP prior.
-    """
-    if X.ndim == 1:
-        X = X[:, None]
-
-    assert X.ndim == 2
-
-    cov = kernel(X)
-    ts = np.random.default_rng(seed=random_seed).multivariate_normal(
-        mean=np.zeros(X.shape[0]), cov=cov, method=method
-    )
-
-    return ts
-
-
-def generate_time_series(max_kernels: int = 5):
-    """Generate a synthetic time series from KernelSynth.
-
-    Parameters
-    ----------
-    max_kernels, optional
-        The maximum number of base kernels to use for each time series, by default 5
-
-    Returns
-    -------
-        A time series generated by KernelSynth.
-    """
-    while True:
-        X = np.linspace(0, 1, LENGTH)
-
-        # Randomly select upto max_kernels kernels from the KERNEL_BANK
-        selected_kernels = np.random.choice(
-            KERNEL_BANK, np.random.randint(1, max_kernels + 1), replace=True
-        )
-
-        # Combine the sampled kernels using random binary operators
-        kernel = functools.reduce(random_binary_map, selected_kernels)
-
-        # Sample a time series from the GP prior
-        try:
-            ts = sample_from_gp_prior(kernel=kernel, X=X)
-        except np.linalg.LinAlgError as err:
-            print("Error caught:", err)
-            continue
-
-        # The timestamp is arbitrary
-        return {"start": np.datetime64("2000-01-01 00:00", "s"), "target": ts.squeeze()}
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("-N", "--num-series", type=int, default=1000_000)
-    parser.add_argument("-J", "--max-kernels", type=int, default=5)
-    args = parser.parse_args()
-    path = Path(__file__).parent / "kernelsynth-data.arrow"
-
-    generated_dataset = Parallel(n_jobs=-1)(
-        delayed(generate_time_series)(max_kernels=args.max_kernels)
-        for _ in tqdm(range(args.num_series))
-    )
-
-    ArrowWriter(compression="lz4").write_to_file(
-        generated_dataset,
-        path=path,
-    )
--- a/scripts/training/configs/chronos-gpt2.yaml
+++ b/scripts/training/configs/chronos-gpt2.yaml
@ -1,35 +0,0 @@
-training_data_paths:
- "/home/ubuntu/tsmixup-data.arrow"
- "/home/ubuntu/kernelsynth-data.arrow"
-probability:
- 0.9
- 0.1
-context_length: 512
-prediction_length: 64
-min_past: 60
-max_steps: 200_000
-save_steps: 100_000
-log_steps: 500
-per_device_train_batch_size: 32
-learning_rate: 0.001
-optim: adamw_torch_fused
-num_samples: 20
-shuffle_buffer_length: 100_000
-gradient_accumulation_steps: 1
-model_id: openai-community/gpt2
-model_type: causal
-random_init: false
-tie_embeddings: false
-output_dir: ./output/
-tf32: true
-torch_compile: true
-tokenizer_class: "MeanScaleUniformBins"
-tokenizer_kwargs:
-  low_limit: -15.0
-  high_limit: 15.0
-n_tokens: 4096
-lr_scheduler_type: linear
-warmup_ratio: 0.0
-dataloader_num_workers: 1
-max_missing_prop: 0.1
-use_eos_token: true
--- a/scripts/training/configs/chronos-t5-base.yaml
+++ b/scripts/training/configs/chronos-t5-base.yaml
@ -1,35 +0,0 @@
-training_data_paths:
- "/home/ubuntu/tsmixup-data.arrow"
- "/home/ubuntu/kernelsynth-data.arrow"
-probability:
- 0.9
- 0.1
-context_length: 512
-prediction_length: 64
-min_past: 60
-max_steps: 200_000
-save_steps: 100_000
-log_steps: 500
-per_device_train_batch_size: 32
-learning_rate: 0.001
-optim: adamw_torch_fused
-num_samples: 20
-shuffle_buffer_length: 100_000
-gradient_accumulation_steps: 1
-model_id: google/t5-efficient-base
-model_type: seq2seq
-random_init: true
-tie_embeddings: true
-output_dir: ./output/
-tf32: true
-torch_compile: true
-tokenizer_class: "MeanScaleUniformBins"
-tokenizer_kwargs:
-  low_limit: -15.0
-  high_limit: 15.0
-n_tokens: 4096
-lr_scheduler_type: linear
-warmup_ratio: 0.0
-dataloader_num_workers: 1
-max_missing_prop: 0.9
-use_eos_token: true
--- a/scripts/training/configs/chronos-t5-large.yaml
+++ b/scripts/training/configs/chronos-t5-large.yaml
@ -1,35 +0,0 @@
-training_data_paths:
- "/home/ubuntu/tsmixup-data.arrow"
- "/home/ubuntu/kernelsynth-data.arrow"
-probability:
- 0.9
- 0.1
-context_length: 512
-prediction_length: 64
-min_past: 60
-max_steps: 200_000
-save_steps: 100_000
-log_steps: 500
-per_device_train_batch_size: 8
-learning_rate: 0.001
-optim: adamw_torch_fused
-num_samples: 20
-shuffle_buffer_length: 100_000
-gradient_accumulation_steps: 4
-model_id: google/t5-efficient-large
-model_type: seq2seq
-random_init: true
-tie_embeddings: true
-output_dir: ./output/
-tf32: true
-torch_compile: true
-tokenizer_class: "MeanScaleUniformBins"
-tokenizer_kwargs:
-  low_limit: -15.0
-  high_limit: 15.0
-n_tokens: 4096
-lr_scheduler_type: linear
-warmup_ratio: 0.0
-dataloader_num_workers: 1
-max_missing_prop: 0.9
-use_eos_token: true
--- a/scripts/training/configs/chronos-t5-mini.yaml
+++ b/scripts/training/configs/chronos-t5-mini.yaml
@ -1,35 +0,0 @@
-training_data_paths:
- "/home/ubuntu/tsmixup-data.arrow"
- "/home/ubuntu/kernelsynth-data.arrow"
-probability:
- 0.9
- 0.1
-context_length: 512
-prediction_length: 64
-min_past: 60
-max_steps: 200_000
-save_steps: 100_000
-log_steps: 500
-per_device_train_batch_size: 32
-learning_rate: 0.001
-optim: adamw_torch_fused
-num_samples: 20
-shuffle_buffer_length: 100_000
-gradient_accumulation_steps: 1
-model_id: google/t5-efficient-mini
-model_type: seq2seq
-random_init: true
-tie_embeddings: true
-output_dir: ./output/
-tf32: true
-torch_compile: true
-tokenizer_class: "MeanScaleUniformBins"
-tokenizer_kwargs:
-  low_limit: -15.0
-  high_limit: 15.0
-n_tokens: 4096
-lr_scheduler_type: linear
-warmup_ratio: 0.0
-dataloader_num_workers: 1
-max_missing_prop: 0.9
-use_eos_token: true
--- a/scripts/training/configs/chronos-t5-small.yaml
+++ b/scripts/training/configs/chronos-t5-small.yaml
@ -1,35 +0,0 @@
-training_data_paths:
- "/home/ubuntu/tsmixup-data.arrow"
- "/home/ubuntu/kernelsynth-data.arrow"
-probability:
- 0.9
- 0.1
-context_length: 512
-prediction_length: 64
-min_past: 60
-max_steps: 200_000
-save_steps: 100_000
-log_steps: 500
-per_device_train_batch_size: 32
-learning_rate: 0.001
-optim: adamw_torch_fused
-num_samples: 20
-shuffle_buffer_length: 100_000
-gradient_accumulation_steps: 1
-model_id: google/t5-efficient-small
-model_type: seq2seq
-random_init: true
-tie_embeddings: true
-output_dir: ./output/
-tf32: true
-torch_compile: true
-tokenizer_class: "MeanScaleUniformBins"
-tokenizer_kwargs:
-  low_limit: -15.0
-  high_limit: 15.0
-n_tokens: 4096
-lr_scheduler_type: linear
-warmup_ratio: 0.0
-dataloader_num_workers: 1
-max_missing_prop: 0.9
-use_eos_token: true
--- a/scripts/training/configs/chronos-t5-tiny.yaml
+++ b/scripts/training/configs/chronos-t5-tiny.yaml
@ -1,35 +0,0 @@
-training_data_paths:
- "/home/ubuntu/tsmixup-data.arrow"
- "/home/ubuntu/kernelsynth-data.arrow"
-probability:
- 0.9
- 0.1
-context_length: 512
-prediction_length: 64
-min_past: 60
-max_steps: 200_000
-save_steps: 100_000
-log_steps: 500
-per_device_train_batch_size: 32
-learning_rate: 0.001
-optim: adamw_torch_fused
-num_samples: 20
-shuffle_buffer_length: 100_000
-gradient_accumulation_steps: 1
-model_id: google/t5-efficient-tiny
-model_type: seq2seq
-random_init: true
-tie_embeddings: true
-output_dir: ./output/
-tf32: true
-torch_compile: true
-tokenizer_class: "MeanScaleUniformBins"
-tokenizer_kwargs:
-  low_limit: -15.0
-  high_limit: 15.0
-n_tokens: 4096
-lr_scheduler_type: linear
-warmup_ratio: 0.0
-dataloader_num_workers: 1
-max_missing_prop: 0.9
-use_eos_token: true
--- a/scripts/training/train.py
+++ b/scripts/training/train.py
@ -1,702 +0,0 @@
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-import ast
-import logging
-import os
-import re
-import sys
-import json
-import itertools
-import random
-from copy import deepcopy
-from pathlib import Path
-from functools import partial
-from typing import List, Iterator, Optional, Dict
-
-import typer
-from typer_config import use_yaml_config
-import numpy as np
-import torch
-import torch.distributed as dist
-from torch.utils.data import IterableDataset, get_worker_info
-import transformers
-from transformers import (
-    AutoModelForSeq2SeqLM,
-    AutoModelForCausalLM,
-    AutoConfig,
-    T5Config,
-    Trainer,
-    TrainingArguments,
-)
-import accelerate
-import gluonts
-from gluonts.dataset.common import FileDataset
-from gluonts.itertools import Cyclic, Map, Filter
-from gluonts.transform import (
-    FilterTransformation,
-    TestSplitSampler,
-    ValidationSplitSampler,
-    InstanceSplitter,
-    ExpectedNumInstanceSampler,
-    MissingValueImputation,
-    LeavesMissingValues,
-    LastValueImputation,
-)
-
-from chronos import ChronosConfig, ChronosTokenizer
-
-
-app = typer.Typer(pretty_exceptions_enable=False)
-
-
-def is_main_process() -> bool:
-    """
-    Check if we're on the main process.
-    """
-    if not dist.is_torchelastic_launched():
-        return True
-    return int(os.environ["RANK"]) == 0
-
-
-def log_on_main(msg: str, logger: logging.Logger, log_level: int = logging.INFO):
-    """
-    Log the given message using the given logger, if we're on the main process.
-    """
-    if is_main_process():
-        logger.log(log_level, msg)
-
-
-def get_training_job_info() -> Dict:
-    """
-    Returns info about this training job.
-    """
-    job_info = {}
-
-    # CUDA info
-    job_info["cuda_available"] = torch.cuda.is_available()
-    if torch.cuda.is_available():
-        job_info["device_count"] = torch.cuda.device_count()
-
-        job_info["device_names"] = {
-            idx: torch.cuda.get_device_name(idx)
-            for idx in range(torch.cuda.device_count())
-        }
-        job_info["mem_info"] = {
-            idx: torch.cuda.mem_get_info(device=idx)
-            for idx in range(torch.cuda.device_count())
-        }
-
-    # DDP info
-    job_info["torchelastic_launched"] = dist.is_torchelastic_launched()
-
-    if dist.is_torchelastic_launched():
-        job_info["world_size"] = dist.get_world_size()
-
-    # Versions
-    job_info["python_version"] = sys.version.replace("\n", " ")
-    job_info["torch_version"] = torch.__version__
-    job_info["numpy_version"] = np.__version__
-    job_info["gluonts_version"] = gluonts.__version__
-    job_info["transformers_version"] = transformers.__version__
-    job_info["accelerate_version"] = accelerate.__version__
-
-    return job_info
-
-
-def save_training_info(ckpt_path: Path, training_config: Dict):
-    """
-    Save info about this training job in a json file for documentation.
-    """
-    assert ckpt_path.is_dir()
-    with open(ckpt_path / "training_info.json", "w") as fp:
-        json.dump(
-            {"training_config": training_config, "job_info": get_training_job_info()},
-            fp,
-            indent=4,
-        )
-
-
-def get_next_path(
-    base_fname: str,
-    base_dir: Path,
-    file_type: str = "yaml",
-    separator: str = "-",
-):
-    """
-    Gets the next available path in a directory. For example, if `base_fname="results"`
-    and `base_dir` has files ["results-0.yaml", "results-1.yaml"], this function returns
-    "results-2.yaml".
-    """
-    if file_type == "":
-        # Directory
-        items = filter(
-            lambda x: x.is_dir() and re.match(f"^{base_fname}{separator}\\d+$", x.stem),
-            base_dir.glob("*"),
-        )
-    else:
-        # File
-        items = filter(
-            lambda x: re.match(f"^{base_fname}{separator}\\d+$", x.stem),
-            base_dir.glob(f"*.{file_type}"),
-        )
-    run_nums = list(
-        map(lambda x: int(x.stem.replace(base_fname + separator, "")), items)
-    ) + [-1]
-
-    next_num = max(run_nums) + 1
-    fname = f"{base_fname}{separator}{next_num}" + (
-        f".{file_type}" if file_type != "" else ""
-    )
-
-    return base_dir / fname
-
-
-def load_model(
-    model_id="google/t5-efficient-tiny",
-    model_type="seq2seq",
-    vocab_size=4096,
-    random_init=False,
-    tie_embeddings=False,
-    pad_token_id=0,
-    eos_token_id=1,
-):
-    """
-    Load the specified HuggingFace model, adjusting the vocabulary
-    size, special token IDs, and initialization options.
-
-    This allows to set a model up for training on a new vocabulary
-    of tokens.
-    """
-    assert model_type in ["seq2seq", "causal"]
-    AutoModelClass = (
-        AutoModelForSeq2SeqLM if model_type == "seq2seq" else AutoModelForCausalLM
-    )
-    if random_init:
-        log_on_main("Using random initialization", logger)
-        config = AutoConfig.from_pretrained(model_id)
-        if isinstance(config, T5Config):
-            # The default initializer_factor (1.0) in transformers is too large
-            config.initializer_factor = 0.05
-        config.tie_word_embeddings = tie_embeddings
-        model = AutoModelClass.from_config(config)
-    else:
-        log_on_main(f"Using pretrained initialization from {model_id}", logger)
-        model = AutoModelClass.from_pretrained(model_id)
-
-    model.resize_token_embeddings(vocab_size)
-
-    model.config.pad_token_id = model.generation_config.pad_token_id = pad_token_id
-    model.config.eos_token_id = model.generation_config.eos_token_id = eos_token_id
-
-    return model
-
-
-def has_enough_observations(
-    entry: dict, min_length: int = 0, max_missing_prop: float = 1.0
-) -> bool:
-    """
-    Check if the given entry has enough observations in the ``"target"`` attribute.
-
-    Parameters
-    ----------
-    entry
-        The data entry (dictionary) to be tested.
-    min_length
-        The minimum length the ``"target"`` attribute must have.
-    max_missing_prop
-        The maximum proportion of missing data allowed in the ``"target"``
-        attribute.
-    """
-    if (
-        len(entry["target"]) >= min_length
-        and np.isnan(entry["target"]).mean() <= max_missing_prop
-    ):
-        return True
-    return False
-
-
-class PseudoShuffledIterableDataset(IterableDataset):
-    """
-    Shuffle entries from an iterable by temporarily accumulating them
-    in an intermediate buffer.
-
-    Parameters
-    ----------
-    base_dataset
-        The original iterable object, representing the dataset.
-    shuffle_buffer_length
-        Size of the buffer use to shuffle entries from the base dataset.
-    """
-
-    def __init__(self, base_dataset, shuffle_buffer_length: int = 100) -> None:
-        super().__init__()
-        self.base_dataset = base_dataset
-        self.shuffle_buffer_length = shuffle_buffer_length
-        self.generator = torch.Generator()
-
-    def __iter__(self):
-        shuffle_buffer = []
-
-        for element in self.base_dataset:
-            shuffle_buffer.append(element)
-            if len(shuffle_buffer) >= self.shuffle_buffer_length:
-                idx = torch.randint(
-                    len(shuffle_buffer), size=(), generator=self.generator
-                )
-                yield shuffle_buffer.pop(idx)
-
-        while shuffle_buffer:
-            idx = torch.randint(len(shuffle_buffer), size=(), generator=self.generator)
-            yield shuffle_buffer.pop(idx)
-
-
-class ShuffleMixin:
-    """
-    Mix-in class that datasets can inherit from to get
-    shuffling functionality.
-    """
-
-    def shuffle(self, shuffle_buffer_length: int = 100):
-        return PseudoShuffledIterableDataset(self, shuffle_buffer_length)
-
-
-class ChronosDataset(IterableDataset, ShuffleMixin):
-    """
-    Dataset wrapper, using a ``ChronosTokenizer`` to turn data from a time series
-    into a HuggingFace-compatible set of ``input_ids``, ``attention_mask`` and
-    ``labels``.
-
-    Entries from the original datasets are assumed to have a ``"start"`` attribute
-    (of type ``pd.Period``), and a ``"target"`` attribute (of type ``np.ndarray``).
-
-    Parameters
-    ----------
-    datasets
-        Datasets containing the original time series data.
-    probabilities
-        In training mode, data will be sampled from each of the original datasets
-        with these probabilities.
-    tokenizer
-        Tokenizer to be used to turn sequences of real numbers into token IDs.
-    context_length
-        Samples context will be limited to this length.
-    prediction_length
-        Samples labels will be limited to this length.
-    drop_prob
-        In training mode, observations from a sample will be turned into ``np.nan``,
-        i.e. turned into missing values, with this probability.
-    min_past
-        Data samples will be considered only if there's at least ``min_past``-many
-        historical observations.
-    mode
-        One of ``"training"``, ``"validation"``, or ``"test"``.
-    np_dtype
-        Numpy float data type.
-    """
-
-    def __init__(
-        self,
-        datasets: list,
-        probabilities: List[float],
-        tokenizer: ChronosTokenizer,
-        context_length: int = 512,
-        prediction_length: int = 64,
-        drop_prob: float = 0.2,
-        min_past: Optional[int] = None,
-        model_type: str = "seq2seq",
-        imputation_method: Optional[MissingValueImputation] = None,
-        mode: str = "training",
-        np_dtype=np.float32,
-    ) -> None:
-        super().__init__()
-
-        assert len(probabilities) == len(datasets)
-        assert mode in ("training", "validation", "test")
-        assert model_type in ("seq2seq", "causal")
-
-        self.datasets = datasets
-        self.probabilities = probabilities
-        self.tokenizer = tokenizer
-        self.context_length = context_length
-        self.prediction_length = prediction_length
-        self.drop_prob = drop_prob if model_type == "seq2seq" else 0.0
-        self.min_past = min_past or prediction_length
-        self.model_type = model_type
-        self.imputation_method = imputation_method or LeavesMissingValues()
-        self.mode = mode
-        self.np_dtype = np_dtype
-
-    def preprocess_entry(self, entry: dict, mode: str) -> dict:
-        entry = {f: entry[f] for f in ["start", "target"]}
-        entry["target"] = np.asarray(entry["target"], dtype=self.np_dtype)
-        assert entry["target"].ndim == 1, f"got {entry['target'].ndim=}, expected 1"
-
-        if self.model_type == "causal":
-            # Causal models do not play nice with missing values, so it is
-            # recommended to use an imputation method, e.g., LastValueImputation
-            entry["target"] = self.imputation_method(entry["target"])
-
-        if mode == "training" and self.drop_prob > 0:
-            target = entry["target"].copy()
-            drop_p = np.random.uniform(low=0.0, high=self.drop_prob)
-            mask = np.random.choice(
-                [True, False], size=len(target), p=[drop_p, 1 - drop_p]
-            )
-            target[mask] = np.nan
-            entry["target"] = target
-
-        return entry
-
-    def _create_instance_splitter(self, mode: str):
-        assert mode in ["training", "test", "validation"]
-
-        instance_sampler = {
-            "training": ExpectedNumInstanceSampler(
-                num_instances=1.0,
-                min_instances=1,
-                min_past=self.min_past,
-                min_future=self.prediction_length,
-            ),
-            "test": TestSplitSampler(),
-            "validation": ValidationSplitSampler(min_future=self.prediction_length),
-        }[mode]
-
-        return InstanceSplitter(
-            target_field="target",
-            is_pad_field="is_pad",
-            start_field="start",
-            forecast_start_field="forecast_start",
-            instance_sampler=instance_sampler,
-            past_length=self.context_length,
-            future_length=self.prediction_length,
-            dummy_value=np.nan,
-        )
-
-    def create_training_data(self, data):
-        data = Cyclic(data)
-        split_transform = self._create_instance_splitter(
-            "training"
-        ) + FilterTransformation(
-            condition=lambda entry: (~np.isnan(entry["past_target"])).sum() > 0
-        )
-        data = split_transform.apply(data, is_train=True)
-        return data
-
-    def create_test_data(self, data):
-        data = self._create_instance_splitter("test").apply(data, is_train=False)
-        return data
-
-    def create_validation_data(self, data):
-        data = self._create_instance_splitter("validation").apply(data, is_train=False)
-        return data
-
-    def to_hf_format(self, entry: dict) -> dict:
-        past_target = torch.tensor(entry["past_target"]).unsqueeze(0)
-        input_ids, attention_mask, scale = self.tokenizer.context_input_transform(
-            past_target
-        )
-        future_target = torch.tensor(entry["future_target"]).unsqueeze(0)
-        labels, labels_mask = self.tokenizer.label_input_transform(future_target, scale)
-        labels[labels_mask == 0] = -100
-
-        if self.model_type == "causal":
-            # The InstanceSplitter pads time series on the left to be equal to the
-            # context_length. However, certain models (e.g., GPT2) with absolute
-            # position embeddings should not be trained with left padding.
-            # The following piece of code moves padding from left to right.
-
-            assert input_ids.shape[-1] == entry["past_is_pad"].shape[0]
-
-            # Find the index where padding starts
-            pad_start_idx = np.searchsorted(1 - entry["past_is_pad"], 1)
-            padded_input_ids, obs_input_ids = torch.tensor_split(
-                input_ids, [pad_start_idx], dim=-1
-            )
-            padded_attention_mask, obs_attention_mask = torch.tensor_split(
-                attention_mask, [pad_start_idx], dim=-1
-            )
-
-            # Move padding to the right
-            input_ids = torch.cat(
-                [
-                    obs_input_ids,
-                    labels,
-                    padded_input_ids,
-                ],
-                axis=-1,
-            )
-            attention_mask = torch.cat(
-                [
-                    obs_attention_mask,
-                    labels_mask,
-                    padded_attention_mask,
-                ],
-                axis=-1,
-            )
-
-            # labels for causal models are same as the input_ids.
-            # Internally transformers shifts the labels by one during training.
-            labels = input_ids.clone()
-            input_ids[~attention_mask] = self.tokenizer.config.pad_token_id
-            labels[~attention_mask] = -100
-
-        return {
-            "input_ids": input_ids.squeeze(0),
-            "attention_mask": attention_mask.squeeze(0),
-            "labels": labels.squeeze(0),
-        }
-
-    def __iter__(self) -> Iterator:
-        preprocessed_datasets = [
-            Map(
-                partial(self.preprocess_entry, mode=self.mode),
-                dataset,
-            )
-            for dataset in self.datasets
-        ]
-
-        if self.mode == "training":
-            iterables = [
-                self.create_training_data(dataset) for dataset in preprocessed_datasets
-            ]
-        elif self.mode == "test":
-            iterables = [
-                self.create_test_data(dataset) for dataset in preprocessed_datasets
-            ]
-        else:
-            iterables = [
-                self.create_validation_data(dataset)
-                for dataset in preprocessed_datasets
-            ]
-
-        worker_info = get_worker_info()
-        if worker_info is None:
-            probs = list(self.probabilities)
-        else:
-            worker_id = worker_info.id
-            num_workers = worker_info.num_workers
-            iterables = list(itertools.islice(iterables, worker_id, None, num_workers))
-            probs = list(
-                itertools.islice(self.probabilities, worker_id, None, num_workers)
-            )
-
-        probs = [prob / sum(probs) for prob in probs]
-
-        iterators = list(map(iter, iterables))
-        if self.mode == "training":
-            while True:
-                idx = np.random.choice(range(len(iterators)), p=probs)
-                try:
-                    yield self.to_hf_format(next(iterators[idx]))
-                except StopIteration:
-                    probs[idx] = 0
-                    if sum(probs) == 0:
-                        return
-                    probs = [prob / sum(probs) for prob in probs]
-        else:
-            for entry in itertools.chain(*iterators):
-                yield self.to_hf_format(entry)
-
-
-@app.command()
-@use_yaml_config(param_name="config")
-def main(
-    training_data_paths: str,
-    probability: Optional[str] = None,
-    context_length: int = 512,
-    prediction_length: int = 64,
-    min_past: int = 64,
-    max_steps: int = 200_000,
-    save_steps: int = 50_000,
-    log_steps: int = 500,
-    per_device_train_batch_size: int = 32,
-    learning_rate: float = 1e-3,
-    optim: str = "adamw_torch_fused",
-    shuffle_buffer_length: int = 100,
-    gradient_accumulation_steps: int = 2,
-    model_id: str = "google/t5-efficient-tiny",
-    model_type: str = "seq2seq",
-    random_init: bool = False,
-    tie_embeddings: bool = False,
-    output_dir: str = "./output/",
-    tf32: bool = True,
-    torch_compile: bool = True,
-    tokenizer_class: str = "MeanScaleUniformBins",
-    tokenizer_kwargs: str = "{'low_limit': -15.0, 'high_limit': 15.0}",
-    n_tokens: int = 4096,
-    n_special_tokens: int = 2,
-    pad_token_id: int = 0,
-    eos_token_id: int = 1,
-    use_eos_token: bool = True,
-    lr_scheduler_type: str = "linear",
-    warmup_ratio: float = 0.0,
-    dataloader_num_workers: int = 1,
-    max_missing_prop: float = 0.9,
-    num_samples: int = 20,
-    temperature: float = 1.0,
-    top_k: int = 50,
-    top_p: float = 1.0,
-    seed: Optional[int] = None,
-):
-    if tf32 and not (
-        torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8
-    ):
-        # TF32 floating point format is available only on NVIDIA GPUs
-        # with compute capability 8 and above. See link for details.
-        # https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#compute-capability-8-x
-        log_on_main(
-            "TF32 format is only available on devices with compute capability >= 8. "
-            "Setting tf32 to False.",
-            logger,
-        )
-        tf32 = False
-
-    if seed is None:
-        seed = random.randint(0, 2**32)
-
-    log_on_main(f"Using SEED: {seed}", logger)
-    transformers.set_seed(seed=seed)
-
-    raw_training_config = deepcopy(locals())
-    output_dir = Path(output_dir)
-    training_data_paths = ast.literal_eval(training_data_paths)
-    assert isinstance(training_data_paths, list)
-
-    if isinstance(probability, str):
-        probability = ast.literal_eval(probability)
-    elif probability is None:
-        probability = [1.0 / len(training_data_paths)] * len(training_data_paths)
-    assert isinstance(probability, list)
-
-    assert len(training_data_paths) == len(probability)
-
-    if dataloader_num_workers > len(training_data_paths):
-        log_on_main(
-            f"Setting the number of data loader workers to {len(training_data_paths)}, "
-            f"instead of {dataloader_num_workers}.",
-            logger,
-        )
-        dataloader_num_workers = len(training_data_paths)
-
-    if isinstance(tokenizer_kwargs, str):
-        tokenizer_kwargs = ast.literal_eval(tokenizer_kwargs)
-    assert isinstance(tokenizer_kwargs, dict)
-
-    assert model_type in ["seq2seq", "causal"]
-
-    output_dir = get_next_path("run", base_dir=output_dir, file_type="")
-
-    log_on_main(f"Logging dir: {output_dir}", logger)
-    log_on_main(
-        f"Loading and filtering {len(training_data_paths)} datasets "
-        f"for training: {training_data_paths}",
-        logger,
-    )
-
-    log_on_main(
-        f"Mixing probabilities: {probability}",
-        logger,
-    )
-
-    train_datasets = [
-        Filter(
-            partial(
-                has_enough_observations,
-                min_length=min_past + prediction_length,
-                max_missing_prop=max_missing_prop,
-            ),
-            FileDataset(path=Path(data_path), freq="h"),
-        )
-        for data_path in training_data_paths
-    ]
-
-    log_on_main("Initializing model", logger)
-
-    model = load_model(
-        model_id=model_id,
-        model_type=model_type,
-        vocab_size=n_tokens,
-        random_init=random_init,
-        tie_embeddings=tie_embeddings,
-        pad_token_id=pad_token_id,
-        eos_token_id=eos_token_id,
-    )
-
-    chronos_config = ChronosConfig(
-        tokenizer_class=tokenizer_class,
-        tokenizer_kwargs=tokenizer_kwargs,
-        n_tokens=n_tokens,
-        n_special_tokens=n_special_tokens,
-        pad_token_id=pad_token_id,
-        eos_token_id=eos_token_id,
-        use_eos_token=use_eos_token,
-        model_type=model_type,
-        context_length=context_length,
-        prediction_length=prediction_length,
-        num_samples=num_samples,
-        temperature=temperature,
-        top_k=top_k,
-        top_p=top_p,
-    )
-
-    # Add extra items to model config so that it's saved in the ckpt
-    model.config.chronos_config = chronos_config.__dict__
-
-    shuffled_train_dataset = ChronosDataset(
-        datasets=train_datasets,
-        probabilities=probability,
-        tokenizer=chronos_config.create_tokenizer(),
-        context_length=context_length,
-        prediction_length=prediction_length,
-        min_past=min_past,
-        model_type=model_type,
-        imputation_method=LastValueImputation() if model_type == "causal" else None,
-        mode="training",
-    ).shuffle(shuffle_buffer_length=shuffle_buffer_length)
-
-    # Define training args
-    training_args = TrainingArguments(
-        output_dir=str(output_dir),
-        per_device_train_batch_size=per_device_train_batch_size,
-        learning_rate=learning_rate,
-        lr_scheduler_type=lr_scheduler_type,
-        warmup_ratio=warmup_ratio,
-        optim=optim,
-        logging_dir=str(output_dir / "logs"),
-        logging_strategy="steps",
-        logging_steps=log_steps,
-        save_strategy="steps",
-        save_steps=save_steps,
-        report_to=["tensorboard"],
-        max_steps=max_steps,
-        gradient_accumulation_steps=gradient_accumulation_steps,
-        dataloader_num_workers=dataloader_num_workers,
-        tf32=tf32,  # remove this if not using Ampere GPUs (e.g., A100)
-        torch_compile=torch_compile,
-        ddp_find_unused_parameters=False,
-        remove_unused_columns=False,
-    )
-
-    # Create Trainer instance
-    trainer = Trainer(
-        model=model,
-        args=training_args,
-        train_dataset=shuffled_train_dataset,
-    )
-    log_on_main("Training", logger)
-
-    trainer.train()
-
-    if is_main_process():
-        model.save_pretrained(output_dir / "checkpoint-final")
-        save_training_info(
-            output_dir / "checkpoint-final", training_config=raw_training_config
-        )
-
-
-if __name__ == "__main__":
-    logging.basicConfig(format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
-    logger = logging.getLogger(__file__)
-    logger.setLevel(logging.INFO)
-    app()
--- a/src/chronos/init.py
+++ b/src/chronos/init.py
@ -1,24 +0,0 @@
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from .base import BaseChronosPipeline, ForecastType
-from .chronos import (
-    ChronosConfig,
-    ChronosModel,
-    ChronosPipeline,
-    ChronosTokenizer,
-    MeanScaleUniformBins,
-)
-from .chronos_bolt import ChronosBoltConfig, ChronosBoltPipeline
-
-__all__ = [
-    "BaseChronosPipeline",
-    "ForecastType",
-    "ChronosConfig",
-    "ChronosModel",
-    "ChronosPipeline",
-    "ChronosTokenizer",
-    "MeanScaleUniformBins",
-    "ChronosBoltConfig",
-    "ChronosBoltPipeline",
-]
--- a/src/chronos/base.py
+++ b/src/chronos/base.py
@ -1,164 +0,0 @@
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-# Authors: Caner Turkmen <atturkm@amazon.com>, Abdul Fatir Ansari <ansarnd@amazon.com>, Lorenzo Stella <stellalo@amazon.com>
-# Original source:
-# https://github.com/autogluon/autogluon/blob/f57beb26cb769c6e0d484a6af2b89eab8aee73a8/timeseries/src/autogluon/timeseries/models/chronos/pipeline/base.py
-
-from enum import Enum
-from pathlib import Path
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
-
-import torch
-
-if TYPE_CHECKING:
-    from transformers import PreTrainedModel
-
-from .utils import left_pad_and_stack_1D
-
-
-class ForecastType(Enum):
-    SAMPLES = "samples"
-    QUANTILES = "quantiles"
-
-
-class PipelineRegistry(type):
-    REGISTRY: Dict[str, "PipelineRegistry"] = {}
-
-    def __new__(cls, name, bases, attrs):
-        """See, https://github.com/faif/python-patterns."""
-        new_cls = type.__new__(cls, name, bases, attrs)
-        if name is not None:
-            cls.REGISTRY[name] = new_cls
-
-        return new_cls
-
-
-class BaseChronosPipeline(metaclass=PipelineRegistry):
-    forecast_type: ForecastType
-    dtypes = {"bfloat16": torch.bfloat16, "float32": torch.float32}
-
-    def __init__(self, inner_model: "PreTrainedModel"):
-        """
-        Parameters
-        ----------
-        inner_model : PreTrainedModel
-            A hugging-face transformers PreTrainedModel, e.g., T5ForConditionalGeneration
-        """
-        # for easy access to the inner HF-style model
-        self.inner_model = inner_model
-
-    def _prepare_and_validate_context(
-        self, context: Union[torch.Tensor, List[torch.Tensor]]
-    ):
-        if isinstance(context, list):
-            context = left_pad_and_stack_1D(context)
-        assert isinstance(context, torch.Tensor)
-        if context.ndim == 1:
-            context = context.unsqueeze(0)
-        assert context.ndim == 2
-
-        return context
-
-    def predict(
-        self,
-        context: Union[torch.Tensor, List[torch.Tensor]],
-        prediction_length: Optional[int] = None,
-        **kwargs,
-    ):
-        """
-        Get forecasts for the given time series. Predictions will be
-        returned in fp32 on the cpu.
-
-        Parameters
-        ----------
-        context
-            Input series. This is either a 1D tensor, or a list
-            of 1D tensors, or a 2D tensor whose first dimension
-            is batch. In the latter case, use left-padding with
-            ``torch.nan`` to align series of different lengths.
-        prediction_length
-            Time steps to predict. Defaults to a model-dependent
-            value if not given.
-
-        Returns
-        -------
-        forecasts
-            Tensor containing forecasts. The layout and meaning
-            of the forecasts values depends on ``self.forecast_type``.
-        """
-        raise NotImplementedError()
-
-    def predict_quantiles(
-        self,
-        context: Union[torch.Tensor, List[torch.Tensor]],
-        prediction_length: Optional[int] = None,
-        quantile_levels: List[float] = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9],
-        **kwargs,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Get quantile and mean forecasts for given time series.
-        Predictions will be returned in fp32 on the cpu.
-
-        Parameters
-        ----------
-        context : Union[torch.Tensor, List[torch.Tensor]]
-            Input series. This is either a 1D tensor, or a list
-            of 1D tensors, or a 2D tensor whose first dimension
-            is batch. In the latter case, use left-padding with
-            ``torch.nan`` to align series of different lengths.
-        prediction_length : Optional[int], optional
-            Time steps to predict. Defaults to a model-dependent
-            value if not given.
-        quantile_levels : List[float], optional
-            Quantile levels to compute, by default [0.1, 0.2, ..., 0.9]
-
-        Returns
-        -------
-        quantiles
-            Tensor containing quantile forecasts. Shape
-            (batch_size, prediction_length, num_quantiles)
-        mean
-            Tensor containing mean (point) forecasts. Shape
-            (batch_size, prediction_length)
-        """
-        raise NotImplementedError()
-
-    @classmethod
-    def from_pretrained(
-        cls,
-        pretrained_model_name_or_path: Union[str, Path],
-        *model_args,
-        **kwargs,
-    ):
-        """
-        Load the model, either from a local path or from the HuggingFace Hub.
-        Supports the same arguments as ``AutoConfig`` and ``AutoModel``
-        from ``transformers``.
-        """
-        from transformers import AutoConfig
-
-        torch_dtype = kwargs.get("torch_dtype", "auto")
-        if torch_dtype != "auto" and isinstance(torch_dtype, str):
-            kwargs["torch_dtype"] = cls.dtypes[torch_dtype]
-
-        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
-        is_valid_config = hasattr(config, "chronos_pipeline_class") or hasattr(
-            config, "chronos_config"
-        )
-
-        if not is_valid_config:
-            raise ValueError("Not a Chronos config file")
-
-        pipeline_class_name = getattr(
-            config, "chronos_pipeline_class", "ChronosPipeline"
-        )
-        class_ = PipelineRegistry.REGISTRY.get(pipeline_class_name)
-        if class_ is None:
-            raise ValueError(
-                f"Trying to load unknown pipeline class: {pipeline_class_name}"
-            )
-
-        return class_.from_pretrained(  # type: ignore[attr-defined]
-            pretrained_model_name_or_path, *model_args, **kwargs
-        )
--- a/src/chronos/chronos.py
+++ b/src/chronos/chronos.py
@ -1,582 +0,0 @@
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-# Authors: Abdul Fatir Ansari <ansarnd@amazon.com>, Lorenzo Stella <stellalo@amazon.com>, Caner Turkmen <atturkm@amazon.com>
-
-import logging
-from dataclasses import dataclass
-from typing import Any, Dict, List, Literal, Optional, Tuple, Union
-
-import torch
-import torch.nn as nn
-from transformers import (
-    AutoConfig,
-    AutoModelForCausalLM,
-    AutoModelForSeq2SeqLM,
-    GenerationConfig,
-    PreTrainedModel,
-)
-
-import chronos
-from chronos.base import BaseChronosPipeline, ForecastType
-from chronos.utils import left_pad_and_stack_1D
-
-logger = logging.getLogger(__file__)
-
-
-@dataclass
-class ChronosConfig:
-    """
-    This class holds all the configuration parameters to be used
-    by ``ChronosTokenizer`` and ``ChronosModel``.
-    """
-
-    tokenizer_class: str
-    tokenizer_kwargs: Dict[str, Any]
-    context_length: int
-    prediction_length: int
-    n_tokens: int
-    n_special_tokens: int
-    pad_token_id: int
-    eos_token_id: int
-    use_eos_token: bool
-    model_type: Literal["causal", "seq2seq"]
-    num_samples: int
-    temperature: float
-    top_k: int
-    top_p: float
-
-    def __post_init__(self):
-        assert (
-            self.pad_token_id < self.n_special_tokens
-            and self.eos_token_id < self.n_special_tokens
-        ), f"Special token id's must be smaller than {self.n_special_tokens=}"
-
-    def create_tokenizer(self) -> "ChronosTokenizer":
-        class_ = getattr(chronos, self.tokenizer_class)
-        return class_(**self.tokenizer_kwargs, config=self)
-
-
-class ChronosTokenizer:
-    """
-    A ``ChronosTokenizer`` definines how time series are mapped into token IDs
-    and back.
-
-    For details, see the ``input_transform`` and ``output_transform`` methods,
-    which concrete classes must implement.
-    """
-
-    def context_input_transform(
-        self,
-        context: torch.Tensor,
-    ) -> Tuple:
-        """
-        Turn a batch of time series into token IDs, attention map, and tokenizer_state.
-
-        Parameters
-        ----------
-        context
-            A tensor shaped (batch_size, time_length), containing the
-            timeseries to forecast. Use left-padding with ``torch.nan``
-            to align time series of different lengths.
-
-        Returns
-        -------
-        token_ids
-            A tensor of integers, shaped (batch_size, time_length + 1)
-            if ``config.use_eos_token`` and (batch_size, time_length)
-            otherwise, containing token IDs for the input series.
-        attention_mask
-            A boolean tensor, same shape as ``token_ids``, indicating
-            which input observations are not ``torch.nan`` (i.e. not
-            missing nor padding).
-        tokenizer_state
-            An object that can be passed to ``label_input_transform``
-            and ``output_transform``. Contains the relevant information
-            to decode output samples into real values,
-            such as location and scale parameters.
-        """
-        raise NotImplementedError()
-
-    def label_input_transform(self, label: torch.Tensor, tokenizer_state: Any) -> Tuple:
-        """
-        Turn a batch of label slices of time series into token IDs and attention map
-        using the ``tokenizer_state`` provided by ``context_input_transform``.
-
-        Parameters
-        ----------
-        context
-            A tensor shaped (batch_size, time_length), containing the
-            timeseries to forecast. Use left-padding with ``torch.nan``
-            to align time series of different lengths.
-        tokenizer_state
-            An object returned by ``context_input_transform`` containing
-            relevant information to preprocess data, such as location and
-            scale. The nature of this depends on the specific tokenizer.
-            This is used for tokenizing the label, in order to use the same
-            scaling used to tokenize the context.
-
-        Returns
-        -------
-        token_ids
-            A tensor of integers, shaped (batch_size, time_length + 1)
-            if ``config.use_eos_token`` and (batch_size, time_length)
-            otherwise, containing token IDs for the input series.
-        attention_mask
-            A boolean tensor, same shape as ``token_ids``, indicating
-            which input observations are not ``torch.nan`` (i.e. not
-            missing nor padding).
-        """
-        raise NotImplementedError()
-
-    def output_transform(
-        self, samples: torch.Tensor, tokenizer_state: Any
-    ) -> torch.Tensor:
-        """
-        Turn a batch of sample token IDs into real values.
-
-        Parameters
-        ----------
-        samples
-            A tensor of integers, shaped (batch_size, num_samples, time_length),
-            containing token IDs of sample trajectories.
-        tokenizer_state
-            An object returned by ``input_transform`` containing
-            relevant context to decode samples, such as location and scale.
-            The nature of this depends on the specific tokenizer.
-
-        Returns
-        -------
-        forecasts
-            A real tensor, shaped (batch_size, num_samples, time_length),
-            containing forecasted sample paths.
-        """
-        raise NotImplementedError()
-
-
-class MeanScaleUniformBins(ChronosTokenizer):
-    def __init__(
-        self, low_limit: float, high_limit: float, config: ChronosConfig
-    ) -> None:
-        self.config = config
-        self.centers = torch.linspace(
-            low_limit,
-            high_limit,
-            config.n_tokens - config.n_special_tokens - 1,
-        )
-        self.boundaries = torch.concat(
-            (
-                torch.tensor([-1e20], device=self.centers.device),
-                (self.centers[1:] + self.centers[:-1]) / 2,
-                torch.tensor([1e20], device=self.centers.device),
-            )
-        )
-
-    def _input_transform(
-        self, context: torch.Tensor, scale: Optional[torch.Tensor] = None
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        context = context.to(dtype=torch.float32)
-        attention_mask = ~torch.isnan(context)
-
-        if scale is None:
-            scale = torch.nansum(
-                torch.abs(context) * attention_mask, dim=-1
-            ) / torch.nansum(attention_mask, dim=-1)
-            scale[~(scale > 0)] = 1.0
-
-        scaled_context = context / scale.unsqueeze(dim=-1)
-        token_ids = (
-            torch.bucketize(
-                input=scaled_context,
-                boundaries=self.boundaries,
-                # buckets are open to the right, see:
-                # https://pytorch.org/docs/2.1/generated/torch.bucketize.html#torch-bucketize
-                right=True,
-            )
-            + self.config.n_special_tokens
-        )
-
-        token_ids.clamp_(0, self.config.n_tokens - 1)
-
-        token_ids[~attention_mask] = self.config.pad_token_id
-
-        return token_ids, attention_mask, scale
-
-    def _append_eos_token(
-        self, token_ids: torch.Tensor, attention_mask: torch.Tensor
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        batch_size = token_ids.shape[0]
-        eos_tokens = torch.full((batch_size, 1), fill_value=self.config.eos_token_id)
-        token_ids = torch.concat((token_ids, eos_tokens), dim=1)
-        eos_mask = torch.full((batch_size, 1), fill_value=True)
-        attention_mask = torch.concat((attention_mask, eos_mask), dim=1)
-
-        return token_ids, attention_mask
-
-    def context_input_transform(
-        self, context: torch.Tensor
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        length = context.shape[-1]
-
-        if length > self.config.context_length:
-            context = context[..., -self.config.context_length :]
-
-        token_ids, attention_mask, scale = self._input_transform(context=context)
-
-        if self.config.use_eos_token and self.config.model_type == "seq2seq":
-            token_ids, attention_mask = self._append_eos_token(
-                token_ids=token_ids, attention_mask=attention_mask
-            )
-
-        return token_ids, attention_mask, scale
-
-    def label_input_transform(
-        self, label: torch.Tensor, scale: torch.Tensor
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        length = label.shape[-1]
-
-        assert length == self.config.prediction_length
-        token_ids, attention_mask, _ = self._input_transform(context=label, scale=scale)
-
-        if self.config.use_eos_token:
-            token_ids, attention_mask = self._append_eos_token(
-                token_ids=token_ids, attention_mask=attention_mask
-            )
-
-        return token_ids, attention_mask
-
-    def output_transform(
-        self, samples: torch.Tensor, scale: torch.Tensor
-    ) -> torch.Tensor:
-        scale_unsqueezed = scale.unsqueeze(-1).unsqueeze(-1)
-        indices = torch.clamp(
-            samples - self.config.n_special_tokens - 1,
-            min=0,
-            max=len(self.centers) - 1,
-        )
-        return self.centers[indices] * scale_unsqueezed
-
-
-class ChronosModel(nn.Module):
-    """
-    A ``ChronosModel`` wraps a ``PreTrainedModel`` object from ``transformers``
-    and uses it to predict sample paths for time series tokens.
-
-    Parameters
-    ----------
-    config
-        The configuration to use.
-    model
-        The pretrained model to use.
-    """
-
-    def __init__(self, config: ChronosConfig, model: PreTrainedModel) -> None:
-        super().__init__()
-        self.config = config
-        self.model = model
-
-    @property
-    def device(self):
-        return self.model.device
-
-    def encode(
-        self,
-        input_ids: torch.Tensor,
-        attention_mask: torch.Tensor,
-    ):
-        """
-        Extract the encoder embedding for the given token sequences.
-
-        Parameters
-        ----------
-        input_ids
-            Tensor of indices of input sequence tokens in the vocabulary
-            with shape (batch_size, sequence_length).
-        attention_mask
-            A mask tensor of the same shape as input_ids to avoid attending
-            on padding or missing tokens.
-
-        Returns
-        -------
-        embedding
-            A tensor of encoder embeddings with shape
-            (batch_size, sequence_length, d_model).
-        """
-        assert (
-            self.config.model_type == "seq2seq"
-        ), "Encoder embeddings are only supported for encoder-decoder models"
-        return self.model.encoder(
-            input_ids=input_ids, attention_mask=attention_mask
-        ).last_hidden_state
-
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        attention_mask: torch.Tensor,
-        prediction_length: Optional[int] = None,
-        num_samples: Optional[int] = None,
-        temperature: Optional[float] = None,
-        top_k: Optional[int] = None,
-        top_p: Optional[float] = None,
-    ) -> torch.Tensor:
-        """
-        Predict future sample tokens for the given token sequences.
-
-        Arguments ``prediction_length``, ``num_samples``, ``temperature``,
-        ``top_k``, ``top_p`` can be used to customize the model inference,
-        and default to the corresponding attributes in ``self.config`` if
-        not provided.
-
-        Returns
-        -------
-        samples
-            A tensor of integers, shaped (batch_size, num_samples, time_length),
-            containing forecasted sample paths.
-        """
-        if prediction_length is None:
-            prediction_length = self.config.prediction_length
-        if num_samples is None:
-            num_samples = self.config.num_samples
-        if temperature is None:
-            temperature = self.config.temperature
-        if top_k is None:
-            top_k = self.config.top_k
-        if top_p is None:
-            top_p = self.config.top_p
-
-        preds = self.model.generate(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            generation_config=GenerationConfig(
-                min_new_tokens=prediction_length,
-                max_new_tokens=prediction_length,
-                do_sample=True,
-                num_return_sequences=num_samples,
-                eos_token_id=self.config.eos_token_id,
-                pad_token_id=self.config.pad_token_id,
-                temperature=temperature,
-                top_k=top_k,
-                top_p=top_p,
-            ),
-        )
-
-        if self.config.model_type == "seq2seq":
-            preds = preds[..., 1:]  # remove the decoder start token
-        else:
-            assert self.config.model_type == "causal"
-            assert preds.size(-1) == input_ids.size(-1) + prediction_length
-            preds = preds[..., -prediction_length:]
-
-        return preds.reshape(input_ids.size(0), num_samples, -1)
-
-
-class ChronosPipeline(BaseChronosPipeline):
-    """
-    A ``ChronosPipeline`` uses the given tokenizer and model to forecast
-    input time series.
-
-    Use the ``from_pretrained`` class method to load serialized models.
-    Use the ``predict`` method to get forecasts.
-
-    Parameters
-    ----------
-    tokenizer
-        The tokenizer object to use.
-    model
-        The model to use.
-    """
-
-    tokenizer: ChronosTokenizer
-    model: ChronosModel
-    forecast_type: ForecastType = ForecastType.SAMPLES
-
-    def __init__(self, tokenizer, model):
-        super().__init__(inner_model=model.model)
-        self.tokenizer = tokenizer
-        self.model = model
-
-    def _prepare_and_validate_context(
-        self, context: Union[torch.Tensor, List[torch.Tensor]]
-    ):
-        if isinstance(context, list):
-            context = left_pad_and_stack_1D(context)
-        assert isinstance(context, torch.Tensor)
-        if context.ndim == 1:
-            context = context.unsqueeze(0)
-        assert context.ndim == 2
-
-        return context
-
-    @torch.no_grad()
-    def embed(
-        self, context: Union[torch.Tensor, List[torch.Tensor]]
-    ) -> Tuple[torch.Tensor, Any]:
-        """
-        Get encoder embeddings for the given time series.
-
-        Parameters
-        ----------
-        context
-            Input series. This is either a 1D tensor, or a list
-            of 1D tensors, or a 2D tensor whose first dimension
-            is batch. In the latter case, use left-padding with
-            ``torch.nan`` to align series of different lengths.
-
-        Returns
-        -------
-        embeddings, tokenizer_state
-            A tuple of two tensors: the encoder embeddings and the tokenizer_state,
-            e.g., the scale of the time series in the case of mean scaling.
-            The encoder embeddings are shaped (batch_size, context_length, d_model)
-            or (batch_size, context_length + 1, d_model), where context_length
-            is the size of the context along the time axis if a 2D tensor was provided
-            or the length of the longest time series, if a list of 1D tensors was
-            provided, and the extra 1 is for EOS.
-        """
-        context_tensor = self._prepare_and_validate_context(context=context)
-        token_ids, attention_mask, tokenizer_state = (
-            self.tokenizer.context_input_transform(context_tensor)
-        )
-        embeddings = self.model.encode(
-            input_ids=token_ids.to(self.model.device),
-            attention_mask=attention_mask.to(self.model.device),
-        ).cpu()
-        return embeddings, tokenizer_state
-
-    def predict(  # type: ignore[override]
-        self,
-        context: Union[torch.Tensor, List[torch.Tensor]],
-        prediction_length: Optional[int] = None,
-        num_samples: Optional[int] = None,
-        temperature: Optional[float] = None,
-        top_k: Optional[int] = None,
-        top_p: Optional[float] = None,
-        limit_prediction_length: bool = False,
-    ) -> torch.Tensor:
-        """
-        Get forecasts for the given time series.
-
-        Refer to the base method (``BaseChronosPipeline.predict``)
-        for details on shared parameters.
-
-        Additional parameters
-        ---------------------
-        num_samples
-            Number of sample paths to predict. Defaults to what
-            specified in ``self.model.config``.
-        temperature
-            Temperature to use for generating sample tokens.
-            Defaults to what specified in ``self.model.config``.
-        top_k
-            Top-k parameter to use for generating sample tokens.
-            Defaults to what specified in ``self.model.config``.
-        top_p
-            Top-p parameter to use for generating sample tokens.
-            Defaults to what specified in ``self.model.config``.
-        limit_prediction_length
-            Force prediction length smaller or equal than the
-            built-in prediction length from the model. False by
-            default. When true, fail loudly if longer predictions
-            are requested, otherwise longer predictions are allowed.
-
-        Returns
-        -------
-        samples
-            Tensor of sample forecasts, of shape
-            (batch_size, num_samples, prediction_length).
-        """
-        context_tensor = self._prepare_and_validate_context(context=context)
-
-        if prediction_length is None:
-            prediction_length = self.model.config.prediction_length
-
-        if prediction_length > self.model.config.prediction_length:
-            msg = (
-                f"We recommend keeping prediction length <= {self.model.config.prediction_length}. "
-                "The quality of longer predictions may degrade since the model is not optimized for it. "
-            )
-            if limit_prediction_length:
-                msg += "You can turn off this check by setting `limit_prediction_length=False`."
-                raise ValueError(msg)
-            logger.warning(msg)
-
-        predictions = []
-        remaining = prediction_length
-
-        while remaining > 0:
-            token_ids, attention_mask, scale = self.tokenizer.context_input_transform(
-                context_tensor
-            )
-            samples = self.model(
-                token_ids.to(self.model.device),
-                attention_mask.to(self.model.device),
-                min(remaining, self.model.config.prediction_length),
-                num_samples,
-                temperature,
-                top_k,
-                top_p,
-            )
-            prediction = self.tokenizer.output_transform(
-                samples.to(scale.device), scale
-            )
-
-            predictions.append(prediction)
-            remaining -= prediction.shape[-1]
-
-            if remaining <= 0:
-                break
-
-            context_tensor = torch.cat(
-                [context_tensor, prediction.median(dim=1).values], dim=-1
-            )
-
-        return torch.cat(predictions, dim=-1).to(dtype=torch.float32, device="cpu")
-
-    def predict_quantiles(
-        self,
-        context: Union[torch.Tensor, List[torch.Tensor]],
-        prediction_length: Optional[int] = None,
-        quantile_levels: List[float] = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9],
-        **predict_kwargs,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Refer to the base method (``BaseChronosPipeline.predict_quantiles``).
-        """
-        prediction_samples = (
-            self.predict(context, prediction_length=prediction_length, **predict_kwargs)
-            .detach()
-            .swapaxes(1, 2)
-        )
-        mean = prediction_samples.mean(dim=-1)
-        quantiles = torch.quantile(
-            prediction_samples,
-            q=torch.tensor(quantile_levels, dtype=prediction_samples.dtype),
-            dim=-1,
-        ).permute(1, 2, 0)
-
-        return quantiles, mean
-
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        """
-        Load the model, either from a local path or from the HuggingFace Hub.
-        Supports the same arguments as ``AutoConfig`` and ``AutoModel``
-        from ``transformers``.
-        """
-
-        config = AutoConfig.from_pretrained(*args, **kwargs)
-
-        assert hasattr(config, "chronos_config"), "Not a Chronos config file"
-
-        chronos_config = ChronosConfig(**config.chronos_config)
-
-        if chronos_config.model_type == "seq2seq":
-            inner_model = AutoModelForSeq2SeqLM.from_pretrained(*args, **kwargs)
-        else:
-            assert chronos_config.model_type == "causal"
-            inner_model = AutoModelForCausalLM.from_pretrained(*args, **kwargs)
-
-        return cls(
-            tokenizer=chronos_config.create_tokenizer(),
-            model=ChronosModel(config=chronos_config, model=inner_model),
-        )
--- a/src/chronos/chronos_bolt.py
+++ b/src/chronos/chronos_bolt.py
@ -1,640 +0,0 @@
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-# Authors: Abdul Fatir Ansari <ansarnd@amazon.com>, Caner Turkmen <atturkm@amazon.com>, Lorenzo Stella <stellalo@amazon.com>
-# Original source:
-# https://github.com/autogluon/autogluon/blob/f57beb26cb769c6e0d484a6af2b89eab8aee73a8/timeseries/src/autogluon/timeseries/models/chronos/pipeline/chronos_bolt.py
-
-import copy
-import logging
-import warnings
-from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.nn as nn
-from transformers import AutoConfig
-from transformers.models.t5.modeling_t5 import (
-    ACT2FN,
-    T5Config,
-    T5LayerNorm,
-    T5PreTrainedModel,
-    T5Stack,
-)
-from transformers.utils import ModelOutput
-
-from .base import BaseChronosPipeline, ForecastType
-
-logger = logging.getLogger(__file__)
-
-
-@dataclass
-class ChronosBoltConfig:
-    context_length: int
-    prediction_length: int
-    input_patch_size: int
-    input_patch_stride: int
-    quantiles: List[float]
-    use_reg_token: bool = False
-
-
-@dataclass
-class ChronosBoltOutput(ModelOutput):
-    loss: Optional[torch.Tensor] = None
-    quantile_preds: Optional[torch.Tensor] = None
-    attentions: Optional[torch.Tensor] = None
-    cross_attentions: Optional[torch.Tensor] = None
-
-
-class Patch(nn.Module):
-    def __init__(self, patch_size: int, patch_stride: int) -> None:
-        super().__init__()
-        self.patch_size = patch_size
-        self.patch_stride = patch_stride
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        length = x.shape[-1]
-
-        if length % self.patch_size != 0:
-            padding_size = (
-                *x.shape[:-1],
-                self.patch_size - (length % self.patch_size),
-            )
-            padding = torch.full(
-                size=padding_size, fill_value=torch.nan, dtype=x.dtype, device=x.device
-            )
-            x = torch.concat((padding, x), dim=-1)
-
-        x = x.unfold(dimension=-1, size=self.patch_size, step=self.patch_stride)
-        return x
-
-
-class InstanceNorm(nn.Module):
-    """
-    See, also, RevIN. Apply standardization along the last dimension.
-    """
-
-    def __init__(self, eps: float = 1e-5) -> None:
-        super().__init__()
-        self.eps = eps
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        loc_scale: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
-        if loc_scale is None:
-            loc = torch.nan_to_num(torch.nanmean(x, dim=-1, keepdim=True), nan=0.0)
-            scale = torch.nan_to_num(
-                torch.nanmean((x - loc).square(), dim=-1, keepdim=True).sqrt(), nan=1.0
-            )
-            scale = torch.where(scale == 0, torch.abs(loc) + self.eps, scale)
-        else:
-            loc, scale = loc_scale
-
-        return (x - loc) / scale, (loc, scale)
-
-    def inverse(
-        self, x: torch.Tensor, loc_scale: Tuple[torch.Tensor, torch.Tensor]
-    ) -> torch.Tensor:
-        loc, scale = loc_scale
-        return x * scale + loc
-
-
-class ResidualBlock(nn.Module):
-    def __init__(
-        self,
-        in_dim: int,
-        h_dim: int,
-        out_dim: int,
-        act_fn_name: str,
-        dropout_p: float = 0.0,
-        use_layer_norm: bool = False,
-    ) -> None:
-        super().__init__()
-
-        self.dropout = nn.Dropout(dropout_p)
-        self.hidden_layer = nn.Linear(in_dim, h_dim)
-        self.act = ACT2FN[act_fn_name]
-        self.output_layer = nn.Linear(h_dim, out_dim)
-        self.residual_layer = nn.Linear(in_dim, out_dim)
-
-        self.use_layer_norm = use_layer_norm
-        if use_layer_norm:
-            self.layer_norm = T5LayerNorm(out_dim)
-
-    def forward(self, x: torch.Tensor):
-        hid = self.act(self.hidden_layer(x))
-        out = self.dropout(self.output_layer(hid))
-        res = self.residual_layer(x)
-
-        out = out + res
-
-        if self.use_layer_norm:
-            return self.layer_norm(out)
-        return out
-
-
-class ChronosBoltModelForForecasting(T5PreTrainedModel):
-    _keys_to_ignore_on_load_missing = [
-        r"input_patch_embedding\.",
-        r"output_patch_embedding\.",
-    ]
-    _keys_to_ignore_on_load_unexpected = [r"lm_head.weight"]
-    _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight"]
-
-    def __init__(self, config: T5Config):
-        assert hasattr(config, "chronos_config"), "Not a Chronos config file"
-
-        super().__init__(config)
-        self.model_dim = config.d_model
-
-        self.chronos_config = ChronosBoltConfig(**config.chronos_config)
-
-        # Only decoder_start_id (and optionally REG token)
-        if self.chronos_config.use_reg_token:
-            config.reg_token_id = 1
-
-        config.vocab_size = 2 if self.chronos_config.use_reg_token else 1
-        self.shared = nn.Embedding(config.vocab_size, config.d_model)
-
-        # Input patch embedding layer
-        self.input_patch_embedding = ResidualBlock(
-            in_dim=self.chronos_config.input_patch_size * 2,
-            h_dim=config.d_ff,
-            out_dim=config.d_model,
-            act_fn_name=config.dense_act_fn,
-            dropout_p=config.dropout_rate,
-        )
-
-        # patching layer
-        self.patch = Patch(
-            patch_size=self.chronos_config.input_patch_size,
-            patch_stride=self.chronos_config.input_patch_stride,
-        )
-
-        # instance normalization, also referred to as "scaling" in Chronos and GluonTS
-        self.instance_norm = InstanceNorm()
-
-        encoder_config = copy.deepcopy(config)
-        encoder_config.is_decoder = False
-        encoder_config.use_cache = False
-        encoder_config.is_encoder_decoder = False
-        self.encoder = T5Stack(encoder_config, self.shared)
-
-        self._init_decoder(config)
-
-        self.num_quantiles = len(self.chronos_config.quantiles)
-        quantiles = torch.tensor(self.chronos_config.quantiles, dtype=self.dtype)
-        self.register_buffer("quantiles", quantiles, persistent=False)
-
-        self.output_patch_embedding = ResidualBlock(
-            in_dim=config.d_model,
-            h_dim=config.d_ff,
-            out_dim=self.num_quantiles * self.chronos_config.prediction_length,
-            act_fn_name=config.dense_act_fn,
-            dropout_p=config.dropout_rate,
-        )
-
-        # Initialize weights and apply final processing
-        self.post_init()
-
-        # Model parallel
-        self.model_parallel = False
-        self.device_map = None
-
-    def _init_weights(self, module):
-        super()._init_weights(module)
-        """Initialize the weights"""
-        factor = self.config.initializer_factor
-        if isinstance(module, (self.__class__)):
-            module.shared.weight.data.normal_(mean=0.0, std=factor * 1.0)
-        elif isinstance(module, ResidualBlock):
-            module.hidden_layer.weight.data.normal_(
-                mean=0.0,
-                std=factor * ((self.chronos_config.input_patch_size * 2) ** -0.5),
-            )
-            if (
-                hasattr(module.hidden_layer, "bias")
-                and module.hidden_layer.bias is not None
-            ):
-                module.hidden_layer.bias.data.zero_()
-
-            module.residual_layer.weight.data.normal_(
-                mean=0.0,
-                std=factor * ((self.chronos_config.input_patch_size * 2) ** -0.5),
-            )
-            if (
-                hasattr(module.residual_layer, "bias")
-                and module.residual_layer.bias is not None
-            ):
-                module.residual_layer.bias.data.zero_()
-
-            module.output_layer.weight.data.normal_(
-                mean=0.0, std=factor * ((self.config.d_ff) ** -0.5)
-            )
-            if (
-                hasattr(module.output_layer, "bias")
-                and module.output_layer.bias is not None
-            ):
-                module.output_layer.bias.data.zero_()
-
-    def encode(
-        self, context: torch.Tensor, mask: Optional[torch.Tensor] = None
-    ) -> Tuple[
-        torch.Tensor, Tuple[torch.Tensor, torch.Tensor], torch.Tensor, torch.Tensor
-    ]:
-        mask = (
-            mask.to(context.dtype)
-            if mask is not None
-            else torch.isnan(context).logical_not().to(context.dtype)
-        )
-
-        batch_size, _ = context.shape
-        if context.shape[-1] > self.chronos_config.context_length:
-            context = context[..., -self.chronos_config.context_length :]
-            mask = mask[..., -self.chronos_config.context_length :]
-
-        # scaling
-        context, loc_scale = self.instance_norm(context)
-
-        # the scaling op above is done in 32-bit precision,
-        # then the context is moved to model's dtype
-        context = context.to(self.dtype)
-        mask = mask.to(self.dtype)
-
-        # patching
-        patched_context = self.patch(context)
-        patched_mask = torch.nan_to_num(self.patch(mask), nan=0.0)
-        patched_context = torch.where(patched_mask > 0.0, patched_context, 0.0)
-        # concat context and mask along patch dim
-        patched_context = torch.cat([patched_context, patched_mask], dim=-1)
-
-        # attention_mask = 1 if at least one item in the patch is observed
-        attention_mask = (
-            patched_mask.sum(dim=-1) > 0
-        )  # (batch_size, patched_seq_length)
-
-        input_embeds = self.input_patch_embedding(patched_context)
-
-        if self.chronos_config.use_reg_token:
-            # Append [REG]
-            reg_input_ids = torch.full(
-                (batch_size, 1),
-                self.config.reg_token_id,
-                device=input_embeds.device,
-            )
-            reg_embeds = self.shared(reg_input_ids)
-            input_embeds = torch.cat([input_embeds, reg_embeds], dim=-2)
-            attention_mask = torch.cat(
-                [
-                    attention_mask.to(self.dtype),
-                    torch.ones_like(reg_input_ids).to(self.dtype),
-                ],
-                dim=-1,
-            )
-
-        encoder_outputs = self.encoder(
-            attention_mask=attention_mask,
-            inputs_embeds=input_embeds,
-        )
-
-        return encoder_outputs[0], loc_scale, input_embeds, attention_mask
-
-    def forward(
-        self,
-        context: torch.Tensor,
-        mask: Optional[torch.Tensor] = None,
-        target: Optional[torch.Tensor] = None,
-        target_mask: Optional[torch.Tensor] = None,
-    ) -> ChronosBoltOutput:
-        batch_size = context.size(0)
-
-        hidden_states, loc_scale, input_embeds, attention_mask = self.encode(
-            context=context, mask=mask
-        )
-        sequence_output = self.decode(input_embeds, attention_mask, hidden_states)
-
-        quantile_preds_shape = (
-            batch_size,
-            self.num_quantiles,
-            self.chronos_config.prediction_length,
-        )
-        quantile_preds = self.output_patch_embedding(sequence_output).view(
-            *quantile_preds_shape
-        )
-
-        loss = None
-        if target is not None:
-            # normalize target
-            target, _ = self.instance_norm(target, loc_scale)
-            target = target.unsqueeze(1)  # type: ignore
-            assert self.chronos_config.prediction_length >= target.shape[-1]
-
-            target = target.to(quantile_preds.device)
-            target_mask = (
-                target_mask.unsqueeze(1).to(quantile_preds.device)
-                if target_mask is not None
-                else ~torch.isnan(target)
-            )
-            target[~target_mask] = 0.0
-
-            # pad target and target_mask if they are shorter than model's prediction_length
-            if self.chronos_config.prediction_length > target.shape[-1]:
-                padding_shape = (
-                    *target.shape[:-1],
-                    self.chronos_config.prediction_length - target.shape[-1],
-                )
-                target = torch.cat(
-                    [target, torch.zeros(padding_shape).to(target)], dim=-1
-                )
-                target_mask = torch.cat(
-                    [target_mask, torch.zeros(padding_shape).to(target_mask)], dim=-1
-                )
-
-            loss = (
-                2
-                * torch.abs(
-                    (target - quantile_preds)
-                    * (
-                        (target <= quantile_preds).float()
-                        - self.quantiles.view(1, self.num_quantiles, 1)
-                    )
-                )
-                * target_mask.float()
-            )
-            loss = loss.mean(dim=-2)  # Mean over prediction horizon
-            loss = loss.sum(dim=-1)  # Sum over quantile levels
-            loss = loss.mean()  # Mean over batch
-
-        # Unscale predictions
-        quantile_preds = self.instance_norm.inverse(
-            quantile_preds.view(batch_size, -1),
-            loc_scale,
-        ).view(*quantile_preds_shape)
-
-        return ChronosBoltOutput(
-            loss=loss,
-            quantile_preds=quantile_preds,
-        )
-
-    def _init_decoder(self, config):
-        decoder_config = copy.deepcopy(config)
-        decoder_config.is_decoder = True
-        decoder_config.is_encoder_decoder = False
-        decoder_config.num_layers = config.num_decoder_layers
-        self.decoder = T5Stack(decoder_config, self.shared)
-
-    def decode(
-        self,
-        input_embeds,
-        attention_mask,
-        hidden_states,
-        output_attentions=False,
-    ):
-        """
-        Parameters
-        ----------
-        input_embeds: torch.Tensor
-            Patched and embedded inputs. Shape (batch_size, patched_context_length, d_model)
-        attention_mask: torch.Tensor
-            Attention mask for the patched context. Shape (batch_size, patched_context_length), type: torch.int64
-        hidden_states: torch.Tensor
-            Hidden states returned by the encoder. Shape (batch_size, patched_context_length, d_model)
-
-        Returns
-        -------
-        last_hidden_state
-            Last hidden state returned by the decoder, of shape (batch_size, 1, d_model)
-        """
-        batch_size = input_embeds.shape[0]
-        decoder_input_ids = torch.full(
-            (batch_size, 1),
-            self.config.decoder_start_token_id,
-            device=input_embeds.device,
-        )
-        decoder_outputs = self.decoder(
-            input_ids=decoder_input_ids,
-            encoder_hidden_states=hidden_states,
-            encoder_attention_mask=attention_mask,
-            output_attentions=output_attentions,
-            return_dict=True,
-        )
-
-        return decoder_outputs.last_hidden_state  # sequence_outputs, b x 1 x d_model
-
-
-class ChronosBoltPipeline(BaseChronosPipeline):
-    forecast_type: ForecastType = ForecastType.QUANTILES
-    default_context_length: int = 2048
-
-    def __init__(self, model: ChronosBoltModelForForecasting):
-        super().__init__(inner_model=model)
-        self.model = model
-
-    @property
-    def quantiles(self) -> List[float]:
-        return self.model.config.chronos_config["quantiles"]
-
-    @torch.no_grad()
-    def embed(
-        self, context: Union[torch.Tensor, List[torch.Tensor]]
-    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
-        """
-        Get encoder embeddings for the given time series.
-
-        Parameters
-        ----------
-        context
-            Input series. This is either a 1D tensor, or a list
-            of 1D tensors, or a 2D tensor whose first dimension
-            is batch. In the latter case, use left-padding with
-            ``torch.nan`` to align series of different lengths.
-
-        Returns
-        -------
-        embeddings, loc_scale
-            A tuple of two items: the encoder embeddings and the loc_scale,
-            i.e., the mean and std of the original time series.
-            The encoder embeddings are shaped (batch_size, num_patches + 1, d_model),
-            where num_patches is the number of patches in the time series
-            and the extra 1 is for the [REG] token (if used by the model).
-        """
-        context_tensor = self._prepare_and_validate_context(context=context)
-        model_context_length = self.model.config.chronos_config["context_length"]
-
-        if context_tensor.shape[-1] > model_context_length:
-            context_tensor = context_tensor[..., -model_context_length:]
-
-        context_tensor = context_tensor.to(
-            device=self.model.device,
-            dtype=torch.float32,
-        )
-        embeddings, loc_scale, *_ = self.model.encode(context=context_tensor)
-        return embeddings.cpu(), (
-            loc_scale[0].squeeze(-1).cpu(),
-            loc_scale[1].squeeze(-1).cpu(),
-        )
-
-    def predict(  # type: ignore[override]
-        self,
-        context: Union[torch.Tensor, List[torch.Tensor]],
-        prediction_length: Optional[int] = None,
-        limit_prediction_length: bool = False,
-    ) -> torch.Tensor:
-        """
-        Get forecasts for the given time series.
-
-        Refer to the base method (``BaseChronosPipeline.predict``)
-        for details on shared parameters.
-        Additional parameters
-        ---------------------
-        limit_prediction_length
-            Force prediction length smaller or equal than the
-            built-in prediction length from the model. False by
-            default. When true, fail loudly if longer predictions
-            are requested, otherwise longer predictions are allowed.
-
-        Returns
-        -------
-        torch.Tensor
-            Forecasts of shape (batch_size, num_quantiles, prediction_length)
-            where num_quantiles is the number of quantiles the model has been
-            trained to output. For official Chronos-Bolt models, the value of
-            num_quantiles is 9 for [0.1, 0.2, ..., 0.9]-quantiles.
-
-        Raises
-        ------
-        ValueError
-            When limit_prediction_length is True and the prediction_length is
-            greater than model's trainig prediction_length.
-        """
-        context_tensor = self._prepare_and_validate_context(context=context)
-
-        model_context_length = self.model.config.chronos_config["context_length"]
-        model_prediction_length = self.model.config.chronos_config["prediction_length"]
-        if prediction_length is None:
-            prediction_length = model_prediction_length
-
-        if prediction_length > model_prediction_length:
-            msg = (
-                f"We recommend keeping prediction length <= {model_prediction_length}. "
-                "The quality of longer predictions may degrade since the model is not optimized for it. "
-            )
-            if limit_prediction_length:
-                msg += "You can turn off this check by setting `limit_prediction_length=False`."
-                raise ValueError(msg)
-            warnings.warn(msg)
-
-        predictions = []
-        remaining = prediction_length
-
-        # We truncate the context here because otherwise batches with very long
-        # context could take up large amounts of GPU memory unnecessarily.
-        if context_tensor.shape[-1] > model_context_length:
-            context_tensor = context_tensor[..., -model_context_length:]
-
-        # TODO: We unroll the forecast of Chronos Bolt greedily with the full forecast
-        # horizon that the model was trained with (i.e., 64). This results in variance collapsing
-        # every 64 steps.
-        context_tensor = context_tensor.to(
-            device=self.model.device,
-            dtype=torch.float32,
-        )
-        while remaining > 0:
-            with torch.no_grad():
-                prediction = self.model(
-                    context=context_tensor,
-                ).quantile_preds.to(context_tensor)
-
-            predictions.append(prediction)
-            remaining -= prediction.shape[-1]
-
-            if remaining <= 0:
-                break
-
-            central_idx = torch.abs(torch.tensor(self.quantiles) - 0.5).argmin()
-            central_prediction = prediction[:, central_idx]
-
-            context_tensor = torch.cat([context_tensor, central_prediction], dim=-1)
-
-        return torch.cat(predictions, dim=-1)[..., :prediction_length].to(
-            dtype=torch.float32, device="cpu"
-        )
-
-    def predict_quantiles(
-        self,
-        context: Union[torch.Tensor, List[torch.Tensor]],
-        prediction_length: Optional[int] = None,
-        quantile_levels: List[float] = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9],
-        **predict_kwargs,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Refer to the base method (``BaseChronosPipeline.predict_quantiles``).
-        """
-        # shape (batch_size, prediction_length, len(training_quantile_levels))
-        predictions = (
-            self.predict(context, prediction_length=prediction_length, **predict_kwargs)
-            .detach()
-            .swapaxes(1, 2)
-        )
-
-        training_quantile_levels = self.quantiles
-
-        if set(quantile_levels).issubset(set(training_quantile_levels)):
-            # no need to perform intra/extrapolation
-            quantiles = predictions[
-                ..., [training_quantile_levels.index(q) for q in quantile_levels]
-            ]
-        else:
-            # we rely on torch for interpolating quantiles if quantiles that
-            # Chronos Bolt was trained on were not provided
-            if min(quantile_levels) < min(training_quantile_levels) or max(
-                quantile_levels
-            ) > max(training_quantile_levels):
-                logger.warning(
-                    f"\tQuantiles to be predicted ({quantile_levels}) are not within the range of "
-                    f"quantiles that Chronos-Bolt was trained on ({training_quantile_levels}). "
-                    "Quantile predictions will be set to the minimum/maximum levels at which Chronos-Bolt "
-                    "was trained on. This may significantly affect the quality of the predictions."
-                )
-
-            # TODO: this is a hack that assumes the model's quantiles during training (training_quantile_levels)
-            # made up an equidistant grid along the quantile dimension. i.e., they were (0.1, 0.2, ..., 0.9).
-            # While this holds for official Chronos-Bolt models, this may not be true in the future, and this
-            # function may have to be revised.
-            augmented_predictions = torch.cat(
-                [predictions[..., [0]], predictions, predictions[..., [-1]]],
-                dim=-1,
-            )
-            quantiles = torch.quantile(
-                augmented_predictions,
-                q=torch.tensor(quantile_levels, dtype=augmented_predictions.dtype),
-                dim=-1,
-            ).permute(1, 2, 0)
-        # NOTE: the median is returned as the mean here
-        mean = predictions[:, :, training_quantile_levels.index(0.5)]
-        return quantiles, mean
-
-    @classmethod
-    def from_pretrained(cls, *args, **kwargs):
-        """
-        Load the model, either from a local path or from the HuggingFace Hub.
-        Supports the same arguments as ``AutoConfig`` and ``AutoModel``
-        from ``transformers``.
-        """
-
-        config = AutoConfig.from_pretrained(*args, **kwargs)
-        assert hasattr(config, "chronos_config"), "Not a Chronos config file"
-
-        architecture = config.architectures[0]
-        class_ = globals().get(architecture)
-
-        if class_ is None:
-            logger.warning(
-                f"Unknown architecture: {architecture}, defaulting to ChronosBoltModelForForecasting"
-            )
-            class_ = ChronosBoltModelForForecasting
-
-        model = class_.from_pretrained(*args, **kwargs)
-        return cls(model=model)
--- a/src/chronos/utils.py
+++ b/src/chronos/utils.py
@ -1,20 +0,0 @@
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-
-from typing import List
-
-import torch
-
-
-def left_pad_and_stack_1D(tensors: List[torch.Tensor]) -> torch.Tensor:
-    max_len = max(len(c) for c in tensors)
-    padded = []
-    for c in tensors:
-        assert isinstance(c, torch.Tensor)
-        assert c.ndim == 1
-        padding = torch.full(
-            size=(max_len - len(c),), fill_value=torch.nan, device=c.device
-        )
-        padded.append(torch.concat((padding, c), dim=-1))
-    return torch.stack(padded)
--- a/src/chronosx/chronosx.py
+++ b/src/chronosx/chronosx.py
@ -0,0 +1,706 @@
+import inspect
+import logging
+import numpy as np
+import torch
+import itertools
+
+from chronos import BaseChronosPipeline
+from chronos import ChronosConfig
+from chronos import ChronosPipeline
+from pathlib import Path
+from tqdm.auto import tqdm
+from transformers.generation.configuration_utils import GenerationConfig
+from transformers.utils.import_utils import is_accelerate_available
+from transformers import (
+    AutoConfig,
+    T5ForConditionalGeneration,
+    Trainer,
+    TrainingArguments,
+)
+from typing import Any, Dict, Optional
+
+from gluonts.itertools import batcher
+from gluonts.model.forecast import SampleForecast
+from gluonts.transform import (
+    InstanceSplitter,
+    TestSplitSampler,
+)
+
+from chronosx.utils.transform import create_transformation
+
+if is_accelerate_available():
+    from accelerate.hooks import AlignDevicesHook, add_hook_to_module
+
+from chronosx.injection_blocks.block_mapping import injection_blocks_map
+from chronosx.utils.utils import count_parameters, compute_metrics, log_on_main
+from chronosx.utils.prepare_covariates import prepare_covariates
+
+logging.basicConfig(format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__file__)
+logger.setLevel(logging.INFO)
+
+
+class ChronosX(T5ForConditionalGeneration):
+    def __init__(self, *args, **kwargs):
+        T5ForConditionalGeneration.__init__(self, *args, **kwargs)
+
+        self.output_hidden_states = False
+
+        if self.input_injection_class:
+            self.input_injection_block = self.input_injection_class(
+                hidden_dim=self.hidden_dim,
+                model_dim=self.model_dim,
+                num_covariates=self.num_covariates,
+                num_layers=self.num_layers,
+            )
+            self.input_injection_block_decoder = self.input_injection_class(
+                hidden_dim=self.hidden_dim,
+                model_dim=self.model_dim,
+                num_covariates=self.num_covariates,
+                num_layers=self.num_layers,
+            )
+        else:
+            self.input_injection_block = None
+            self.input_injection_block_decoder = None
+
+        if self.output_injection_class:
+            self.output_injection_block = self.output_injection_class(
+                hidden_dim=self.hidden_dim,
+                model_dim=self.model_dim,
+                num_covariates=self.num_covariates,
+                num_layers=self.num_layers,
+                vocab_size=self.vocab_size,
+            )
+            self.output_hidden_states = True
+        else:
+            self.output_injection_block = None
+
+    @classmethod
+    def set_state(
+        cls,
+        num_covariates,
+        covariate_injection,
+        hidden_dim,
+        num_layers,
+        vocab_size,
+        model_dim,
+    ):
+
+        cls.num_covariates = num_covariates
+        cls.covariate_injection = covariate_injection
+        cls.hidden_dim = hidden_dim
+        cls.num_layers = num_layers
+        cls.vocab_size = vocab_size
+        cls.model_dim = model_dim
+
+        if cls.covariate_injection in injection_blocks_map.keys():
+            (cls.input_injection_class, cls.output_injection_class) = (
+                injection_blocks_map[cls.covariate_injection]
+            )
+        else:
+            cls.input_injection_class = None
+            cls.output_injection_class = None
+
+        return cls
+
+    def initialize_blocks(self):
+        if self.input_injection_block:
+            self.input_injection_block.initialize_modules()
+            self.input_injection_block_decoder.initialize_modules()
+
+        if self.output_injection_block:
+            self.output_injection_block.initialize_modules()
+
+    def freeze(self, layer_name=None):
+        for name, param in self.named_parameters():
+            if layer_name == "all" or layer_name in name:
+                param.requires_grad = False
+
+    def unfreeze(self, layer_name=None):
+        for name, param in self.named_parameters():
+            if layer_name == "all" or layer_name in name:
+                param.requires_grad = True
+
+    def _inject_at_input(
+        self,
+        input_ids: torch.Tensor = None,
+        decoder_input_ids: torch.Tensor = None,
+        past_covariates: torch.Tensor = None,
+        future_covariates: torch.Tensor = None,
+        labels: torch.Tensor = None,
+    ):
+        # input injection
+        inputs_embeds = None
+        decoder_inputs_embeds = None
+        if input_ids is not None:
+            inputs_embeds = self.encoder.embed_tokens(input_ids)
+            inputs_embeds = self.input_injection_block(inputs_embeds, past_covariates)
+            input_ids = None
+
+        if decoder_input_ids is None and labels is not None:
+            decoder_input_ids = self._shift_right(labels)
+
+        if future_covariates is not None and decoder_input_ids is not None:
+            decoder_inputs_embeds = self.decoder.embed_tokens(decoder_input_ids)
+
+            # shifting covariates
+            shifted_future_covariates = self._shift_right(
+                future_covariates.transpose(1, 2)
+            ).transpose(1, 2)
+
+            decoder_inputs_embeds = self.input_injection_block(
+                decoder_inputs_embeds, shifted_future_covariates, is_decoder=True
+            )
+            decoder_input_ids = None
+
+        return (input_ids, decoder_input_ids, inputs_embeds, decoder_inputs_embeds)
+
+    def _inject_at_output(self, output, labels, future_covariates):
+        last_hidden_state = output.decoder_hidden_states[-1]
+        if future_covariates is not None:
+            output.logits, output.loss = self.output_injection_block(
+                future_covariates=future_covariates,
+                labels=labels,
+                logits=output.logits,
+                last_hidden_state=last_hidden_state,
+            )
+
+        return output
+
+    def forward(
+        self,
+        input_ids: torch.Tensor = None,
+        decoder_input_ids: torch.Tensor = None,
+        inputs_embeds: torch.Tensor = None,
+        decoder_inputs_embeds: torch.Tensor = None,
+        past_covariates: torch.Tensor = None,
+        future_covariates: torch.Tensor = None,
+        labels: torch.Tensor = None,
+        **kwargs,
+    ):
+
+        if self.input_injection_block is not None:
+            (input_ids, decoder_input_ids, inputs_embeds, decoder_inputs_embeds) = (
+                self._inject_at_input(
+                    input_ids=input_ids,
+                    decoder_input_ids=decoder_input_ids,
+                    past_covariates=past_covariates,
+                    future_covariates=future_covariates,
+                    labels=labels,
+                )
+            )
+
+        # Removing key 'num_items_in_batch' from kwargs.
+        # This is necessary as recent versions of transformers make the code break with it.
+        kwargs_filtered = kwargs.copy()
+        if "num_items_in_batch" in kwargs_filtered:
+            kwargs_filtered.pop("num_items_in_batch")
+
+        output = super(ChronosX, self).forward(
+            input_ids=input_ids,
+            decoder_input_ids=decoder_input_ids,
+            inputs_embeds=inputs_embeds,
+            decoder_inputs_embeds=decoder_inputs_embeds,
+            labels=labels,
+            output_hidden_states=self.output_hidden_states,
+            **kwargs_filtered,
+        )
+
+        if self.output_injection_block is not None:
+            output = self._inject_at_output(
+                output=output, labels=labels, future_covariates=future_covariates
+            )
+
+        return output
+
+    def generate(self, **kwargs):
+        if self.output_injection_block is not None:
+            self.output_injection_block.restart_generator_counter()
+
+        if self.input_injection_block is not None:
+            self.input_injection_block.restart_generator_counter()
+            self.input_injection_block_decoder.restart_generator_counter()
+
+        output = super(ChronosX, self).generate(**kwargs)
+
+        if self.output_injection_block is not None:
+            self.output_injection_block.generating = False
+
+        if self.input_injection_block is not None:
+            self.input_injection_block.generating = False
+            self.input_injection_block_decoder.generating = False
+
+        return output
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        head_mask=None,
+        decoder_head_mask=None,
+        decoder_attention_mask=None,
+        cross_attn_head_mask=None,
+        use_cache=None,
+        encoder_outputs=None,
+        future_covariates=None,
+        **kwargs,
+    ):
+
+        kwargs_augmented = {
+            **kwargs,
+            "past_key_values": past_key_values,
+            "attention_mask": attention_mask,
+            "head_mask": head_mask,
+            "decoder_head_mask": decoder_head_mask,
+            "decoder_attention_mask": decoder_attention_mask,
+            "cross_attn_head_mask": cross_attn_head_mask,
+            "use_cache": use_cache,
+            "encoder_outputs": encoder_outputs,
+        }
+
+        output = super(ChronosX, self).prepare_inputs_for_generation(
+            input_ids,
+            **kwargs_augmented,
+        )
+
+        output.update({"future_covariates": future_covariates})
+
+        return output
+
+    def _prepare_encoder_decoder_kwargs_for_generation(
+        self,
+        inputs_tensor: torch.Tensor,
+        model_kwargs,
+        model_input_name: Optional[str],
+        generation_config: GenerationConfig,
+    ) -> Dict[str, Any]:
+        # 1. get encoder
+        encoder = self.get_encoder()
+        # Compatibility with Accelerate big model inference: we need the encoder to outputs stuff on the same device
+        # as the inputs.
+        if hasattr(self, "hf_device_map"):
+            if hasattr(encoder, "_hf_hook"):
+                encoder._hf_hook.io_same_device = True
+            else:
+                add_hook_to_module(encoder, AlignDevicesHook(io_same_device=True))
+
+        # 2. Prepare encoder args and encoder kwargs from model kwargs and generation config.
+        irrelevant_prefix = ["decoder_", "cross_attn", "use_cache"]
+        encoder_kwargs = {
+            argument: value
+            for argument, value in model_kwargs.items()
+            if not any(argument.startswith(p) for p in irrelevant_prefix)
+        }
+        encoder_signature = set(inspect.signature(encoder.forward).parameters)
+        encoder_accepts_wildcard = (
+            "kwargs" in encoder_signature or "model_kwargs" in encoder_signature
+        )
+        if not encoder_accepts_wildcard:
+            encoder_kwargs = {
+                argument: value
+                for argument, value in encoder_kwargs.items()
+                if argument in encoder_signature
+            }
+        encoder_kwargs["output_attentions"] = generation_config.output_attentions
+        encoder_kwargs["output_hidden_states"] = generation_config.output_hidden_states
+
+        # 3. make sure that encoder returns `ModelOutput`
+        model_input_name = (
+            model_input_name if model_input_name is not None else self.main_input_name
+        )
+        encoder_kwargs["return_dict"] = True
+        encoder_kwargs[model_input_name] = inputs_tensor
+        model_kwargs["encoder_outputs"]: ModelOutput = encoder(**encoder_kwargs)  # type: ignore
+
+        # IMPORTANT: here is the place where we provide the updated token embeddings for Input Injection Block
+        # Remember that we inject covariates on the token embeddings. This means that
+        # we do not have to give to the encoder the input_ids but rather the updated token embeddings.
+        # That's why we remove input_ids from *encoder_kwargs* and rather use 'inputs_embeds'
+        if self.input_injection_block is not None:
+            inputs_embeds = self.encoder.embed_tokens(encoder_kwargs["input_ids"])
+            inputs_embeds = self.input_injection_block(
+                inputs_embeds, model_kwargs["past_covariates"]
+            )
+            encoder_kwargs.pop("input_ids")
+            encoder_kwargs["inputs_embeds"] = inputs_embeds
+            model_kwargs["encoder_outputs"]: ModelOutput = encoder(**encoder_kwargs)  # type: ignore
+
+        return model_kwargs
+
+
+class ChronosXPipeline(ChronosPipeline):
+    def __init__(
+        self,
+        prediction_length: int,
+        num_covariates: int,
+        covariate_injection: str = "IIB+OIB",
+        device_map: str = "cuda",  # use "cpu" for CPU inference
+        hidden_dim: int = 256,
+        num_layers: int = 1,
+        pretrained_model_name_or_path="amazon/chronos-t5-small",
+        layers_to_unfreeze: str = "injection_block",
+    ):
+
+        pretrained_kwargs = {"device_map": device_map}
+
+        pipeline = BaseChronosPipeline.from_pretrained(
+            pretrained_model_name_or_path=pretrained_model_name_or_path,
+            **pretrained_kwargs,
+        )
+
+        chronos_config = AutoConfig.from_pretrained(
+            pretrained_model_name_or_path,
+            **pretrained_kwargs,
+        )
+
+        self.prediction_length = prediction_length
+        self.num_covariates = num_covariates
+        self.covariate_injection = covariate_injection
+        self.hidden_dim = hidden_dim
+        self.num_layers = num_layers
+        self.pretrained_model_name_or_path = pretrained_model_name_or_path
+        self.pretrained_model = pipeline.model
+        self.pretrained_model_tokenizer = pipeline.tokenizer
+        self.tokenizer = self.create_tokenizer()
+        self.vocab_size = chronos_config.vocab_size
+        self.model_dim = chronos_config.d_model
+        self.layers_to_unfreeze = layers_to_unfreeze
+        self.chronosx = ChronosX.set_state(
+            num_covariates=self.num_covariates,
+            covariate_injection=self.covariate_injection,
+            hidden_dim=self.hidden_dim,
+            num_layers=self.num_layers,
+            vocab_size=self.vocab_size,
+            model_dim=self.model_dim,
+        ).from_pretrained(
+            self.pretrained_model_name_or_path,
+            **pretrained_kwargs,
+        )
+
+    def create_tokenizer(self):
+        inputs = self.pretrained_model_tokenizer.config.__dict__
+        inputs.update({"prediction_length": self.prediction_length})
+        chronos_config = ChronosConfig(**inputs)
+        return chronos_config.create_tokenizer()
+
+    def prepare_model_for_finetuning(self):
+
+        self.chronosx.initialize_blocks()
+
+        self.chronosx.config.pad_token_id = (
+            self.chronosx.generation_config.pad_token_id
+        ) = 0
+
+        self.chronosx.config.eos_token_id = (
+            self.chronosx.generation_config.eos_token_id
+        ) = 1
+
+        if self.layers_to_unfreeze is not None:
+            self.chronosx.freeze(layer_name="all")
+            self.chronosx.unfreeze(layer_name=self.layers_to_unfreeze)
+
+    def load_pretrained_zero_shot_model(self, device_map: str = "cuda"):
+        return ChronosX.set_state(
+            covariate_injection=None,
+            num_covariates=self.num_covariates,
+            hidden_dim=self.hidden_dim,
+            num_layers=self.num_layers,
+            vocab_size=self.vocab_size,
+            model_dim=self.model_dim,
+        ).from_pretrained(
+            self.pretrained_model_name_or_path,
+            **{"device_map": device_map},
+        )
+
+    def evaluate_model_on_validation_set(
+        self,
+        quantized_val_dataset,
+        output_dir=Path(__file__).parent / "output" / "group0" / "finetune",
+        dataloader_num_workers=1,
+        tf32=False,
+        torch_compile=0,
+        per_device_eval_batch_size=8,
+        eval_accumulation_steps=4,
+        device_map: str = "cuda",
+        covariate_injection: str = None,
+    ):
+
+        output_dir.mkdir(exist_ok=True, parents=True)
+
+        training_args = TrainingArguments(
+            output_dir=output_dir,
+            dataloader_num_workers=dataloader_num_workers,
+            tf32=tf32,  # remove this if not using Ampere GPUs (e.g., A100)
+            torch_compile=torch_compile,
+            per_device_eval_batch_size=per_device_eval_batch_size,
+            eval_accumulation_steps=eval_accumulation_steps,
+        )
+
+        pretrained_model_zeroshot = ChronosX.set_state(
+            covariate_injection=covariate_injection,
+            num_covariates=self.num_covariates,
+            hidden_dim=self.hidden_dim,
+            num_layers=self.num_layers,
+            vocab_size=self.vocab_size,
+            model_dim=self.model_dim,
+        ).from_pretrained(
+            self.pretrained_model_name_or_path,
+            **{"device_map": device_map},
+        )
+
+        # Create Trainer instance
+        self.pretrained_model.eval()
+        trainer = Trainer(
+            model=pretrained_model_zeroshot,
+            args=training_args,
+            eval_dataset=quantized_val_dataset,
+            compute_metrics=compute_metrics,
+        )
+
+        valiation_loss = trainer.evaluate(quantized_val_dataset)
+        self.pretrained_model.train()
+
+        return valiation_loss["eval_loss"]
+
+    def train(
+        self,
+        output_dir=Path(__file__).parent / "output" / "group0" / "finetune",
+        per_device_train_batch_size=32,
+        learning_rate=0.01,
+        lr_scheduler_type="linear",
+        warmup_ratio=0.0,
+        optim="adamw_torch_fused",
+        log_steps=20,
+        save_steps=100,
+        max_steps=5000,
+        gradient_accumulation_steps=2,
+        dataloader_num_workers=1,
+        tf32=False,
+        torch_compile=0,
+        eval_steps=100,
+        per_device_eval_batch_size=8,
+        eval_accumulation_steps=4,
+        load_best_model_at_end=True,
+        save_total_limit=5,
+        quantized_train_dataset=None,
+        quantized_val_dataset=None,
+        seed=None,
+    ):
+
+        output_dir.mkdir(exist_ok=True, parents=True)
+
+        training_args = TrainingArguments(
+            output_dir=output_dir,
+            per_device_train_batch_size=per_device_train_batch_size,
+            learning_rate=learning_rate,
+            lr_scheduler_type=lr_scheduler_type,
+            warmup_ratio=warmup_ratio,
+            optim=optim,
+            logging_dir=str(output_dir / "logs"),
+            logging_strategy="steps",
+            logging_steps=log_steps,
+            save_strategy="steps",
+            save_steps=save_steps,
+            report_to="tensorboard",
+            max_steps=max_steps,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            dataloader_num_workers=dataloader_num_workers,
+            tf32=tf32,  # remove this if not using Ampere GPUs (e.g., A100)
+            torch_compile=torch_compile,
+            ddp_find_unused_parameters=False,
+            metric_for_best_model="val_loss",
+            evaluation_strategy="steps",
+            eval_steps=eval_steps,
+            per_device_eval_batch_size=per_device_eval_batch_size,
+            eval_accumulation_steps=eval_accumulation_steps,
+            load_best_model_at_end=load_best_model_at_end,
+            logging_first_step=True,
+            greater_is_better=False,
+            save_total_limit=save_total_limit,
+            # seed=seed,
+        )
+
+        # Create Trainer instance
+        if quantized_val_dataset is None:
+            training_args.eval_strategy = "no"
+
+        trainer = Trainer(
+            model=self.chronosx,
+            args=training_args,
+            train_dataset=quantized_train_dataset,
+            eval_dataset=(
+                {"val": quantized_val_dataset} if quantized_val_dataset else None
+            ),
+            compute_metrics=compute_metrics,
+        )
+
+        # prepare training for finetuning
+        trainer._signature_columns = [
+            "labels",
+            "attention_mask",
+            "input_ids",
+            "past_covariates",
+            "future_covariates",
+            "decoder_input_ids",
+        ]
+
+        log_on_main("Training", logger)
+
+        self.prepare_model_for_finetuning()
+
+        parameter_count = count_parameters(self.chronosx)
+        print(f"parameter_count: {parameter_count}")
+
+        print("Model parameters:")
+        for name, params in self.chronosx.named_parameters():
+            print(name, params.requires_grad)
+
+        trainer.train()
+
+        if quantized_val_dataset:
+            val_loss_finetuned_model = trainer.evaluate(quantized_val_dataset)
+        else:
+            val_loss_finetuned_model = {"eval_loss": np.nan}
+
+        pretrained_model_path = output_dir / "final-checkpoint"
+        self.chronosx.save_pretrained(pretrained_model_path)
+
+        return val_loss_finetuned_model["eval_loss"]
+
+    def predict(
+        self, context: list[torch.Tensor], covariates: list[dict], num_samples: int = 20
+    ):
+        context_tensor = self._prepare_and_validate_context(context=context)
+        token_ids, attention_mask, scale = self.tokenizer.context_input_transform(
+            context_tensor
+        )
+
+        prepared_covariates = [prepare_covariates(entry) for entry in covariates]
+        future_covariates = torch.tensor(
+            [entry["future_covariates"] for entry in prepared_covariates]
+        )
+        past_covariates = torch.tensor(
+            [entry["past_covariates"] for entry in prepared_covariates]
+        )
+
+        preds = self.chronosx.generate(
+            input_ids=token_ids.to(self.chronosx.device),
+            attention_mask=attention_mask.to(self.chronosx.device),
+            generation_config=GenerationConfig(
+                min_new_tokens=self.prediction_length,
+                max_new_tokens=self.prediction_length,
+                do_sample=True,
+                num_return_sequences=num_samples,
+                eos_token_id=self.chronosx.config.eos_token_id,
+                pad_token_id=self.chronosx.config.pad_token_id,
+            ),
+            future_covariates=future_covariates.to(self.chronosx.device),
+            past_covariates=past_covariates.to(self.chronosx.device),
+        )
+
+        preds = preds[..., 1:]  # remove the decoder start token
+        preds = preds.reshape(token_ids.size(0), num_samples, -1)
+        preds = self.tokenizer.output_transform(preds.to("cpu"), scale.to("cpu"))
+
+        return preds.to(dtype=torch.float32, device="cpu")
+
+    def finetune(
+        self,
+        output_dir,
+        quantized_train_dataset,
+        lr=0.01,
+        quantized_val_dataset=None,
+        skip_pretrained_validation=False,
+        max_steps=20,
+        seed=None,
+    ):
+
+        if not skip_pretrained_validation:
+            assert quantized_val_dataset is not None
+            val_loss_pretrained_model = self.evaluate_model_on_validation_set(
+                covariate_injection=None,
+                quantized_val_dataset=quantized_val_dataset,
+                output_dir=output_dir,
+            )
+        else:
+            print(f"val_loss_pretrained_model is set up to nan since skip_pretrained_validation is given as False")
+            val_loss_pretrained_model = np.nan
+
+        eval_loss_finetuned_model = self.train(
+            learning_rate=lr,
+            quantized_train_dataset=quantized_train_dataset,
+            quantized_val_dataset=quantized_val_dataset,
+            max_steps=max_steps,
+            output_dir=output_dir,
+            seed=seed,
+        )
+
+        if quantized_val_dataset is None:
+            print(f"No Validation set is provided")
+
+        print(
+            f"lr:{lr} - val_loss_pretrained_model: {val_loss_pretrained_model:.4f} - eval_loss_finetuned_model:{eval_loss_finetuned_model:.4f}"
+        )
+
+        if val_loss_pretrained_model < eval_loss_finetuned_model:
+            model = self.load_pretrained_zero_shot_model()
+            val_loss = val_loss_pretrained_model
+        else:
+            model = self.chronosx
+            val_loss = eval_loss_finetuned_model
+
+        save_model_path = output_dir / "final-checkpoint"
+        model.save_pretrained(save_model_path)
+
+        return val_loss, save_model_path
+
+
+    def generate_forecasts(
+        self,
+        test_data_input,
+        batch_size=32,
+        context_length=512,
+    ):
+
+        transformation = create_transformation(include_covariates=True)
+        transformed_dataset = transformation(iter(test_data_input), is_train=False)
+        instance_splitter = InstanceSplitter(
+            target_field="target",
+            is_pad_field="is_pad",
+            start_field="start",
+            forecast_start_field="forecast_start",
+            instance_sampler=TestSplitSampler(),
+            past_length=context_length,
+            future_length=self.prediction_length,
+            time_series_fields=[
+                "observed_values",
+                "feat_dynamic_real",
+            ],
+            dummy_value=np.nan,
+        )
+        iterables = [instance_splitter.apply(transformed_dataset, is_train=False)]
+        iterators = list(map(iter, iterables))
+        test_data_input_transformed = itertools.chain(*iterators)
+
+        forecast_outputs = []
+        for batch in tqdm(batcher(test_data_input_transformed, batch_size=batch_size)):
+            context = [torch.tensor(entry["past_target"]) for entry in batch]
+            covariates = [
+                {
+                    "future_feat_dynamic_real": entry.get("future_feat_dynamic_real", None),
+                    "past_feat_dynamic_real": entry.get("past_feat_dynamic_real", None),
+                }
+                for entry in batch
+            ]
+            forecast_outputs.append(self.predict(context, covariates).numpy())
+
+        forecast_outputs = np.concatenate(forecast_outputs)
+
+        # Convert forecast samples into gluonts Forecast objects
+        forecasts = []
+        for item, ts in zip(forecast_outputs, test_data_input):
+            forecast_start_date = ts["start"] + len(ts["target"])
+            forecasts.append(SampleForecast(samples=item, start_date=forecast_start_date))
+
+        return forecasts
--- a/src/chronosx/injection_blocks/init.py
+++ b/src/chronosx/injection_blocks/init.py
--- a/src/chronosx/injection_blocks/basic_modules.py
+++ b/src/chronosx/injection_blocks/basic_modules.py
@ -0,0 +1,52 @@
+from torch import nn
+import math
+
+
+class FeedForwardNN(nn.Module):
+    def __init__(self, input_size, hidden_sizes, output_size):
+        super(FeedForwardNN, self).__init__()
+        self.input_size = input_size
+        self.hidden_sizes = hidden_sizes
+        self.output_size = output_size
+
+        layers = []
+        prev_size = input_size
+        for size in hidden_sizes:
+            layers.append(nn.Linear(prev_size, size))
+            layers.append(nn.ReLU())
+            prev_size = size
+
+        layers.append(nn.Linear(prev_size, output_size))
+
+        self.network = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.network(x)
+
+
+class InjectionBlock(nn.Module):
+    name = "generic"
+
+    def __init__(
+        self,
+    ):
+        super(InjectionBlock, self).__init__()
+        self.counter = 0
+        self.generating = False
+
+    def restart_generator_counter(self):
+        self.counter = 0
+        self.generating = True
+
+    def initialize_modules(self, modules=None):
+        if modules is None:
+            modules = self.modules()
+        for module in modules:
+            if isinstance(module, nn.Linear):
+                for name, param in module.named_parameters():
+                    if "weight" in name:
+                        nn.init.kaiming_uniform_(param, a=math.sqrt(5))
+                    elif "bias" in name:
+                        fan_in, _ = nn.init._calculate_fan_in_and_fan_out(module.weight)
+                        bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+                        nn.init.uniform_(param, -bound, bound)
--- a/src/chronosx/injection_blocks/block_mapping.py
+++ b/src/chronosx/injection_blocks/block_mapping.py
@ -0,0 +1,8 @@
+from chronosx.injection_blocks.input_injection_block import InputInjectionBlock
+from chronosx.injection_blocks.output_injection_block import OutputInjectionBlock
+
+injection_blocks_map = {
+    "IIB": (InputInjectionBlock, None),
+    "OIB": (None, OutputInjectionBlock),
+    "IIB+OIB": (InputInjectionBlock, OutputInjectionBlock),
+}
--- a/src/chronosx/injection_blocks/input_injection_block.py
+++ b/src/chronosx/injection_blocks/input_injection_block.py
@ -0,0 +1,42 @@
+from torch import nn
+import torch
+
+from chronosx.injection_blocks.basic_modules import InjectionBlock, FeedForwardNN
+
+
+class InputInjectionBlock(InjectionBlock):
+    name = "input_injection_block"
+
+    def __init__(
+        self,
+        hidden_dim: int = 256,
+        model_dim: int = 512,
+        num_covariates: int = 0,
+        num_layers: int = 1,
+    ):
+        super(InputInjectionBlock, self).__init__()
+
+        self.hidden_dim = hidden_dim
+        self.model_dim = model_dim
+        self.num_covariates = num_covariates
+        self.num_layers = num_layers
+
+        self.cov_in = nn.Linear(self.num_covariates, self.hidden_dim)
+        self.concat_dim = self.hidden_dim * 2
+
+        self.concat_layer = FeedForwardNN(
+            self.concat_dim, [self.hidden_dim] * self.num_layers, self.model_dim
+        )
+        self.emb_in = nn.Linear(self.model_dim, self.hidden_dim)
+
+    def forward(self, input_embeds, past_covariates, is_decoder=False):
+
+        x = self.emb_in(input_embeds)
+        if self.generating and is_decoder:
+            past_covariates = past_covariates[:, self.counter, :].unsqueeze(1)
+            self.counter += 1
+        x_cov = self.cov_in(past_covariates)
+        x = torch.cat([x, x_cov], axis=-1)
+        x = nn.ReLU()(x)
+
+        return input_embeds + self.concat_layer(x)
--- a/src/chronosx/injection_blocks/output_injection_block.py
+++ b/src/chronosx/injection_blocks/output_injection_block.py
@ -0,0 +1,64 @@
+from torch import nn
+from torch.nn import CrossEntropyLoss
+import torch
+
+from chronosx.injection_blocks.basic_modules import InjectionBlock, FeedForwardNN
+
+
+class OutputInjectionBlock(InjectionBlock):
+    name = "output_injection_block"
+
+    def __init__(
+        self,
+        hidden_dim: int = 256,
+        model_dim: int = 512,
+        num_covariates: int = 0,
+        num_layers: int = 1,
+        vocab_size: int = 4096,
+    ):
+        super(OutputInjectionBlock, self).__init__()
+
+        self.hidden_dim = hidden_dim
+        self.model_dim = model_dim
+        self.num_covariates = num_covariates
+        self.num_layers = num_layers
+        self.vocab_size = vocab_size
+
+        # self.num_layers = num_layers
+        self.concat_dim = 2 * self.hidden_dim
+        self.cov_out = nn.Linear(self.num_covariates, self.hidden_dim)
+        self.concat_layer = FeedForwardNN(
+            self.concat_dim, [self.hidden_dim] * self.num_layers, self.vocab_size
+        )
+        self.hidden_state_out = nn.Linear(self.model_dim, self.hidden_dim)
+        self.loss_fct = CrossEntropyLoss(ignore_index=-100)
+
+    def compute_loss(self, logits, labels):
+        labels = labels.to(logits.device)
+        loss = self.loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+        return loss
+
+    def forward(self, future_covariates, labels, logits, last_hidden_state):
+        x = self.hidden_state_out(last_hidden_state)
+
+        if self.generating:
+            x_cov = future_covariates[:, self.counter].unsqueeze(1)
+        else:
+            x_cov = future_covariates
+
+        if x_cov.flatten().isnan().sum() > 0:
+            print(1)
+
+        x_cov = self.cov_out(x_cov)  # -> here is the problem
+        x = torch.concatenate([x, x_cov], axis=-1)
+
+        x = nn.ReLU()(x)
+        logits = self.concat_layer(x) + logits
+
+        if self.training:
+            loss = self.compute_loss(logits, labels)
+        else:
+            loss = None
+            self.counter += 1
+
+        return logits, loss
--- a/src/chronosx/utils/init.py
+++ b/src/chronosx/utils/init.py
@ -0,0 +1,9 @@
+from .chronos_dataset import ChronosDataset
+from .hf_data_loader import load_and_split_dataset
+from .prepare_covariates import prepare_covariates
+
+__all__ = [
+    "ChronosDataset",
+    "load_and_split_dataset",
+    "prepare_covariates"
+]
--- a/src/chronosx/utils/chronos_dataset.py
+++ b/src/chronosx/utils/chronos_dataset.py
@ -0,0 +1,412 @@
+import itertools
+import numpy as np
+import torch
+import typer
+
+from chronos import ChronosTokenizer
+from chronosx.utils.transform import DropValues
+from gluonts.dataset.field_names import FieldName
+from gluonts.itertools import Cyclic
+from torch.utils.data import IterableDataset, get_worker_info
+from typing import List, Iterator, Optional
+from gluonts.transform import (
+    FilterTransformation,
+    TestSplitSampler,
+    ValidationSplitSampler,
+    InstanceSplitter,
+    ExpectedNumInstanceSampler,
+    MissingValueImputation,
+    Transformation,
+    AddObservedValuesIndicator,
+    SelectFields,
+    VstackFeatures,
+    Chain,
+    LeavesMissingValues,
+)
+
+
+
+
+app = typer.Typer(pretty_exceptions_enable=False)
+
+
+class PseudoShuffledIterableDataset(IterableDataset):
+    """
+    Shuffle entries from an iterable by temporarily accumulating them
+    in an intermediate buffer.
+
+    Parameters
+    ----------
+    base_dataset
+        The original iterable object, representing the dataset.
+    shuffle_buffer_length
+        Size of the buffer use to shuffle entries from the base dataset.
+    """
+
+    def __init__(
+        self, base_dataset, shuffle_buffer_length: int = 100, random_seed: int = None
+    ) -> None:
+        super().__init__()
+        self.base_dataset = base_dataset
+        self.shuffle_buffer_length = shuffle_buffer_length
+        self.generator = torch.Generator()
+        if random_seed:
+            self.generator.manual_seed(random_seed)
+
+    def __iter__(self):
+        shuffle_buffer = []
+
+        for element in self.base_dataset:
+            shuffle_buffer.append(element)
+            if len(shuffle_buffer) >= self.shuffle_buffer_length:
+                idx = torch.randint(
+                    len(shuffle_buffer), size=(), generator=self.generator
+                )
+                yield shuffle_buffer.pop(idx)
+
+        while shuffle_buffer:
+            idx = torch.randint(len(shuffle_buffer), size=(), generator=self.generator)
+            yield shuffle_buffer.pop(idx)
+
+
+class ShuffleMixin:
+    """
+    Mix-in class that datasets can inherit from to get
+    shuffling functionality.
+    """
+
+    def shuffle(self, shuffle_buffer_length: int = 100, random_seed: int = None):
+        return PseudoShuffledIterableDataset(self, shuffle_buffer_length, random_seed)
+
+
+class ChronosDataset(IterableDataset, ShuffleMixin):
+    """
+    Dataset wrapper, using a ``ChronosTokenizer`` to turn data from a time series
+    into a HuggingFace-compatible set of ``input_ids``, ``attention_mask`` and
+    ``labels``.
+
+    Entries from the original datasets are assumed to have a ``"start"`` attribute
+    (of type ``pd.Period``), and a ``"target"`` attribute (of type ``np.ndarray``).
+
+    Parameters
+    ----------
+    datasets
+        Datasets containing the original time series data.
+    probabilities
+        In training mode, data will be sampled from each of the original datasets
+        with these probabilities.
+    tokenizer
+        Tokenizer to be used to turn sequences of real numbers into token IDs.
+    context_length
+        Samples context will be limited to this length.
+    prediction_length
+        Samples labels will be limited to this length.
+    drop_prob
+        In training mode, observations from a sample will be turned into ``np.nan``,
+        i.e. turned into missing values, with this probability.
+    min_past
+        Data samples will be considered only if there's at least ``min_past``-many
+        historical observations.
+    mode
+        One of ``"training"``, ``"validation"``, or ``"test"``.
+    np_dtype
+        Numpy float data type.
+    """
+
+    def __init__(
+        self,
+        datasets: list,
+        probabilities: List[float],
+        tokenizer: ChronosTokenizer,
+        context_length: int = 512,
+        prediction_length: int = 64,
+        drop_prob: float = 0.2,
+        min_past: Optional[int] = None,
+        model_type: str = "seq2seq",
+        imputation_method: Optional[MissingValueImputation] = None,
+        mode: str = "training",
+        np_dtype=np.float32,
+        include_covariates: bool = True,
+    ) -> None:
+        super().__init__()
+
+        assert len(probabilities) == len(datasets)
+        assert mode in ("training", "validation", "test")
+        assert model_type in ("seq2seq", "causal")
+
+        self.datasets = datasets
+        self.probabilities = probabilities
+        self.tokenizer = tokenizer
+        self.context_length = context_length
+        self.prediction_length = prediction_length
+        self.drop_prob = drop_prob if model_type == "seq2seq" else 0.0
+        self.min_past = min_past or prediction_length
+        self.model_type = model_type
+        self.imputation_method = imputation_method or LeavesMissingValues()
+        self.mode = mode
+        self.np_dtype = np_dtype
+        self.include_covariates = include_covariates
+        self.time_series_fields = [
+            FieldName.OBSERVED_VALUES,
+        ]
+
+        if self.include_covariates:
+            self.time_series_fields.append(FieldName.FEAT_DYNAMIC_REAL)
+
+    def preprocess_entry(self, entry: dict, mode: str) -> dict:
+        entry = {f: entry[f] for f in ["start", "target"]}
+        entry["target"] = np.asarray(entry["target"], dtype=self.np_dtype)
+        assert entry["target"].ndim == 1, f"got {entry['target'].ndim=}, expected 1"
+
+        if self.model_type == "causal":
+            # Causal models do not play nice with missing values, so it is
+            # recommended to use an imputation method, e.g., LastValueImputation
+            entry["target"] = self.imputation_method(entry["target"])
+
+        if mode == "training" and self.drop_prob > 0:
+            target = entry["target"].copy()
+            drop_p = np.random.uniform(low=0.0, high=self.drop_prob)
+            mask = np.random.choice(
+                [True, False], size=len(target), p=[drop_p, 1 - drop_p]
+            )
+            target[mask] = np.nan
+            entry["target"] = target
+
+        return entry
+
+    # to adapt for covariates
+    def create_transformation(self) -> Transformation:
+
+        fields = [
+            FieldName.START,
+            FieldName.TARGET,
+        ]
+        if self.include_covariates:
+            fields.append(FieldName.FEAT_DYNAMIC_REAL)
+
+        return (
+            SelectFields(
+                fields,
+                allow_missing=True,
+            )
+            + DropValues(drop_prob=0.2, target_field=FieldName.TARGET)
+            + AddObservedValuesIndicator(
+                target_field=FieldName.TARGET,
+                output_field=FieldName.OBSERVED_VALUES,
+            )
+        )
+
+    def _create_instance_splitter(self, mode: str):
+        assert mode in ["training", "test", "validation"]
+
+        instance_sampler = {
+            "training": ExpectedNumInstanceSampler(
+                num_instances=1.0,
+                min_instances=1,
+                min_past=self.min_past,
+                min_future=self.prediction_length,
+            ),
+            "test": TestSplitSampler(),
+            "validation": ValidationSplitSampler(min_future=self.prediction_length),
+        }[mode]
+
+        return InstanceSplitter(
+            target_field="target",
+            is_pad_field="is_pad",
+            start_field="start",
+            forecast_start_field="forecast_start",
+            instance_sampler=instance_sampler,
+            past_length=self.context_length,
+            future_length=self.prediction_length,
+            time_series_fields=self.time_series_fields,  # to adapt for covariates
+            dummy_value=np.nan,
+            # dummy_value=0.0,
+        )
+
+    def create_training_data(self, data):
+        data = Cyclic(data)
+        split_transform = self._create_instance_splitter(
+            "training"
+        ) + FilterTransformation(
+            condition=lambda entry: (~np.isnan(entry["past_target"])).sum() > 0
+        )
+        data = split_transform.apply(data, is_train=True)
+        return data
+
+    def create_test_data(self, data):
+        data = self._create_instance_splitter("test").apply(data, is_train=False)
+        return data
+
+    def create_validation_data(self, data):
+        data = self._create_instance_splitter("validation").apply(data, is_train=False)
+        return data
+
+    def to_hf_format(self, entry: dict) -> dict:
+        past_target = torch.tensor(entry["past_target"]).unsqueeze(0)
+        input_ids, attention_mask, scale = self.tokenizer.context_input_transform(
+            past_target
+        )
+
+        future_target = torch.tensor(entry["future_target"]).unsqueeze(0)
+
+        if self.mode == "test":
+            labels = future_target.detach().clone()
+        else:
+            labels, labels_mask = self.tokenizer.label_input_transform(
+                future_target, scale
+            )
+            labels[labels_mask == 0] = -100
+
+        # normalize covariates by mean of absolute values
+        past_covariates = entry.get("past_feat_dynamic_real", None)
+        if past_covariates is not None:
+            #  missing value imputation for covariates
+            past_covariates[np.isnan(past_covariates)] = 0
+
+            covariates_scale = np.mean(np.abs(past_covariates), axis=0)
+            covariates_scale[covariates_scale < 1] = 1.0
+            past_covariates = past_covariates / covariates_scale
+
+        future_covariates = entry.get("future_feat_dynamic_real", None)
+        if future_covariates is not None:
+            #  missing value imputation for covariates
+            future_covariates[np.isnan(future_covariates)] = 0
+
+            future_covariates = future_covariates / covariates_scale
+
+        # shift covariates by one
+        if past_covariates is not None:
+            num_covariates = past_covariates.shape[-1]
+            past_covariates = np.concatenate(
+                [past_covariates, np.array([[0] * num_covariates])], axis=0
+            )
+
+        if future_covariates is not None:
+            num_covariates = future_covariates.shape[-1]
+            future_covariates = np.concatenate(
+                [future_covariates, np.array([[0] * num_covariates])], axis=0
+            )
+
+        if self.model_type == "causal":
+            # The InstanceSplitter pads time series on the left to be equal to the
+            # context_length. However, certain models (e.g., GPT2) with absolute
+            # position embeddings should not be trained with left padding.
+            # The following piece of code moves padding from left to right.
+
+            assert input_ids.shape[-1] == entry["past_is_pad"].shape[0]
+
+            # Find the index where padding starts
+            pad_start_idx = np.searchsorted(1 - entry["past_is_pad"], 1)
+            padded_input_ids, obs_input_ids = torch.tensor_split(
+                input_ids, [pad_start_idx], dim=-1
+            )
+            padded_attention_mask, obs_attention_mask = torch.tensor_split(
+                attention_mask, [pad_start_idx], dim=-1
+            )
+
+            # Move padding to the right
+            input_ids = torch.cat(
+                [
+                    obs_input_ids,
+                    labels,
+                    padded_input_ids,
+                ],
+                axis=-1,
+            )
+            attention_mask = torch.cat(
+                [
+                    obs_attention_mask,
+                    labels_mask,
+                    padded_attention_mask,
+                ],
+                axis=-1,
+            )
+
+            # labels for causal models are same as the input_ids.
+            # Internally transformers shifts the labels by one during training.
+            labels = input_ids.clone()
+            input_ids[~attention_mask] = self.tokenizer.config.pad_token_id
+            labels[~attention_mask] = -100
+        return {
+            "input_ids": input_ids.squeeze(0),
+            "attention_mask": attention_mask.squeeze(0),
+            "labels": labels.squeeze(0),
+            "past_covariates": past_covariates.astype(np.float32),
+            "future_covariates": future_covariates.astype(np.float32),
+            "scale": scale,
+        }
+
+    def __iter__(self) -> Iterator:
+        transformation = self.create_transformation()
+
+        if self.include_covariates:
+            field_time = (
+                FieldName.FEAT_TIME
+                if not self.include_covariates
+                else FieldName.FEAT_DYNAMIC_REAL
+            )
+            transformation = Chain(
+                [
+                    transformation,
+                    VstackFeatures(
+                        output_field=FieldName.FEAT_DYNAMIC_REAL,
+                        input_fields=[field_time],
+                    ),
+                ]
+            )
+
+        if self.mode == "training":
+            transformed_datasets = [
+                transformation.apply(dataset, is_train=True)
+                for dataset in self.datasets
+            ]
+            iterables = [
+                self.create_training_data(transformed_dataset)
+                for transformed_dataset in transformed_datasets
+            ]
+        elif self.mode == "validation":
+            transformed_datasets = [
+                transformation.apply(dataset, is_train=False)
+                for dataset in self.datasets
+            ]
+            iterables = [
+                self.create_validation_data(transformed_dataset)
+                for transformed_dataset in transformed_datasets
+            ]
+        elif self.mode == "test":
+            transformed_datasets = [
+                transformation.apply(dataset, is_train=False)
+                for dataset in self.datasets
+            ]
+            iterables = [
+                self.create_test_data(transformed_dataset)
+                for transformed_dataset in transformed_datasets
+            ]
+
+        worker_info = get_worker_info()
+        if worker_info is None:
+            probs = list(self.probabilities)
+        else:
+            worker_id = worker_info.id
+            num_workers = worker_info.num_workers
+            iterables = list(itertools.islice(iterables, worker_id, None, num_workers))
+            probs = list(
+                itertools.islice(self.probabilities, worker_id, None, num_workers)
+            )
+
+        probs = [prob / sum(probs) for prob in probs]
+
+        iterators = list(map(iter, iterables))
+        if self.mode == "training":
+            while True:
+                idx = np.random.choice(range(len(iterators)), p=probs)
+                try:
+                    yield self.to_hf_format(next(iterators[idx]))
+                except StopIteration:
+                    probs[idx] = 0
+                    if sum(probs) == 0:
+                        return
+                    probs = [prob / sum(probs) for prob in probs]
+        else:
+            for entry in itertools.chain(*iterators):
+                yield self.to_hf_format(entry)
--- a/src/chronosx/utils/epf.py
+++ b/src/chronosx/utils/epf.py
@ -0,0 +1,56 @@
+from pathlib import Path
+from urllib import request
+import datasets
+import pandas as pd
+import tempfile
+
+
+MAIN_URL = "https://raw.githubusercontent.com/Nixtla/transfer-learning-time-series/f055eec3266eda77e183b62e9f07eff0bc6c155b/datasets/electricity.csv"
+EX_URL = "https://raw.githubusercontent.com/Nixtla/transfer-learning-time-series/f055eec3266eda77e183b62e9f07eff0bc6c155b/datasets/exogenous-vars-electricity.csv"
+
+
+def load_data(id: str):
+    with tempfile.TemporaryDirectory() as dir_path:
+        data_path = Path(dir_path) / "data.csv"
+
+        request.urlretrieve(MAIN_URL, data_path)
+        main_data = pd.read_csv(data_path)
+        main_data = main_data[main_data.unique_id == id][["ds", "y"]]
+
+        request.urlretrieve(EX_URL, data_path)
+        ex_data = pd.read_csv(data_path)
+        ex_data = ex_data[ex_data.unique_id == id]
+
+    time_data = pd.date_range(
+        start=main_data.ds.values[0], end=main_data.ds.values[-1], freq="H"
+    )
+    data = pd.merge(main_data, ex_data, how="left", on="ds")
+
+    assert time_data.shape[0] == data.shape[0]
+
+    return data
+
+
+def generate(id: str):
+    data = load_data(id)
+    df = data[["ds", "y", "unique_id", "Exogenous1", "Exogenous2"]]
+    df = df.rename(
+        columns={
+            "ds": "timestamp",
+            "y": "target",
+            "unique_id": "id",
+            "Exogenous1": "exogenous1",
+            "Exogenous2": "exogenous2",
+        }
+    )
+
+    ds = datasets.Dataset.from_list([df.to_dict("list")])
+
+    return ds
+
+
+def get_epf_dataset(dataset_name: str):
+    id = dataset_name.split("_")[-2].upper()
+    ds = generate(id)
+
+    return ds
--- a/src/chronosx/utils/hf_data_loader.py
+++ b/src/chronosx/utils/hf_data_loader.py
@ -0,0 +1,92 @@
+import datasets
+import numpy as np
+import pandas as pd
+from gluonts.dataset.field_names import FieldName
+from gluonts.dataset.split import split
+from typing import List
+from .m5 import get_m5_dataset
+from .epf import get_epf_dataset
+
+
+def to_gluonts_univariate(
+    hf_dataset: datasets.Dataset,
+    series_fields: List[str],
+    covariates_fields: List[str] = None,
+):
+    if isinstance(series_fields, str):
+        series_fields = [series_fields]
+
+    # dataset_length = hf_dataset.info.splits["train"].num_examples * len(series_fields)
+    dataset_length = len(hf_dataset) * len(series_fields)
+
+    # Assumes that all time series in the dataset have the same frequency
+    dataset_freq = pd.DatetimeIndex(hf_dataset[0]["timestamp"]).to_period()[0].freqstr
+
+    gts_dataset = []
+    for hf_entry in hf_dataset:
+        for field in series_fields:
+            entry = {
+                "start": pd.Period(
+                    hf_entry["timestamp"][0],
+                    freq=dataset_freq,
+                ),
+                "target": hf_entry[field],
+            }
+
+            if covariates_fields:
+                covariates = np.array([hf_entry[field] for field in covariates_fields])
+                is_nan_covariates = np.isnan(covariates)
+                covariates[is_nan_covariates] = -1
+                covariates = np.vstack([covariates, is_nan_covariates]).astype(
+                    np.float32
+                )
+                entry.update({FieldName.FEAT_DYNAMIC_REAL: covariates})
+
+            gts_dataset.append(entry)
+
+    assert len(gts_dataset) == dataset_length
+
+    return gts_dataset
+
+
+def load_and_split_dataset(backtest_config: dict):
+    hf_repo = backtest_config.get("hf_repo", None)
+    filename = backtest_config.get("filename", None)
+    dataset_name = backtest_config["name"]
+    offset = backtest_config["offset"]
+    prediction_length = backtest_config["prediction_length"]
+    num_rolls = backtest_config["num_rolls"]
+    series_fields = backtest_config.get("series_fields")
+    covariates_fields = backtest_config.get("covariates_fields")
+
+    if dataset_name == "m5_with_covariates":
+        ds = get_m5_dataset()
+
+    elif dataset_name in [
+        "epf_electricity_be_paper",
+        "epf_electricity_de_paper",
+        "epf_electricity_fr_paper",
+        "epf_electricity_np_paper",
+        "epf_electricity_pjm_paper",
+    ]:
+        ds = get_epf_dataset(dataset_name)
+
+    elif 'synthetic_datasets' in dataset_name:
+        ds = datasets.load_from_disk(filename)
+
+    else:
+        ds = datasets.load_dataset(
+            hf_repo, dataset_name, split="train", trust_remote_code=True
+        )
+
+    ds.set_format("numpy")
+
+    gts_dataset = to_gluonts_univariate(
+        ds, covariates_fields=covariates_fields, series_fields=series_fields
+    )
+
+    # Split dataset for evaluation
+    train_dataset, test_template = split(gts_dataset, offset=offset)
+    test_data = test_template.generate_instances(prediction_length, windows=num_rolls)
+
+    return train_dataset, test_data
--- a/src/chronosx/utils/m5.py
+++ b/src/chronosx/utils/m5.py
@ -0,0 +1,55 @@
+from datasetsforecast.m5 import M5
+from joblib import Parallel, delayed
+from pathlib import Path
+from tqdm.auto import tqdm
+import datasets
+import pandas as pd
+import pyarrow as pa
+
+def _process_entry(id_, ts, static_df):
+    entry = {"id": id_}
+    entry.update(ts.drop(columns=["id"]).to_dict("list"))
+    entry.update(static_df.loc[id_].to_dict())
+    return entry
+
+
+def generate():
+    target, covariates, static = M5.load(".")
+
+    encoded_columns = pd.get_dummies(
+        covariates[["event_type_1", "event_type_2"]]
+    ).astype(float)
+
+    covariates = covariates[["snap_CA", "snap_TX", "snap_WI", "sell_price"]]
+    covariates = pd.concat([covariates, encoded_columns], axis=1)
+
+    df = pd.concat([target, covariates], axis=1)
+    df = df.rename(columns={"unique_id": "id", "ds": "timestamp", "y": "target"})
+    df["timestamp"] = df["timestamp"].astype("datetime64[ms]")
+
+    static = static.rename(columns={"unique_id": "id"}).set_index("id")
+
+    processed = Parallel(n_jobs=-1)(
+        delayed(_process_entry)(id_, ts, static) for id_, ts in tqdm(df.groupby("id"))
+    )
+
+    table = pa.Table.from_pylist(processed)
+
+    splits = datasets.SplitDict()
+    splits.add(datasets.SplitInfo('train', num_examples=len(processed)))
+    info = datasets.DatasetInfo(splits=splits)
+
+    dataset = datasets.Dataset(table, info=info, split='train')
+    dataset.save_to_disk(Path(__file__).parent / "m5_with_covariates")
+
+
+def get_m5_dataset(force_computation: bool = False):
+    if (Path(__file__).parent / "m5_with_covariates").is_dir():
+        if force_computation is True:
+            generate()
+    else:
+        generate()
+
+    return datasets.load_from_disk(Path(__file__).parent / "m5_with_covariates")
+
+
--- a/src/chronosx/utils/prepare_covariates.py
+++ b/src/chronosx/utils/prepare_covariates.py
@ -0,0 +1,38 @@
+import numpy as np
+
+
+def prepare_covariates(entry: dict) -> dict:
+
+    past_covariates = entry.get("past_feat_dynamic_real", None)
+    if past_covariates is not None:
+        #  missing value imputation for covariates
+        past_covariates[np.isnan(past_covariates)] = 0
+
+        # normalize covariates by mean of absolute values,
+        covariates_scale = np.mean(np.abs(past_covariates), axis=0)
+        covariates_scale[covariates_scale < 1] = 1.0
+        past_covariates = past_covariates / covariates_scale
+
+        # shift covariates by one
+        num_covariates = past_covariates.shape[-1]
+        past_covariates = np.concatenate(
+            [past_covariates, np.array([[0] * num_covariates])], axis=0
+        )
+
+    future_covariates = entry.get("future_feat_dynamic_real", None)
+    if future_covariates is not None:
+        #  missing value imputation for covariates
+        future_covariates[np.isnan(future_covariates)] = 0
+
+        # normalize covariates by mean of absolute values,
+        future_covariates = future_covariates / covariates_scale
+
+        num_covariates = future_covariates.shape[-1]
+        future_covariates = np.concatenate(
+            [future_covariates, np.array([[0] * num_covariates])], axis=0
+        )
+
+    return {
+        "past_covariates": past_covariates.astype(np.float32),
+        "future_covariates": future_covariates.astype(np.float32),
+    }
--- a/src/chronosx/utils/transform.py
+++ b/src/chronosx/utils/transform.py
@ -0,0 +1,83 @@
+import numpy as np
+from gluonts.dataset.common import DataEntry
+from gluonts.dataset.field_names import FieldName
+from gluonts.transform import (
+    AddObservedValuesIndicator,
+    Chain,
+    MapTransformation,
+    SelectFields,
+    Transformation,
+    VstackFeatures,
+)
+
+
+class DropValues(MapTransformation):
+    def __init__(
+        self,
+        drop_prob: float = 0.0,
+        drop_mode: str = "upto",
+        target_field: str = "target",
+    ):
+        assert drop_mode in ["upto", "exact"]
+        assert 0.0 <= drop_prob <= 1.0
+
+        self.drop_prob = drop_prob
+        self.drop_mode = drop_mode
+        self.target_field = target_field
+
+    def map_transform(self, data: DataEntry, is_train: bool) -> DataEntry:
+        if not is_train:
+            return data
+
+        target = data[self.target_field].copy().astype(float)
+
+        if self.drop_mode == "exact":
+            drop_p = self.drop_prob
+        else:
+            drop_p = np.random.uniform(low=0.0, high=self.drop_prob)
+
+        mask = np.random.choice([True, False], size=len(target), p=[drop_p, 1 - drop_p])
+        target[mask] = np.nan
+        data[self.target_field] = target
+
+        return data
+
+
+def create_transformation(include_covariates: bool) -> Transformation:
+
+    fields = [
+        FieldName.START,
+        FieldName.TARGET,
+    ]
+    if include_covariates:
+        fields.append(FieldName.FEAT_DYNAMIC_REAL)
+
+    transformation = (
+        SelectFields(
+            fields,
+            allow_missing=True,
+        )
+        + DropValues(drop_prob=0.2, target_field=FieldName.TARGET)
+        + AddObservedValuesIndicator(
+            target_field=FieldName.TARGET,
+            output_field=FieldName.OBSERVED_VALUES,
+        )
+    )
+
+    if include_covariates:
+        field_time = (
+            FieldName.FEAT_TIME
+            if not include_covariates
+            else FieldName.FEAT_DYNAMIC_REAL
+        )
+        assert field_time == "feat_dynamic_real"
+        transformation = Chain(
+            [
+                transformation,
+                VstackFeatures(
+                    output_field=FieldName.FEAT_DYNAMIC_REAL, input_fields=[field_time]
+                ),
+            ]
+        )
+
+    return transformation
--- a/src/chronosx/utils/utils.py
+++ b/src/chronosx/utils/utils.py
@ -0,0 +1,54 @@
+from torch import nn
+import logging
+import numpy as np
+import os
+import torch
+import torch.distributed as dist
+
+
+def count_parameters(model):
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+
+def compute_metrics(pred, num_classes=4096):
+    loss = nn.CrossEntropyLoss()
+    logits = torch.FloatTensor(pred.predictions[0])
+    labels = torch.LongTensor(pred.label_ids)
+
+    loss_value = loss(logits.reshape(-1, num_classes), labels.reshape(-1))
+    return {"loss": loss_value}
+
+
+def is_main_process():
+    if not dist.is_torchelastic_launched():
+        return True
+    return int(os.environ["RANK"]) == 0
+
+
+def log_on_main(msg: str, logger: logging.Logger, log_level: int = logging.INFO):
+    if is_main_process():
+        logger.log(log_level, msg)
+
+
+def has_enough_observations(
+    entry: dict, min_length: int = 0, max_missing_prop: float = 1.0
+) -> bool:
+    """
+    Check if the given entry has enough observations in the ``"target"`` attribute.
+
+    Parameters
+    ----------
+    entry
+        The data entry (dictionary) to be tested.
+    min_length
+        The minimum length the ``"target"`` attribute must have.
+    max_missing_prop
+        The maximum proportion of missing data allowed in the ``"target"``
+        attribute.
+    """
+    if (
+        len(entry["target"]) >= min_length
+        and np.isnan(entry["target"]).mean() <= max_missing_prop
+    ):
+        return True
+    return False
--- a/test/init.py
+++ b/test/init.py
@ -1,2 +0,0 @@
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
--- a/test/dummy-chronos-bolt-model/config.json
+++ b/test/dummy-chronos-bolt-model/config.json
@ -1,50 +0,0 @@
-{
-  "architectures": [
-    "ChronosBoltModelForForecasting"
-  ],
-  "chronos_config": {
-    "context_length": 512,
-    "input_patch_size": 16,
-    "input_patch_stride": 16,
-    "prediction_length": 64,
-    "quantiles": [
-      0.1,
-      0.2,
-      0.3,
-      0.4,
-      0.5,
-      0.6,
-      0.7,
-      0.8,
-      0.9
-    ],
-    "use_reg_token": true
-  },
-  "chronos_pipeline_class": "ChronosBoltPipeline",
-  "classifier_dropout": 0.0,
-  "d_ff": 8,
-  "d_kv": 4,
-  "d_model": 8,
-  "decoder_start_token_id": 0,
-  "dense_act_fn": "relu",
-  "dropout_rate": 0.1,
-  "eos_token_id": 1,
-  "feed_forward_proj": "relu",
-  "initializer_factor": 0.05,
-  "is_encoder_decoder": true,
-  "is_gated_act": false,
-  "layer_norm_epsilon": 1e-06,
-  "model_type": "t5",
-  "n_positions": 512,
-  "num_decoder_layers": 4,
-  "num_heads": 4,
-  "num_layers": 4,
-  "pad_token_id": 0,
-  "reg_token_id": 1,
-  "relative_attention_max_distance": 128,
-  "relative_attention_num_buckets": 32,
-  "torch_dtype": "float32",
-  "transformers_version": "4.40.2",
-  "use_cache": true,
-  "vocab_size": 2
-}
--- a/test/dummy-chronos-bolt-model/model.safetensors
+++ b/test/dummy-chronos-bolt-model/model.safetensors
--- a/test/dummy-chronos-model/config.json
+++ b/test/dummy-chronos-model/config.json
@ -1,48 +0,0 @@
-{
-  "architectures": [
-    "T5ForConditionalGeneration"
-  ],
-  "d_ff": 32,
-  "d_kv": 16,
-  "d_model": 64,
-  "decoder_start_token_id": 0,
-  "dense_act_fn": "relu",
-  "dropout_rate": 0.1,
-  "eos_token_id": 1,
-  "feed_forward_proj": "relu",
-  "initializer_factor": 0.05,
-  "is_encoder_decoder": true,
-  "is_gated_act": false,
-  "layer_norm_epsilon": 1e-06,
-  "model_type": "t5",
-  "n_positions": 512,
-  "num_decoder_layers": 1,
-  "num_heads": 1,
-  "num_layers": 1,
-  "pad_token_id": 0,
-  "relative_attention_max_distance": 128,
-  "relative_attention_num_buckets": 32,
-  "torch_dtype": "bfloat16",
-  "transformers_version": "4.31.0",
-  "use_cache": true,
-  "vocab_size": 32,
-  "chronos_config": {
-    "tokenizer_class": "MeanScaleUniformBins",
-    "tokenizer_kwargs": {
-      "low_limit": -15.0,
-      "high_limit": 15.0
-    },
-    "n_tokens": 32,
-    "n_special_tokens": 2,
-    "pad_token_id": 0,
-    "eos_token_id": 1,
-    "use_eos_token": true,
-    "model_type": "seq2seq",
-    "context_length": 512,
-    "prediction_length": 64,
-    "num_samples": 20,
-    "temperature": 1.0,
-    "top_k": 50,
-    "top_p": 1.0
-  }  
-}
--- a/test/dummy-chronos-model/generation_config.json
+++ b/test/dummy-chronos-model/generation_config.json
@ -1,7 +0,0 @@
-{
-  "_from_model_config": true,
-  "decoder_start_token_id": 0,
-  "eos_token_id": 1,
-  "pad_token_id": 0,
-  "transformers_version": "4.31.0"
-}
--- a/test/dummy-chronos-model/pytorch_model.bin
+++ b/test/dummy-chronos-model/pytorch_model.bin
--- a/test/test_chronos.py
+++ b/test/test_chronos.py
@ -1,388 +0,0 @@
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from pathlib import Path
-
-import pytest
-import torch
-
-from chronos import (
-    BaseChronosPipeline,
-    ChronosConfig,
-    ChronosPipeline,
-    MeanScaleUniformBins,
-)
-from test.util import validate_tensor
-
-
-def test_base_chronos_pipeline_loads_from_huggingface():
-    BaseChronosPipeline.from_pretrained("amazon/chronos-t5-tiny", device_map="cpu")
-
-
-@pytest.mark.parametrize("n_numerical_tokens", [5, 10, 27])
-@pytest.mark.parametrize("n_special_tokens", [2, 5, 13])
-def test_tokenizer_consistency(n_numerical_tokens: int, n_special_tokens: int):
-    n_tokens = n_numerical_tokens + n_special_tokens
-
-    config = ChronosConfig(
-        tokenizer_class="MeanScaleUniformBins",
-        tokenizer_kwargs=dict(low_limit=-1.0, high_limit=1.0),
-        n_tokens=n_tokens,
-        n_special_tokens=n_special_tokens,
-        pad_token_id=0,
-        eos_token_id=1,
-        use_eos_token=True,
-        model_type="seq2seq",
-        context_length=512,
-        prediction_length=64,
-        num_samples=20,
-        temperature=1.0,
-        top_k=50,
-        top_p=1.0,
-    )
-
-    tokenizer = config.create_tokenizer()
-    assert isinstance(tokenizer, MeanScaleUniformBins)
-
-    context = tokenizer.centers.unsqueeze(0)  # add batch dimension
-    scale = torch.ones((1,))  # fix the scale to one to turn off scaling
-
-    token_ids, _, _ = tokenizer._input_transform(context, scale=scale)
-
-    samples = tokenizer.output_transform(
-        token_ids.unsqueeze(1),  # add sample dimension
-        scale=scale,
-    )
-
-    assert (samples[0, 0, :] == context).all()
-
-
-@pytest.mark.xfail
-@pytest.mark.parametrize("n_numerical_tokens", [5, 10, 27])
-@pytest.mark.parametrize("n_special_tokens", [2, 5, 13])
-@pytest.mark.parametrize("use_eos_token", [False, True])
-def test_tokenizer_fixed_data(
-    n_numerical_tokens: int, n_special_tokens: int, use_eos_token: bool
-):
-    n_tokens = n_numerical_tokens + n_special_tokens
-    context_length = 3
-
-    config = ChronosConfig(
-        tokenizer_class="MeanScaleUniformBins",
-        tokenizer_kwargs=dict(low_limit=-1.0, high_limit=1.0),
-        n_tokens=n_tokens,
-        n_special_tokens=n_special_tokens,
-        pad_token_id=0,
-        eos_token_id=1,
-        use_eos_token=use_eos_token,
-        model_type="seq2seq",
-        context_length=512,
-        prediction_length=64,
-        num_samples=20,
-        temperature=1.0,
-        top_k=50,
-        top_p=1.0,
-    )
-
-    tokenizer = config.create_tokenizer()
-
-    context = torch.tensor(
-        [
-            [-3.7, 3.7],
-            [-42.0, 42.0],
-        ]
-    )
-    batch_size, _ = context.shape
-
-    token_ids, attention_mask, scale = tokenizer.context_input_transform(context)
-
-    assert token_ids.shape == (batch_size, context_length + 1 * use_eos_token)
-    assert all(token_ids[:, 0] == torch.tensor([0]).repeat(batch_size))
-    assert all(token_ids[:, 1] == torch.tensor([n_special_tokens]).repeat(batch_size))
-    assert all(token_ids[:, 2] == torch.tensor([n_tokens - 1]).repeat(batch_size))
-
-    if use_eos_token:
-        assert all(token_ids[:, 3] == torch.tensor([1]).repeat(batch_size))
-
-    samples = tokenizer.output_transform(
-        torch.arange(n_special_tokens, n_tokens).unsqueeze(0).repeat(batch_size, 1, 1),
-        tokenizer_state=scale,
-    )
-
-    assert (samples[:, 0, [0, -1]] == context).all()
-
-
-@pytest.mark.xfail
-@pytest.mark.parametrize("use_eos_token", [False, True])
-def test_tokenizer_random_data(use_eos_token: bool):
-    context_length = 8
-    n_tokens = 256
-    n_special_tokens = 2
-
-    config = ChronosConfig(
-        tokenizer_class="MeanScaleUniformBins",
-        tokenizer_kwargs=dict(low_limit=-1.0, high_limit=1.0),
-        n_tokens=n_tokens,
-        n_special_tokens=n_special_tokens,
-        pad_token_id=0,
-        eos_token_id=1,
-        use_eos_token=use_eos_token,
-        model_type="seq2seq",
-        context_length=context_length,
-        prediction_length=64,
-        num_samples=20,
-        temperature=1.0,
-        top_k=50,
-        top_p=1.0,
-    )
-
-    tokenizer = config.create_tokenizer()
-
-    context = torch.tensor(
-        [
-            [torch.nan, torch.nan, 1.0, 1.1, torch.nan, 2.0],
-            [3.0, torch.nan, 3.9, 4.0, 4.1, 4.9],
-        ]
-    )
-
-    token_ids, attention_mask, scale = tokenizer.context_input_transform(context)
-
-    assert token_ids.shape == (
-        *context.shape[:-1],
-        context_length + 1 * use_eos_token,
-    )
-    assert attention_mask.shape == (
-        *context.shape[:-1],
-        context_length + 1 * use_eos_token,
-    )
-    assert scale.shape == context.shape[:1]
-
-    sample_ids = torch.randint(low=n_special_tokens, high=n_tokens, size=(2, 10, 4))
-    sample_ids[0, 0, 0] = n_special_tokens
-    sample_ids[-1, -1, -1] = n_tokens - 1
-
-    samples = tokenizer.output_transform(sample_ids, scale)
-
-    assert samples.shape == (2, 10, 4)
-
-
-@pytest.mark.parametrize("model_dtype", [torch.float32, torch.bfloat16])
-@pytest.mark.parametrize("input_dtype", [torch.float32, torch.bfloat16, torch.int64])
-def test_pipeline_predict(model_dtype: torch.dtype, input_dtype: torch.dtype):
-    pipeline = ChronosPipeline.from_pretrained(
-        Path(__file__).parent / "dummy-chronos-model",
-        device_map="cpu",
-        torch_dtype=model_dtype,
-    )
-    context = 10 * torch.rand(size=(4, 16)) + 10
-    context = context.to(dtype=input_dtype)
-
-    # input: tensor of shape (batch_size, context_length)
-
-    samples = pipeline.predict(context, num_samples=12, prediction_length=3)
-    validate_tensor(samples, shape=(4, 12, 3), dtype=torch.float32)
-
-    with pytest.raises(ValueError):
-        samples = pipeline.predict(
-            context, num_samples=7, prediction_length=65, limit_prediction_length=True
-        )
-
-    samples = pipeline.predict(
-        context, num_samples=7, prediction_length=65, limit_prediction_length=False
-    )
-    validate_tensor(samples, shape=(4, 7, 65), dtype=torch.float32)
-
-    # input: batch_size-long list of tensors of shape (context_length,)
-
-    samples = pipeline.predict(list(context), num_samples=12, prediction_length=3)
-    validate_tensor(samples, shape=(4, 12, 3), dtype=torch.float32)
-
-    with pytest.raises(ValueError):
-        samples = pipeline.predict(
-            list(context),
-            num_samples=7,
-            prediction_length=65,
-            limit_prediction_length=True,
-        )
-
-    samples = pipeline.predict(
-        list(context),
-        num_samples=7,
-        prediction_length=65,
-        limit_prediction_length=False,
-    )
-    validate_tensor(samples, shape=(4, 7, 65), dtype=torch.float32)
-
-    # input: tensor of shape (context_length,)
-
-    samples = pipeline.predict(context[0, ...], num_samples=12, prediction_length=3)
-    validate_tensor(samples, shape=(1, 12, 3), dtype=torch.float32)
-
-    with pytest.raises(ValueError):
-        samples = pipeline.predict(
-            context[0, ...],
-            num_samples=7,
-            prediction_length=65,
-            limit_prediction_length=True,
-        )
-
-    samples = pipeline.predict(
-        context[0, ...],
-        num_samples=7,
-        prediction_length=65,
-    )
-    validate_tensor(samples, shape=(1, 7, 65), dtype=torch.float32)
-
-
-@pytest.mark.parametrize("model_dtype", [torch.float32, torch.bfloat16])
-@pytest.mark.parametrize("input_dtype", [torch.float32, torch.bfloat16, torch.int64])
-@pytest.mark.parametrize("prediction_length", [3, 65])
-@pytest.mark.parametrize(
-    "quantile_levels", [[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], [0.1, 0.5, 0.9]]
-)
-def test_pipeline_predict_quantiles(
-    model_dtype: torch.dtype,
-    input_dtype: torch.dtype,
-    prediction_length: int,
-    quantile_levels: list[int],
-):
-    pipeline = ChronosPipeline.from_pretrained(
-        Path(__file__).parent / "dummy-chronos-model",
-        device_map="cpu",
-        torch_dtype=model_dtype,
-    )
-    context = 10 * torch.rand(size=(4, 16)) + 10
-    context = context.to(dtype=input_dtype)
-
-    num_expected_quantiles = len(quantile_levels)
-    # input: tensor of shape (batch_size, context_length)
-
-    quantiles, mean = pipeline.predict_quantiles(
-        context,
-        num_samples=12,
-        prediction_length=prediction_length,
-        quantile_levels=quantile_levels,
-    )
-    validate_tensor(
-        quantiles, (4, prediction_length, num_expected_quantiles), dtype=torch.float32
-    )
-    validate_tensor(mean, (4, prediction_length), dtype=torch.float32)
-
-    # input: batch_size-long list of tensors of shape (context_length,)
-
-    quantiles, mean = pipeline.predict_quantiles(
-        list(context),
-        num_samples=12,
-        prediction_length=prediction_length,
-        quantile_levels=quantile_levels,
-    )
-    validate_tensor(
-        quantiles, (4, prediction_length, num_expected_quantiles), dtype=torch.float32
-    )
-    validate_tensor(mean, (4, prediction_length), dtype=torch.float32)
-
-    # input: tensor of shape (context_length,)
-
-    quantiles, mean = pipeline.predict_quantiles(
-        context[0, ...],
-        num_samples=12,
-        prediction_length=prediction_length,
-        quantile_levels=quantile_levels,
-    )
-    validate_tensor(
-        quantiles, (1, prediction_length, num_expected_quantiles), dtype=torch.float32
-    )
-    validate_tensor(mean, (1, prediction_length), dtype=torch.float32)
-
-
-@pytest.mark.parametrize("model_dtype", [torch.float32, torch.bfloat16])
-@pytest.mark.parametrize("input_dtype", [torch.float32, torch.bfloat16, torch.int64])
-def test_pipeline_embed(model_dtype: torch.dtype, input_dtype: torch.dtype):
-    pipeline = ChronosPipeline.from_pretrained(
-        Path(__file__).parent / "dummy-chronos-model",
-        device_map="cpu",
-        torch_dtype=model_dtype,
-    )
-    d_model = pipeline.model.model.config.d_model
-    context = 10 * torch.rand(size=(4, 16)) + 10
-    context = context.to(dtype=input_dtype)
-    expected_embed_length = 16 + (1 if pipeline.model.config.use_eos_token else 0)
-
-    # input: tensor of shape (batch_size, context_length)
-
-    embedding, scale = pipeline.embed(context)
-    validate_tensor(
-        embedding, shape=(4, expected_embed_length, d_model), dtype=model_dtype
-    )
-    validate_tensor(scale, shape=(4,), dtype=torch.float32)
-
-    # input: batch_size-long list of tensors of shape (context_length,)
-
-    embedding, scale = pipeline.embed(list(context))
-    validate_tensor(
-        embedding, shape=(4, expected_embed_length, d_model), dtype=model_dtype
-    )
-    validate_tensor(scale, shape=(4,), dtype=torch.float32)
-
-    # input: tensor of shape (context_length,)
-    embedding, scale = pipeline.embed(context[0, ...])
-    validate_tensor(
-        embedding, shape=(1, expected_embed_length, d_model), dtype=model_dtype
-    )
-    validate_tensor(scale, shape=(1,), dtype=torch.float32)
-
-
-@pytest.mark.parametrize("n_tokens", [10, 1000, 10000])
-def test_tokenizer_number_of_buckets(n_tokens):
-    config = ChronosConfig(
-        tokenizer_class="MeanScaleUniformBins",
-        tokenizer_kwargs=dict(low_limit=-1.0, high_limit=1.0),
-        n_tokens=n_tokens,
-        n_special_tokens=2,
-        pad_token_id=0,
-        eos_token_id=1,
-        use_eos_token=True,
-        model_type="seq2seq",
-        context_length=512,
-        prediction_length=64,
-        num_samples=20,
-        temperature=1.0,
-        top_k=50,
-        top_p=1.0,
-    )
-    tokenizer = config.create_tokenizer()
-
-    n_numerical_tokens = config.n_tokens - config.n_special_tokens
-
-    # The tokenizer has one bucket too many as a result of an early bug. In order to
-    # keep consistent with the original trained models, this is kept as it is. However,
-    # token ids are clipped to a maximum of `n_tokens - 1` to avoid out-of-bounds errors.
-    assert len(tokenizer.centers) == (n_numerical_tokens - 1)
-    assert len(tokenizer.boundaries) == n_numerical_tokens
-
-
-@pytest.mark.parametrize("n_tokens", [10, 1000, 10000])
-def test_token_clipping(n_tokens):
-    config = ChronosConfig(
-        tokenizer_class="MeanScaleUniformBins",
-        tokenizer_kwargs={"low_limit": -15, "high_limit": 15},
-        n_tokens=n_tokens,
-        n_special_tokens=2,
-        pad_token_id=0,
-        eos_token_id=1,
-        use_eos_token=True,
-        model_type="seq2seq",
-        context_length=512,
-        prediction_length=64,
-        num_samples=20,
-        temperature=1.0,
-        top_k=50,
-        top_p=1.0,
-    )
-    tokenizer = config.create_tokenizer()
-
-    huge_value = 1e22  # this large value is assigned to the largest bucket
-    token_ids, _, _ = tokenizer._input_transform(
-        context=torch.tensor([[huge_value]]), scale=torch.tensor(([1]))
-    )
-    assert token_ids[0, 0] == config.n_tokens - 1  # and it's clipped to n_tokens - 1
--- a/test/test_chronos_bolt.py
+++ b/test/test_chronos_bolt.py
@ -1,301 +0,0 @@
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-from pathlib import Path
-
-import pytest
-import torch
-
-from chronos import BaseChronosPipeline, ChronosBoltPipeline
-from chronos.chronos_bolt import InstanceNorm, Patch
-from test.util import validate_tensor
-
-
-def test_base_chronos_pipeline_loads_from_huggingface():
-    BaseChronosPipeline.from_pretrained("amazon/chronos-bolt-tiny", device_map="cpu")
-
-
-@pytest.mark.parametrize("torch_dtype", [torch.float32, torch.bfloat16])
-@pytest.mark.parametrize("input_dtype", [torch.float32, torch.bfloat16, torch.int64])
-def test_pipeline_predict(torch_dtype: torch.dtype, input_dtype: torch.dtype):
-    pipeline = ChronosBoltPipeline.from_pretrained(
-        Path(__file__).parent / "dummy-chronos-bolt-model",
-        device_map="cpu",
-        torch_dtype=torch_dtype,
-    )
-    context = 10 * torch.rand(size=(4, 16)) + 10
-    context = context.to(dtype=input_dtype)
-    expected_num_quantiles = len(pipeline.quantiles)
-
-    # input: tensor of shape (batch_size, context_length)
-
-    quantiles = pipeline.predict(context, prediction_length=3)
-    validate_tensor(quantiles, (4, expected_num_quantiles, 3), dtype=torch.float32)
-
-    with pytest.raises(ValueError):
-        quantiles = pipeline.predict(
-            context, prediction_length=65, limit_prediction_length=True
-        )
-
-    quantiles = pipeline.predict(context, prediction_length=65)
-    validate_tensor(quantiles, (4, expected_num_quantiles, 65))
-
-    # input: batch_size-long list of tensors of shape (context_length,)
-
-    quantiles = pipeline.predict(list(context), prediction_length=3)
-    validate_tensor(quantiles, (4, expected_num_quantiles, 3), dtype=torch.float32)
-
-    with pytest.raises(ValueError):
-        quantiles = pipeline.predict(
-            list(context),
-            prediction_length=65,
-            limit_prediction_length=True,
-        )
-
-    quantiles = pipeline.predict(list(context), prediction_length=65)
-    validate_tensor(quantiles, (4, expected_num_quantiles, 65), dtype=torch.float32)
-
-    # input: tensor of shape (context_length,)
-
-    quantiles = pipeline.predict(context[0, ...], prediction_length=3)
-    validate_tensor(quantiles, (1, expected_num_quantiles, 3), dtype=torch.float32)
-
-    with pytest.raises(ValueError):
-        quantiles = pipeline.predict(
-            context[0, ...],
-            prediction_length=65,
-            limit_prediction_length=True,
-        )
-
-    quantiles = pipeline.predict(
-        context[0, ...],
-        prediction_length=65,
-    )
-    validate_tensor(quantiles, (1, expected_num_quantiles, 65), dtype=torch.float32)
-
-
-@pytest.mark.parametrize("torch_dtype", [torch.float32, torch.bfloat16])
-@pytest.mark.parametrize("input_dtype", [torch.float32, torch.bfloat16, torch.int64])
-@pytest.mark.parametrize("prediction_length", [3, 65])
-@pytest.mark.parametrize(
-    "quantile_levels", [[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], [0.1, 0.5, 0.9]]
-)
-def test_pipeline_predict_quantiles(
-    torch_dtype: torch.dtype,
-    input_dtype: torch.dtype,
-    prediction_length: int,
-    quantile_levels: list[int],
-):
-    pipeline = ChronosBoltPipeline.from_pretrained(
-        Path(__file__).parent / "dummy-chronos-bolt-model",
-        device_map="cpu",
-        torch_dtype=torch_dtype,
-    )
-    context = 10 * torch.rand(size=(4, 16)) + 10
-    context = context.to(dtype=input_dtype)
-
-    num_expected_quantiles = len(quantile_levels)
-    # input: tensor of shape (batch_size, context_length)
-
-    quantiles, mean = pipeline.predict_quantiles(
-        context,
-        prediction_length=prediction_length,
-        quantile_levels=quantile_levels,
-    )
-    validate_tensor(
-        quantiles, (4, prediction_length, num_expected_quantiles), dtype=torch.float32
-    )
-    validate_tensor(mean, (4, prediction_length), dtype=torch.float32)
-
-    # input: batch_size-long list of tensors of shape (context_length,)
-
-    quantiles, mean = pipeline.predict_quantiles(
-        list(context),
-        prediction_length=prediction_length,
-        quantile_levels=quantile_levels,
-    )
-    validate_tensor(
-        quantiles, (4, prediction_length, num_expected_quantiles), dtype=torch.float32
-    )
-    validate_tensor(mean, (4, prediction_length), dtype=torch.float32)
-
-    # input: tensor of shape (context_length,)
-
-    quantiles, mean = pipeline.predict_quantiles(
-        context[0, ...],
-        prediction_length=prediction_length,
-        quantile_levels=quantile_levels,
-    )
-    validate_tensor(
-        quantiles, (1, prediction_length, num_expected_quantiles), dtype=torch.float32
-    )
-    validate_tensor(mean, (1, prediction_length), dtype=torch.float32)
-
-
-@pytest.mark.parametrize("model_dtype", [torch.float32, torch.bfloat16])
-@pytest.mark.parametrize("input_dtype", [torch.float32, torch.bfloat16, torch.int64])
-def test_pipeline_embed(model_dtype: torch.dtype, input_dtype: torch.dtype):
-    pipeline = ChronosBoltPipeline.from_pretrained(
-        Path(__file__).parent / "dummy-chronos-bolt-model",
-        device_map="cpu",
-        torch_dtype=model_dtype,
-    )
-    d_model = pipeline.model.config.d_model
-    context = 10 * torch.rand(size=(4, 16)) + 10
-    context = context.to(dtype=input_dtype)
-
-    # the patch size of dummy model is 16, so only 1 patch is created
-    expected_embed_length = 1 + (
-        1 if pipeline.model.config.chronos_config["use_reg_token"] else 0
-    )
-
-    # input: tensor of shape (batch_size, context_length)
-
-    embedding, loc_scale = pipeline.embed(context)
-    validate_tensor(
-        embedding, shape=(4, expected_embed_length, d_model), dtype=model_dtype
-    )
-    validate_tensor(loc_scale[0], shape=(4,), dtype=torch.float32)
-    validate_tensor(loc_scale[1], shape=(4,), dtype=torch.float32)
-
-    # input: batch_size-long list of tensors of shape (context_length,)
-
-    embedding, loc_scale = pipeline.embed(list(context))
-    validate_tensor(
-        embedding, shape=(4, expected_embed_length, d_model), dtype=model_dtype
-    )
-    validate_tensor(loc_scale[0], shape=(4,), dtype=torch.float32)
-    validate_tensor(loc_scale[1], shape=(4,), dtype=torch.float32)
-
-    # input: tensor of shape (context_length,)
-    embedding, loc_scale = pipeline.embed(context[0, ...])
-    validate_tensor(
-        embedding, shape=(1, expected_embed_length, d_model), dtype=model_dtype
-    )
-    validate_tensor(loc_scale[0], shape=(1,), dtype=torch.float32)
-    validate_tensor(loc_scale[1], shape=(1,), dtype=torch.float32)
-
-
-# The following tests have been taken from
-# https://github.com/autogluon/autogluon/blob/f57beb26cb769c6e0d484a6af2b89eab8aee73a8/timeseries/tests/unittests/models/chronos/pipeline/test_chronos_bolt.py
-# Author: Caner Turkmen <atturkm@amazon.com>
-
-
-def test_given_even_data_patch_operator_output_is_correct():
-    batch_size = 17
-    patch_len = 16
-
-    patch = Patch(patch_len, patch_len)
-
-    batch = (
-        torch.stack([torch.arange(512)] * batch_size)
-        + torch.arange(batch_size)[:, None]
-    )
-    output = patch(batch)
-
-    assert output.shape == (batch_size, 512 // patch_len, patch_len)
-
-    assert torch.allclose(
-        output[:, 0],
-        torch.stack([torch.arange(patch_len)] * batch_size)
-        + torch.arange(batch_size)[:, None],
-        atol=1e-5,
-    )
-    assert torch.allclose(
-        output[:, 1],
-        torch.stack([torch.arange(patch_len, 2 * patch_len)] * batch_size)
-        + torch.arange(batch_size)[:, None],
-        atol=1e-5,
-    )
-    assert not torch.isnan(output).any()
-
-
-def test_given_even_data_and_strides_patch_operator_output_is_correct():
-    batch_size = 17
-    patch_len, patch_stride = 16, 8
-
-    patch = Patch(patch_len, patch_stride)
-
-    offset = torch.arange(batch_size)[:, None]
-    batch = torch.stack([torch.arange(512)] * batch_size) + offset
-    output = patch(batch)
-
-    assert torch.allclose(
-        output[:, 1],
-        torch.stack([torch.arange(patch_stride, patch_stride + patch_len)] * batch_size)
-        + offset,
-        atol=1e-5,
-    )
-    assert not torch.isnan(output).any()
-
-
-def test_given_uneven_data_patch_operator_pads_and_output_is_correct():
-    batch_size = 17
-    patch_len = 16
-
-    patch = Patch(patch_len, patch_len)
-
-    batch = (
-        torch.stack([torch.arange(512 - patch_len + 1)] * batch_size)
-        + torch.arange(batch_size)[:, None]
-    ).float()
-    output = patch(batch)
-
-    assert output.shape == (batch_size, 512 // patch_len, patch_len)
-
-    # check the first portion is padded
-    assert torch.isnan(output[:, 0, :-1]).all()
-
-    # check nowhere else is nan
-    assert not torch.isnan(output[:, 1:]).any()
-
-
-def test_when_instancenorm_applied_then_standardization_correct():
-    inorm = InstanceNorm()
-
-    input_ = torch.tensor(
-        [
-            [1, 2, 3, 4, 5],
-            [2, 3, 4, 5, 6],
-        ]
-    ).float()
-
-    normalized, (loc, scale) = inorm(input_)
-
-    assert normalized.shape == input_.shape
-    assert torch.allclose(normalized[0], normalized[1])
-    assert torch.allclose(loc.squeeze(), torch.tensor([3.0, 4.0]))
-    assert torch.allclose(scale.squeeze(), torch.tensor(1.41421))
-
-
-def test_when_instancenorm_applied_and_reversed_then_nans_preserved():
-    inorm = InstanceNorm()
-
-    input_ = torch.tensor(
-        [
-            [1, torch.nan, 3, 4, 5],
-            [2, 3, 4, 5, torch.nan],
-        ]
-    ).float()
-
-    normalized, (loc, scale) = inorm(input_)
-    assert torch.allclose(normalized.isnan(), input_.isnan())
-
-    output = inorm.inverse(normalized, (loc, scale))
-    assert torch.allclose(output, input_, equal_nan=True)
-
-
-def test_when_instancenorm_applied_and_reversed_then_output_correct():
-    inorm = InstanceNorm()
-
-    input_ = torch.tensor(
-        [
-            [1, 2, 3, 4, 5],
-            [2, 3, 4, 5, 1000],
-        ]
-    ).float()
-
-    normalized, loc_scale = inorm(input_)
-    output = inorm.inverse(normalized, loc_scale)
-
-    assert torch.allclose(output, input_)
--- a/test/test_utils.py
+++ b/test/test_utils.py
@ -1,29 +0,0 @@
-# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
-# SPDX-License-Identifier: Apache-2.0
-
-import pytest
-import torch
-
-from chronos.utils import left_pad_and_stack_1D
-
-
-@pytest.mark.parametrize(
-    "tensors",
-    [
-        [
-            torch.tensor([2.0, 3.0], dtype=dtype),
-            torch.tensor([4.0, 5.0, 6.0], dtype=dtype),
-            torch.tensor([7.0, 8.0, 9.0, 10.0], dtype=dtype),
-        ]
-        for dtype in [torch.int, torch.float16, torch.float32]
-    ],
-)
-def test_pad_and_stack(tensors: list):
-    stacked_and_padded = left_pad_and_stack_1D(tensors)
-
-    assert stacked_and_padded.dtype == torch.float32
-    assert stacked_and_padded.shape == (len(tensors), max(len(t) for t in tensors))
-
-    ref = torch.concat(tensors).to(dtype=stacked_and_padded.dtype)
-
-    assert torch.sum(torch.nan_to_num(stacked_and_padded, nan=0)) == torch.sum(ref)
--- a/test/util.py
+++ b/test/util.py
@ -1,13 +0,0 @@
-from typing import Optional, Tuple
-
-import torch
-
-
-def validate_tensor(
-    a: torch.Tensor, shape: Tuple[int, ...], dtype: Optional[torch.dtype] = None
-) -> None:
-    assert isinstance(a, torch.Tensor)
-    assert a.shape == shape
-
-    if dtype is not None:
-        assert a.dtype == dtype