Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 12 additions & 11 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,9 @@ classifiers = [
"Topic :: Software Development :: Libraries :: Python Modules",
]
# NOTE: dependency version notes
# - huggingface-hub can't bump to v1+ until transformers bumps to v5
# - transformers>=5 (CVE-2026-1839 fix) requires huggingface-hub>=1; the upper
# bounds were widened together (transformers<6, huggingface-hub<2). Reaching
# transformers 5 also requires tfc-t0>=0.2 (older tfc-t0 capped hub<1).
# - lightning and pytorch-lightning are the same library
# NOTE: dependencies pinned for faster pip installs
# versions chosen from uv's dependency resolution
Expand All @@ -66,7 +68,7 @@ dependencies = [
"fire",
"fsspec>=2025.9.0",
"gluonts[torch]",
"huggingface-hub>=0.36.2,<1.0",
"huggingface-hub>=0.36.2,<2.0",
"hydra-core>=1.3.2",
"lightgbm>=4.6.0",
"lightning-utilities>=0.15.2",
Expand All @@ -91,18 +93,17 @@ dependencies = [
"statsforecast>=2.0.2",
"tabpfn-time-series==1.0.3 ; python_full_version < '3.13'",
"tensorboard>=2.20.0",
"tfc-t0>=0.1.2 ; python_full_version >= '3.11' and python_full_version < '3.14'",
"timecopilot-chronos-forecasting>=0.2.1",
"timecopilot-granite-tsfm>=0.1.2",
"timecopilot-timesfm>=0.2.1",
"timecopilot-tirex>=0.1.0 ; python_full_version >= '3.11'",
"timecopilot-tirex>=0.1.1",
"tfc-t0>=0.2.0 ; python_full_version >= '3.11' and python_full_version < '3.14'",
"timecopilot-chronos-forecasting>=0.2.2",
"timecopilot-granite-tsfm>=0.2.1 ; python_full_version >= '3.11' and python_full_version < '3.14'",
"timecopilot-timesfm>=0.3.0",
"timecopilot-tirex>=0.1.1 ; python_full_version >= '3.11'",
"timecopilot-toto-2>=0.1.1",
"timecopilot-toto>=0.1.6",
"timecopilot-toto>=0.1.7",
"timecopilot-uni2ts>=0.1.2 ; python_full_version < '3.14'",
"torchmetrics>=1.8.2",
"transformers>=4.41,<5 ; python_full_version < '3.13'",
"transformers>=4.48,<5 ; python_full_version >= '3.13'",
"transformers>=4.41,<6 ; python_full_version < '3.13'",
"transformers>=4.48,<6 ; python_full_version >= '3.13'",
"tsfeatures>=0.4.5",
"utilsforecast[plotting]>=0.2.15",
"xgboost>=3.2.0",
Expand Down
23 changes: 13 additions & 10 deletions tests/models/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,7 @@

from timecopilot.models.ensembles.median import MedianEnsemble
from timecopilot.models.foundation.chronos import Chronos
from timecopilot.models.foundation.flowstate import FlowState
from timecopilot.models.foundation.moirai import Moirai
from timecopilot.models.foundation.patchtst_fm import PatchTSTFM
from timecopilot.models.foundation.timesfm import TimesFM
from timecopilot.models.foundation.toto import Toto
from timecopilot.models.ml import (
Expand Down Expand Up @@ -99,14 +97,6 @@ def disable_mps_session(monkeypatch):
Chronos(repo_id="amazon/chronos-bolt-tiny", alias="Chronos-Bolt"),
Chronos(repo_id="amazon/chronos-2", alias="Chronos-2"),
Chronos(repo_id="amazon/chronos-2", alias="Chronos-2", batch_size=2),
FlowState(repo_id="ibm-research/flowstate"),
FlowState(
repo_id="ibm-granite/granite-timeseries-flowstate-r1",
alias="FlowState-Granite",
),
PatchTSTFM(
context_length=2_048,
),
Toto(context_length=256, batch_size=2),
Toto(
repo_id="Datadog/Toto-2.0-4m",
Expand Down Expand Up @@ -150,6 +140,19 @@ def disable_mps_session(monkeypatch):

models.append(T0(context_length=256, batch_size=2))

if (3, 11) <= sys.version_info < (3, 14):
from timecopilot.models.foundation.flowstate import FlowState
from timecopilot.models.foundation.patchtst_fm import PatchTSTFM

models.append(FlowState(repo_id="ibm-research/flowstate"))
models.append(
FlowState(
repo_id="ibm-granite/granite-timeseries-flowstate-r1",
alias="FlowState-Granite",
)
)
models.append(PatchTSTFM(context_length=2_048))

if sys.version_info < (3, 13):
from tabpfn_time_series import TabPFNMode

Expand Down
9 changes: 9 additions & 0 deletions tests/models/foundation/test_flowstate.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,14 @@
import sys

import numpy as np
import pandas as pd
import pytest

if sys.version_info < (3, 11) or sys.version_info >= (3, 14):
pytest.skip(
"FlowState requires Python >= 3.11 and < 3.14",
allow_module_level=True,
)

from timecopilot import TimeCopilotForecaster
from timecopilot.models.foundation.flowstate import FlowState
Expand Down
10 changes: 3 additions & 7 deletions tests/models/foundation/test_timesfm.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,13 +47,9 @@ def test_load_model_from_local_path(mocker, model_class, mock_paths):
expected_path = os.path.join(local_path, "torch_model.ckpt")
mock_loader[0].assert_called_once_with(path=expected_path)
elif model_class is _TimesFMV2_p5:
expected_predictor = mock_loader[
0
].return_value.model.load_checkpoint.return_value
assert predictor is expected_predictor
mock_loader[0].return_value.model.load_checkpoint.assert_called_once_with(
os.path.join(local_path, "model.safetensors")
)
# `from_pretrained` handles both local directories and HF repos.
assert predictor is mock_loader[0].from_pretrained.return_value
mock_loader[0].from_pretrained.assert_called_once_with(local_path)


@pytest.mark.parametrize("model_class, mock_paths", MODEL_PARAMS)
Expand Down
20 changes: 20 additions & 0 deletions tests/models/test_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,26 @@ def test_t0_import_fails():
assert "requires Python >= 3.11 and < 3.14" in str(excinfo.value)


@pytest.mark.skipif(
(3, 11) <= sys.version_info < (3, 14),
reason="FlowState requires Python >= 3.11 and < 3.14",
)
def test_flowstate_import_fails():
with pytest.raises(ImportError) as excinfo:
from timecopilot.models.foundation.flowstate import FlowState # noqa: F401
assert "requires Python >= 3.11 and < 3.14" in str(excinfo.value)


@pytest.mark.skipif(
(3, 11) <= sys.version_info < (3, 14),
reason="PatchTSTFM requires Python >= 3.11 and < 3.14",
)
def test_patchtst_fm_import_fails():
with pytest.raises(ImportError) as excinfo:
from timecopilot.models.foundation.patchtst_fm import PatchTSTFM # noqa: F401
assert "requires Python >= 3.11 and < 3.14" in str(excinfo.value)


@pytest.mark.skipif(
sys.version_info < (3, 13),
reason="Sundial requires Python < 3.13",
Expand Down
6 changes: 5 additions & 1 deletion timecopilot/models/foundation/flowstate.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,9 @@
import sys
from contextlib import contextmanager

if sys.version_info < (3, 11) or sys.version_info >= (3, 14):
raise ImportError("FlowState requires Python >= 3.11 and < 3.14")

import numpy as np
import pandas as pd
import torch
Expand Down Expand Up @@ -121,7 +125,7 @@ def _predict_batch(
prediction_length=h,
scale_factor=scale_factor,
batch_first=False,
).prediction_outputs
).quantile_outputs
fcst = fcst.squeeze(-1).transpose(-1, -2) # now shape is (batch, h, quantiles)
fcst_mean = fcst[..., supported_quantiles.index(0.5)]
fcst_mean_np = fcst_mean.detach().numpy(force=True)
Expand Down
6 changes: 5 additions & 1 deletion timecopilot/models/foundation/patchtst_fm.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,9 @@
import sys
from contextlib import contextmanager

if sys.version_info < (3, 11) or sys.version_info >= (3, 14):
raise ImportError("PatchTSTFM requires Python >= 3.11 and < 3.14")

import numpy as np
import pandas as pd
import torch
Expand Down Expand Up @@ -128,7 +132,7 @@ def _predict_batch(
quantile_levels=quantile_levels,
# scale_factor=scale_factor,
# batch_first=False,
).quantile_predictions
).quantile_outputs
fcst = fcst.squeeze(-1).transpose(-1, -2) # now shape is (batch, h, quantiles)

# may not be the ideal solution, but this should be more adaptable
Expand Down
10 changes: 4 additions & 6 deletions timecopilot/models/foundation/timesfm.py
Original file line number Diff line number Diff line change
Expand Up @@ -142,12 +142,10 @@ def _get_predictor(
self,
prediction_length: int,
) -> TimesFM_2p5_200M_torch:
# automatically detect the best device
# https://github.com/TimeCopilot/timesfm/blob/b810bbdf9f8a1e66396e7bd5cdb3b005e9116d86/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py#L71
if os.path.exists(self.repo_id):
path = os.path.join(self.repo_id, "model.safetensors")
tfm = TimesFM_2p5_200M_torch().model.load_checkpoint(path)
elif repo_exists(self.repo_id):
# `from_pretrained` handles both a local directory containing
# `model.safetensors` and a Hugging Face repo id, and the model picks
# the best available device on load.
if os.path.exists(self.repo_id) or repo_exists(self.repo_id):
tfm = TimesFM_2p5_200M_torch.from_pretrained(self.repo_id)
else:
raise OSError(
Expand Down
Loading
Loading