TimeMoEForecaster

Quickstart

from sktime.forecasting.timemoe import TimeMoEForecaster

estimator = TimeMoEForecaster(model_path: str | None='Maple728/TimeMoE-50M', config: dict=None, seed: int=None, use_source_package: bool=False, ignore_deps: bool=False, context_length: int=1024, stride: int=None, training_args: dict=None, validation_split: float | None=0.2, device: str | None=None, dtype=None)

Parameters(11)

model_path: str or None, default=”Maple728/TimeMoE-50M”

Path to the TimeMOE model. This can be:

• A model ID from the HuggingFace Hub, e.g., “Maple728/TimeMoE-50M”

• A local directory containing the model files, specified as an absolute or relative path to the current working directory The path should point to a directory containing the model weights and configuration files in the format expected by the HuggingFace Transformers library.

• None to initialize from config with random weights (from-scratch)

config: dict, optional

A dictionary specifying the configuration of the TimeMOE model. The available configuration options include hyperparameters that control the prediction behavior, sampling, and hardware utilization.

• input_size: int, default=1

  The size of the input time series.

seed: int, optional (default=None)

Seed for reproducibility.

use_source_package: bool, optional (default=False)

If True, the model will be loaded directly from the source package TimeMoE. This is useful if you want to bypass the local version of the package or when working in an environment where the latest updates from the source package are needed. If False, the model will be loaded from the local version of package maintained in sktime. To install the source package, follow the instructions here [1].

ignore_deps: bool, optional, default=False

If True, dependency checks will be ignored, and the user is expected to handle the installation of required packages manually. If False, the class will enforce the default dependencies required for Chronos.

context_lengthint, optional (default=1024)

Sliding-window length for pretrain. For small datasets, use a shorter length with stride=1.

strideint, optional (default=None)

Sliding-window stride for pretrain. Defaults to context_length.

training_argsdict, optional (default=None)

Keyword arguments used for training. Supports all arguments by transformers.TrainingArguments [Rf74a85329ff4-3].

Additionally, the following arguments are supported: - min_learning_rate: float, default=0

Minimum learning rate for cosine_schedule

validation_splitfloat or None, default=0.2

Fraction of data reserved for evaluation when pretrain is used. If None, no evaluation dataset is created.

devicestr, optional (default=None)

Device placement passed to transformers device_map, for example "cpu", "cuda", or "auto". If None, config["device_map"] is used.

dtypetorch.dtype, optional (default=None)

Torch dtype used when loading the model and preparing prediction inputs. None keeps the checkpoint’s native dtype for from_pretrained (and the initialized dtype for from-scratch models); prediction inputs then follow the loaded model dtype.

Examples

Zero-shot forecasting:
>>> from sktime.forecasting.timemoe import TimeMoEForecaster
>>> from sktime.datasets import load_airline
>>> from sktime.forecasting.model_selection import temporal_train_test_split
>>> y = load_airline ()
>>> forecaster = TimeMoEForecaster ("Maple728/TimeMoE-50M")
>>> forecaster. fit (y_train)
>>> y_pred = forecaster. predict (fh = [1, 2, 3 ], y = y_test) Fine-tuning of a pretrained checkpoint:
>>> from sktime.utils._testing.hierarchical import _make_hierarchical
>>> y_panel = _make_hierarchical (
... hierarchy_levels = (3,), min_timepoints = 64, max_timepoints = 128,
... )
>>> forecaster = TimeMoEForecaster (
... model_path = "Maple728/TimeMoE-50M",
... context_length = 32,
... stride = 1,
... training_args = { "max_steps": 10, "per_device_train_batch_size": 2 },
... )
>>> forecaster. pretrain (y_panel)
>>> forecaster. fit (load_airline ())
>>> y_pred = forecaster. predict (fh = [1, 2, 3 ]) Training from scratch:
>>> forecaster = TimeMoEForecaster (
... model_path = None,
... config = {
... "hidden_size": 64,
... "intermediate_size": 128,
... "num_hidden_layers": 2,
... "num_attention_heads": 4,
... "num_experts": 2,
... "num_experts_per_tok": 1,
... "horizon_lengths": [1 ],
... "max_position_embeddings": 128,
... },
... context_length = 32,
... stride = 1,
... training_args = { "max_steps": 10 },
... )
>>> forecaster. pretrain (y_panel)
>>> forecaster. fit (load_airline ())
>>> y_pred = forecaster. predict (fh = [1, 2, 3 ])

References