Update readme

README.md

@@ -1,136 +1,50 @@
 ---
+title: MOIRAI readme
 license: apache-2.0
 ---
 
-# Moirai-1.0-R-Small
+# Moirai-1.0-R-Small snapshot of [commit](https://huggingface.co/Salesforce/moirai-1.0-R-large/commit/001acc2aca5fb1d5023e2664bbe45471d8baf21a).
+
+A snapshot of the last Moirai-1.0-R-Small weights released under permissive license, hosted by sktime. For use with the sktime Moirai interface to ensure permissive license throughout the stack
 
 Moirai, the Masked Encoder-based Universal Time Series Forecasting Transformer is a Large Time Series Model pre-trained on [LOTSA data](https://huggingface.co/datasets/Salesforce/lotsa_data).
-For more details on the Moirai architecture, training, and results, please refer to the [paper](https://arxiv.org/abs/2402.02592).
+For more details on the Moirai architecture, training, and results, please refer to the [paper](https://arxiv.org/abs/2402.02592).  
 
 <p align="center">
-  <img src="figures/architecture.png" width="100%">
+  <img src="https://huggingface.co/Salesforce/moirai-1.0-R-large/resolve/main/figures/architecture.png", width="100%">
   <br />
   <span>
     Fig. 1: Overall architecture of Moirai. Visualized is a 3-variate time series, where variates 0 and 1 are target variables (i.e. to be forecasted, and variate 2 is a dynamic covariate (values in forecast horizon known). Based on a patch size of 64, each variate is patchified into 3 tokens. The patch embeddings along with sequence and variate id are fed into the Transformer. The shaded patches represent the forecast horizon to be forecasted, whose corresponding output representations are mapped into the mixture distribution parameters.
   </span>
 </p>
 
-## Usage
-
-To perform inference with Moirai, install the uni2ts library from our [GitHub repo](https://github.com/SalesforceAIResearch/uni2ts).
-
-1. Clone repository:
-```shell
-git clone https://github.com/SalesforceAIResearch/uni2ts.git
-cd uni2ts
-```
-
-2) Create virtual environment:
-```shell
-virtualenv venv
-. venv/bin/activate
-```
-
-3) Build from source:
-```shell
-pip install -e '.[notebook]'
-```
-
-4) Create a `.env` file:
-```shell
-touch .env
-```
-
-A simple example to get started:
+## Example
 
 ```python
-import torch
-import pandas as pd
-from gluonts.dataset.pandas import PandasDataset
-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 sktime.datasets import load_tecator
+from sktime.forecasting.moirai_forecaster import MOIRAIForecaster
 
+y, _ = load_tecator(return_X_y=True, return_type="pd-multiindex")
 
-SIZE = "small"  # model size: choose from {'small', 'base', 'large'}
-PDT = 20  # prediction length: any positive integer
-CTX = 200  # context length: any positive integer
-PSZ = "auto"  # patch size: choose from {"auto", 8, 16, 32, 64, 128}
-BSZ = 32  # batch size: any positive integer
-TEST = 100  # test set length: any positive integer
+moirai_forecaster = MOIRAIForecaster(
+    checkpoint_path=f"sktime/moirai-1.0-R-small", 
+    broadcasting=False
+  )
 
-# Read data into pandas DataFrame
-url = (
-    "https://gist.githubusercontent.com/rsnirwan/c8c8654a98350fadd229b00167174ec4"
-    "/raw/a42101c7786d4bc7695228a0f2c8cea41340e18f/ts_wide.csv"
-)
-df = pd.read_csv(url, index_col=0, parse_dates=True)
-
-# Convert into GluonTS dataset
-ds = PandasDataset(dict(df))
-
-# Split into train/test set
-train, test_template = split(
-    ds, offset=-TEST
-)  # assign last TEST time steps as test set
-
-# Construct rolling window evaluation
-test_data = test_template.generate_instances(
-    prediction_length=PDT,  # number of time steps for each prediction
-    windows=TEST // PDT,  # number of windows in rolling window evaluation
-    distance=PDT,  # number of time steps between each window - distance=PDT for non-overlapping windows
-)
-
-# 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-R-{SIZE}", filename="model.ckpt"
-    ),
-    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)
-forecasts = predictor.predict(test_data.input)
-
-input_it = iter(test_data.input)
-label_it = iter(test_data.label)
-forecast_it = iter(forecasts)
-
-inp = next(input_it)
-label = next(label_it)
-forecast = next(forecast_it)
-
-plot_single(
-    inp, 
-    label, 
-    forecast, 
-    context_length=200,
-    name="pred",
-    show_label=True,
-)
+moirai_forecaster.fit(y)
+forecast = moirai_forecaster.predict(fh=range(1, 16))
 ```
 
-## The Moirai Family
-
+## Moirai weights hosted under sktime
 | # Model | # Parameters |
 | :---: | :---: |
-| [Moirai-1.0-R-Small](https://huggingface.co/Salesforce/moirai-1.0-R-small) | 14m |
-| [Moirai-1.0-R-Base](https://huggingface.co/Salesforce/moirai-1.0-R-base) | 91m |
-| [Moirai-1.0-R-Large](https://huggingface.co/Salesforce/moirai-1.0-R-large) | 311m |
+| [Moirai-1.0-R-Small](https://huggingface.co/sktime/moirai-1.0-R-small) | 14m |
+| [Moirai-1.0-R-Base](https://huggingface.co/sktime/moirai-1.0-R-base) | 91m |
+| [Moirai-1.0-R-Large](https://huggingface.co/sktime/moirai-1.0-R-large) | 311m |
 
-## Citation
-
-If you're using Uni2TS in your research or applications, please cite it using this BibTeX:
+Original Weights present at [Salesforce collection](https://huggingface.co/collections/Salesforce/moirai-10-r-models-65c8d3a94c51428c300e0742).
 
+## Citation
 ```markdown
 @article{woo2024unified,
   title={Unified Training of Universal Time Series Forecasting Transformers},