Merge ICML-acc version

.github/ISSUE_TEMPLATE/bug_report.md

@@ -0,0 +1,32 @@
+---
+name: Bug report
+about: Create a report to help us improve
+title: ''
+labels: bug
+assignees: ''
+
+---
+
+**Describe the bug**
+A clear and concise description of what the bug is.
+
+**To Reproduce**
+Please provide a code snippet of a minimal reproducible example for the error. 
+```python
+put code here
+```
+
+**Expected behavior**
+A clear and concise description of what you expected to happen.
+
+**Error message or code output**
+Paste the complete error message, including stack trace, or the undesired output that the above snippet produces.
+```shell
+put error or undesired output here
+```
+
+**Environment**
+* Operating system:
+* Python version:
+* PyTorch version:
+* uni2ts version (have you pulled the latest version?)

.github/ISSUE_TEMPLATE/config.yml

@@ -0,0 +1,5 @@
+blank_issues_enabled: false
+contact_links:
+  - name: Questions and Discussions
+    url: https://github.com/SalesforceAIResearch/uni2ts/discussions
+    about: Use GitHub Discussions to ask and answer questions, exchange ideas, and share learning.

.github/ISSUE_TEMPLATE/feature_request.md

@@ -0,0 +1,20 @@
+---
+name: Feature request
+about: Suggest an idea for this project
+title: ''
+labels: ''
+assignees: ''
+
+---
+
+**Is your feature request related to a problem? Please describe.**
+A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
+
+**Describe the solution you'd like**
+A clear and concise description of what you want to happen.
+
+**Describe alternatives you've considered**
+A clear and concise description of any alternative solutions or features you've considered.
+
+**Additional context**
+Add any other context or screenshots about the feature request here.

.github/workflows/lint.yml

@@ -0,0 +1,21 @@
+# This workflow will install Python dependencies, run tests and lint with a single version of Python
+# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python
+
+name: Lint
+
+on: [push, pull_request]
+
+jobs:
+  black:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - uses: psf/black@stable
+        with:
+          options: "--check --verbose --line-length 88"
+          jupyter: true
+  isort:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - uses: isort/isort-action@v1

.github/workflows/test.yml

@@ -0,0 +1,20 @@
+name: Test
+
+on: [push, pull_request]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python 3.10
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.10'
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e '.[dev,notebook]'
+      - name: Test with pytest
+        run: |
+          pytest test

CITATION.cff

@@ -0,0 +1,38 @@
+cff-version: 1.2.0
+title: Unified Training of Universal Time Series Forecasting Transformers
+message: If you find Chronos models useful for your research, please consider citing the associated paper.
+authors:
+  - family-names: Woo
+    given-names: Gerald
+  - family-names: Liu
+    given-names: Chenghao
+  - family-names: Kumar
+    given-names: Akshat
+  - family-names: Xiong
+    given-names: Caiming
+  - family-names: Savarese
+    given-names: Silvio
+  - family-names: Sahoo
+    given-names: Doyen
+preferred-citation:
+  type: conference-paper
+  authors:
+    - family-names: Woo
+      given-names: Gerald
+    - family-names: Liu
+      given-names: Chenghao
+    - family-names: Kumar
+      given-names: Akshat
+    - family-names: Xiong
+      given-names: Caiming
+    - family-names: Savarese
+      given-names: Silvio
+    - family-names: Sahoo
+      given-names: Doyen
+  title: Unified Training of Universal Time Series Forecasting Transformers
+  year: 2024
+  collection-title: Proceedings of the 41st International Conference on Machine Learning
+  conference:
+    name: International Conference on Machine Learning (ICML)
+    city: Vienna
+    country: Austria

README.md

@@ -7,6 +7,8 @@ Uni2TS also provides tools for fine-tuning, inference, and evaluation for time s
 
 ## 🎉 What's New
 
+* May 2024: The Uni2TS paper has been accepted to ICML 2024 as an Oral presentation!
+
 * Mar 2024: Release of Uni2TS library, along with [Moirai-1.0-R](https://huggingface.co/collections/Salesforce/moirai-10-r-models-65c8d3a94c51428c300e0742) and [LOTSA data](https://huggingface.co/datasets/Salesforce/lotsa_data/)!
 
 ## ✅ TODO
@@ -64,7 +66,7 @@ from gluonts.dataset.split import split
 from huggingface_hub import hf_hub_download
 
 from uni2ts.eval_util.plot import plot_single
-from uni2ts.model.moirai import MoiraiForecast
+from uni2ts.model.moirai import MoiraiForecast, MoiraiModule
 
 
 SIZE = "small"  # model size: choose from {'small', 'base', 'large'}
@@ -97,18 +99,15 @@ test_data = test_template.generate_instances(
 )
 
 # Prepare pre-trained model by downloading model weights from huggingface hub
-model = MoiraiForecast.load_from_checkpoint(
-    checkpoint_path=hf_hub_download(
-        repo_id=f"Salesforce/moirai-1.0-R-{SIZE}", filename="model.ckpt"
-    ),
+model = MoiraiForecast(
+    module=MoiraiModule.from_pretrained(f"Salesforce/moirai-1.0-R-{SIZE}"),
     prediction_length=PDT,
     context_length=CTX,
     patch_size=PSZ,
     num_samples=100,
     target_dim=1,
     feat_dynamic_real_dim=ds.num_feat_dynamic_real,
     past_feat_dynamic_real_dim=ds.num_past_feat_dynamic_real,
-    map_location="cuda:0" if torch.cuda.is_available() else "cpu",
 )
 
 predictor = model.create_predictor(batch_size=BSZ)

cli/conf/eval/model/moirai_1.0_R_base.yaml

@@ -1,10 +1,8 @@
 _target_: uni2ts.model.moirai.MoiraiForecast.load_from_checkpoint
-#checkpoint_path:
-#  _target_: huggingface_hub.hf_hub_download
-#  repo_id: Salesforce/moirai-1.0-R-base
-#  filename: model.ckpt
-
-checkpoint_path: /home/eee/qzz/uni2ts/outputs/finetune/moirai_1.0_R_base/etth1/my_lsf_run/checkpoints/epoch=19-step=200.ckpt
+checkpoint_path:
+  _target_: huggingface_hub.hf_hub_download
+  repo_id: Salesforce/moirai-1.0-R-base
+  filename: model.ckpt
 num_samples: 100
 patch_size: ???
 context_length: ???

cli/conf/eval/model/moirai_1.0_R_large.yaml

@@ -1,8 +1,7 @@
-_target_: uni2ts.model.moirai.MoiraiForecast.load_from_checkpoint
-checkpoint_path:
-  _target_: huggingface_hub.hf_hub_download
-  repo_id: Salesforce/moirai-1.0-R-large
-  filename: model.ckpt
+_target_: uni2ts.model.moirai.MoiraiForecast
+module:
+  _target_: uni2ts.model.moirai.MoiraiModule.from_pretrained
+  pretrained_model_name_or_path: Salesforce/moirai-1.0-R-large
 num_samples: 100
 patch_size: ???
 context_length: ???

cli/conf/eval/model/moirai_1.0_R_small.yaml

@@ -1,9 +1,8 @@
 _target_: uni2ts.model.moirai.MoiraiForecast.load_from_checkpoint
-#checkpoint_path:
-#  _target_: huggingface_hub.hf_hub_download
-#  repo_id: Salesforce/moirai-1.0-R-small
-#  filename: model.ckpt
-checkpoint_path: /home/eee/qzz/uni2ts/outputs/finetune/moirai_1.0_R_small/ettm1/my_lsf_run/checkpoints/epoch=0-step=5-v2.ckpt
+checkpoint_path:
+  _target_: huggingface_hub.hf_hub_download
+  repo_id: Salesforce/moirai-1.0-R-small
+  filename: model.ckpt
 num_samples: 100
 patch_size: ???
 context_length: ???

cli/conf/finetune/default.yaml

@@ -1,6 +1,6 @@
 hydra:
   run:
-    dir: outputs/${hydra:job.name}/${hydra:runtime.choices.model}/${hydra:runtime.choices.data}/${run_name}
+    dir: outputs/finetune/${hydra:runtime.choices.model}/${hydra:runtime.choices.data}/${run_name}
 defaults:
   - model: ???
   - data: ???

cli/conf/finetune/model/moirai_1.0_R_base.yaml

@@ -1,4 +1,8 @@
-_target_: uni2ts.model.moirai.MoiraiFinetune.load_from_checkpoint
+# load a pretrained checkpoint from huggingface hub
+_target_: uni2ts.model.moirai.MoiraiFinetune
+module:
+  _target_: uni2ts.model.moirai.MoiraiModule.from_pretrained
+  pretrained_model_name_or_path: Salesforce/moirai-1.0-R-base
 module_kwargs:
   _target_: builtins.dict
   distr_output:
@@ -26,8 +30,4 @@ weight_decay: 1e-1
 beta1: 0.9
 beta2: 0.98
 num_training_steps: ${mul:${trainer.max_epochs},${train_dataloader.num_batches_per_epoch}}
-num_warmup_steps: 0
-checkpoint_path:
-  _target_: huggingface_hub.hf_hub_download
-  repo_id: Salesforce/moirai-1.0-R-base
-  filename: model.ckpt
+num_warmup_steps: 0

cli/conf/finetune/model/moirai_1.0_R_large.yaml

@@ -1,4 +1,8 @@
-_target_: uni2ts.model.moirai.MoiraiFinetune.load_from_checkpoint
+# load a pretrained checkpoint from huggingface hub
+_target_: uni2ts.model.moirai.MoiraiFinetune
+module:
+  _target_: uni2ts.model.moirai.MoiraiModule.from_pretrained
+  pretrained_model_name_or_path: Salesforce/moirai-1.0-R-large
 module_kwargs:
   _target_: builtins.dict
   distr_output:
@@ -26,8 +30,4 @@ weight_decay: 1e-1
 beta1: 0.9
 beta2: 0.98
 num_training_steps: ${mul:${trainer.max_epochs},${train_dataloader.num_batches_per_epoch}}
-num_warmup_steps: 0
-checkpoint_path:
-  _target_: huggingface_hub.hf_hub_download
-  repo_id: Salesforce/moirai-1.0-R-large
-  filename: model.ckpt
+num_warmup_steps: 0

cli/conf/finetune/model/moirai_1.0_R_small.yaml

@@ -1,4 +1,8 @@
-_target_: uni2ts.model.moirai.MoiraiFinetune.load_from_checkpoint
+# load a pretrained checkpoint from huggingface hub
+_target_: uni2ts.model.moirai.MoiraiFinetune
+module:
+  _target_: uni2ts.model.moirai.MoiraiModule.from_pretrained
+  pretrained_model_name_or_path: Salesforce/moirai-1.0-R-small
 module_kwargs:
   _target_: builtins.dict
   distr_output:
@@ -30,8 +34,4 @@ weight_decay: 1e-1
 beta1: 0.9
 beta2: 0.98
 num_training_steps: ${mul:${trainer.max_epochs},${train_dataloader.num_batches_per_epoch}}
-num_warmup_steps: 0
-checkpoint_path:
-  _target_: huggingface_hub.hf_hub_download
-  repo_id: Salesforce/moirai-1.0-R-small
-  filename: model.ckpt
+num_warmup_steps: 0

cli/conf/finetune/model/moirai_base.yaml

@@ -0,0 +1,31 @@
+# load a pytorch lightning checkpoint
+_target_: uni2ts.model.moirai.MoiraiFinetune.load_from_checkpoint
+module_kwargs:
+  _target_: builtins.dict
+  distr_output:
+    _target_: uni2ts.distribution.MixtureOutput
+    components:
+      - _target_: uni2ts.distribution.StudentTOutput
+      - _target_: uni2ts.distribution.NormalFixedScaleOutput
+      - _target_: uni2ts.distribution.NegativeBinomialOutput
+      - _target_: uni2ts.distribution.LogNormalOutput
+  d_model: 768
+  num_layers: 12
+  patch_sizes: ${as_tuple:[8, 16, 32, 64, 128]}
+  max_seq_len: 512
+  attn_dropout_p: 0.0
+  dropout_p: 0.0
+  scaling: true
+min_patches: 2
+min_mask_ratio: 0.15
+max_mask_ratio: 0.5
+max_dim: 128
+loss_func:
+  _target_: uni2ts.loss.packed.PackedNLLLoss
+lr: 1e-3
+weight_decay: 1e-1
+beta1: 0.9
+beta2: 0.98
+num_training_steps: ${mul:${trainer.max_epochs},${train_dataloader.num_batches_per_epoch}}
+num_warmup_steps: 0
+checkpoint_path: ...

cli/conf/finetune/model/moirai_large.yaml

@@ -0,0 +1,31 @@
+# load a pytorch lightning checkpoint
+_target_: uni2ts.model.moirai.MoiraiFinetune.load_from_checkpoint
+module_kwargs:
+  _target_: builtins.dict
+  distr_output:
+    _target_: uni2ts.distribution.MixtureOutput
+    components:
+      - _target_: uni2ts.distribution.StudentTOutput
+      - _target_: uni2ts.distribution.NormalFixedScaleOutput
+      - _target_: uni2ts.distribution.NegativeBinomialOutput
+      - _target_: uni2ts.distribution.LogNormalOutput
+  d_model: 1024
+  num_layers: 24
+  patch_sizes: ${as_tuple:[8, 16, 32, 64, 128]}
+  max_seq_len: 512
+  attn_dropout_p: 0.0
+  dropout_p: 0.0
+  scaling: true
+min_patches: 2
+min_mask_ratio: 0.15
+max_mask_ratio: 0.5
+max_dim: 128
+loss_func:
+  _target_: uni2ts.loss.packed.PackedNLLLoss
+lr: 1e-3
+weight_decay: 1e-1
+beta1: 0.9
+beta2: 0.98
+num_training_steps: ${mul:${trainer.max_epochs},${train_dataloader.num_batches_per_epoch}}
+num_warmup_steps: 0
+checkpoint_path: ...