tirex-mirror 2025.6.9.dev3__tar.gz → 2025.8.28__tar.gz

{tirex_mirror-2025.6.9.dev3/src/tirex_mirror.egg-info → tirex_mirror-2025.8.28}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tirex-mirror
-Version: 2025.6.9.dev3
+Version: 2025.8.28
 Summary: Unofficial mirror of NX-AI/tirex for packaging
 Author-email: Arpad Rozsas <rozsasarpi@gmail.com>
 License: NXAI COMMUNITY LICENSE AGREEMENT
@@ -80,10 +80,13 @@ Provides-Extra: gluonts
 Requires-Dist: gluonts; extra == "gluonts"
 Provides-Extra: hfdataset
 Requires-Dist: datasets; extra == "hfdataset"
+Provides-Extra: test
+Requires-Dist: fev; extra == "test"
 Provides-Extra: all
 Requires-Dist: ipykernel; extra == "all"
 Requires-Dist: gluonts; extra == "all"
 Requires-Dist: datasets; extra == "all"
+Requires-Dist: fev; extra == "all"
 Dynamic: license-file
 
 # tirex-mirror

{tirex_mirror-2025.6.9.dev3 → tirex_mirror-2025.8.28}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tirex-mirror"
-version = "2025.06.09dev3"
+version = "2025.08.28"
 description = "Unofficial mirror of NX-AI/tirex for packaging"
 readme = "README.md"
 requires-python = ">=3.11"
@@ -26,7 +26,8 @@ Issues = "https://github.com/rozsasarpi/tirex-mirror/issues"
 notebooks = [ "ipykernel",]
 gluonts = [ "gluonts",]
 hfdataset = [ "datasets",]
-all = [ "ipykernel", "gluonts", "datasets",]
+test = [ "fev",]
+all = [ "ipykernel", "gluonts", "datasets", "fev",]
 
 [tool.docformatter]
 diff = false
@@ -43,6 +44,13 @@ exclude = [ ".eggs", ".git", ".ipynb_checkpoints", ".pytest_cache", ".ruff_cache
 line-length = 120
 target-version = "py311"
 
+[tool.pytest.ini_options]
+pythonpath = "."
+log_cli = true
+log_cli_level = "INFO"
+log_format = "%(asctime)s %(levelname)s %(message)s"
+log_date_format = "%Y-%m-%d %H:%M:%S"
+
 [tool.ruff.format]
 quote-style = "double"
 

{tirex_mirror-2025.6.9.dev3 → tirex_mirror-2025.8.28}/src/tirex/api_adapter/forecast.py

@@ -8,20 +8,6 @@ import torch
 
 from .standard_adapter import ContextType, get_batches
 
-try:
-    from .gluon import format_gluonts_output, get_gluon_batches
-
-    _GLUONTS_AVAILABLE = True
-except ImportError:
-    _GLUONTS_AVAILABLE = False
-
-try:
-    from .hf_data import get_hfdata_batches
-
-    _HF_DATASETS_AVAILABLE = True
-except ImportError:
-    _HF_DATASETS_AVAILABLE = False
-
 
 DEF_TARGET_COLUMN = "target"
 DEF_META_COLUMNS = ("start", "item_id")
@@ -39,7 +25,9 @@ def _format_output(
     elif output_type == "numpy":
         return quantiles.cpu().numpy(), means.cpu().numpy()
     elif output_type == "gluonts":
-        if not _GLUONTS_AVAILABLE:
+        try:
+            from .gluon import format_gluonts_output
+        except ImportError:
             raise ValueError("output_type glutonts needs GluonTs but GluonTS is not available (not installed)!")
         return format_gluonts_output(quantiles, means, sample_meta, quantile_levels)
     else:
@@ -171,8 +159,11 @@ class ForecastModel(ABC):
                                           autogluon data processing function.
         """
         assert batch_size >= 1, "Batch size must be >= 1"
-        if not _GLUONTS_AVAILABLE:
+        try:
+            from .gluon import get_gluon_batches
+        except ImportError:
             raise ValueError("forecast_gluon glutonts needs GluonTs but GluonTS is not available (not installed)!")
+
         batches = get_gluon_batches(gluonDataset, batch_size, **data_kwargs)
         return _gen_forecast(
             self._forecast_quantiles, batches, output_type, quantile_levels, yield_per_batch, **predict_kwargs
@@ -199,10 +190,13 @@ class ForecastModel(ABC):
                                           datasets data processing function.
         """
         assert batch_size >= 1, "Batch size must be >= 1"
-        if not _HF_DATASETS_AVAILABLE:
+        try:
+            from .hf_data import get_hfdata_batches
+        except ImportError:
             raise ValueError(
                 "forecast_hfdata glutonts needs HuggingFace datasets but datasets is not available (not installed)!"
             )
+
         batches = get_hfdata_batches(hf_dataset, batch_size, **data_kwargs)
         return _gen_forecast(
             self._forecast_quantiles, batches, output_type, quantile_levels, yield_per_batch, **predict_kwargs

{tirex_mirror-2025.6.9.dev3 → tirex_mirror-2025.8.28/src/tirex_mirror.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tirex-mirror
-Version: 2025.6.9.dev3
+Version: 2025.8.28
 Summary: Unofficial mirror of NX-AI/tirex for packaging
 Author-email: Arpad Rozsas <rozsasarpi@gmail.com>
 License: NXAI COMMUNITY LICENSE AGREEMENT
@@ -80,10 +80,13 @@ Provides-Extra: gluonts
 Requires-Dist: gluonts; extra == "gluonts"
 Provides-Extra: hfdataset
 Requires-Dist: datasets; extra == "hfdataset"
+Provides-Extra: test
+Requires-Dist: fev; extra == "test"
 Provides-Extra: all
 Requires-Dist: ipykernel; extra == "all"
 Requires-Dist: gluonts; extra == "all"
 Requires-Dist: datasets; extra == "all"
+Requires-Dist: fev; extra == "all"
 Dynamic: license-file
 
 # tirex-mirror

{tirex_mirror-2025.6.9.dev3 → tirex_mirror-2025.8.28}/src/tirex_mirror.egg-info/SOURCES.txt

@@ -21,5 +21,7 @@ src/tirex_mirror.egg-info/SOURCES.txt
 src/tirex_mirror.egg-info/dependency_links.txt
 src/tirex_mirror.egg-info/requires.txt
 src/tirex_mirror.egg-info/top_level.txt
+tests/test_chronos_zs.py
+tests/test_forecast.py
 tests/test_forecast_adapter.py
 tests/test_standard_adapter.py

{tirex_mirror-2025.6.9.dev3 → tirex_mirror-2025.8.28}/src/tirex_mirror.egg-info/requires.txt

@@ -14,6 +14,7 @@ tqdm
 ipykernel
 gluonts
 datasets
+fev
 
 [gluonts]
 gluonts
@@ -23,3 +24,6 @@ datasets
 
 [notebooks]
 ipykernel
+
+[test]
+fev

tirex_mirror-2025.8.28/tests/test_chronos_zs.py

@@ -0,0 +1,85 @@
+from tirex import ForecastModel, load_model
+import time
+import datasets
+import fev
+import pytest
+import os
+import math
+
+def geometric_mean(s):
+    return math.prod(s) ** (1/len(s))
+
+def eval_task(model, task):
+    past_data, _ = task.get_input_data(trust_remote_code=True)
+    quantile_levels = task.quantile_levels
+    past_data = past_data.with_format("torch").cast_column(
+        task.target_column, datasets.Sequence(datasets.Value("float32"))
+    )[task.target_column]
+    loaded_data = [t for t in past_data]
+
+    start_time = time.monotonic()
+    quantiles, means = model.forecast(loaded_data, quantile_levels=quantile_levels, prediction_length=task.horizon)
+    inference_time = time.monotonic() - start_time
+    predictions_dict = {"predictions": means}
+    for idx, level in enumerate(quantile_levels):
+        predictions_dict[str(level)] = quantiles[:, :, idx]  # [num_items, horizon]
+
+    predictions = datasets.Dataset.from_dict(predictions_dict)
+    return predictions, inference_time
+
+
+@pytest.fixture
+def tirex_model() -> ForecastModel:
+    return load_model("NX-AI/TiRex")
+
+
+@pytest.fixture
+def benchmark():
+    url = "https://raw.githubusercontent.com/autogluon/fev/refs/heads/main/benchmarks/chronos_zeroshot/tasks.yaml"
+    return fev.Benchmark.from_yaml(url)
+
+
+def test_chronos_single(tirex_model, benchmark):
+    task_name = "monash_australian_electricity"
+    task = [task for task in benchmark.tasks if task.dataset_config == task_name][0]
+    predictions, inference_time = eval_task(tirex_model, task)
+    evaluation_summary = task.evaluation_summary(
+        predictions,
+        model_name="TiRex",
+        inference_time_s=inference_time,
+    )
+
+    assert evaluation_summary["WQL"] < 0.055, "WQL on the electricity task needs to be less than 0.055"
+    assert evaluation_summary["MASE"] < 0.99, "MASE on the electricity task needs to be less than 0.99"
+
+
+@pytest.mark.skipif(os.getenv("CI"), reason="Skip full chromos benchmarking in the CI")
+def test_chronos_all(tirex_model, benchmark):
+    tasks_wql = []
+    tasks_mase = []
+    for task in benchmark.tasks:
+        predictions, inference_time = eval_task(tirex_model, task)
+        evaluation_summary = task.evaluation_summary(
+            predictions,
+            model_name="TiRex",
+            inference_time_s=inference_time,
+        )
+        tasks_wql.append(evaluation_summary["WQL"])
+        tasks_mase.append(evaluation_summary["MASE"])
+    
+    # Calculated from the geometric mean of the WQL and MASE data of the seasonal_naive model
+    # https://github.com/autogluon/fev/blob/main/benchmarks/chronos_zeroshot/results/seasonal_naive.csv
+    agg_wql_baseline = 0.1460642781226389
+    agg_mase_baseline = 1.6708210897174531
+
+    agg_wql = geometric_mean(tasks_wql)
+    agg_mase = geometric_mean(tasks_mase)
+
+    print(f"WQL: {agg_wql / agg_wql_baseline:.3f}")
+    print(f"MASE: {agg_mase / agg_mase_baseline:.3f}")
+
+    tolerance = 0.01
+
+    # Values from Tirex paper: https://arxiv.org/pdf/2505.23719
+    assert agg_wql / agg_wql_baseline < 0.59 + tolerance, "WQL on chromos needs to be less than 0.60"
+    assert agg_mase / agg_mase_baseline < 0.78 + tolerance, "MASE on chromos needs to be less than 0.79"

tirex_mirror-2025.8.28/tests/test_forecast.py

@@ -0,0 +1,35 @@
+from tirex import ForecastModel, load_model
+import pytest
+from pathlib import Path
+import numpy as np
+import torch
+
+
+def load_tensor_from_file(path):
+    base_path = Path(__file__).parent.resolve() / "data"
+    return torch.from_numpy(np.genfromtxt(base_path / path, dtype=np.float32))
+
+
+@pytest.fixture
+def tirex_model() -> ForecastModel:
+    return load_model("NX-AI/TiRex")
+
+
+def test_forecast_air_traffic(tirex_model):
+    context = load_tensor_from_file("air_passengers.csv")[:-12]
+
+    quantiles, mean = tirex_model.forecast(context, prediction_length=24)
+
+    ref_data = load_tensor_from_file("air_passengers_forecast_ref.csv")
+    
+    assert torch.allclose(mean, ref_data), "Forecasted tensor has to match reference data."
+
+
+def test_forecast_seattle_5T(tirex_model):
+    context = load_tensor_from_file("loop_seattle_5T.csv")[:-512]
+
+    quantiles, mean = tirex_model.forecast(context, prediction_length=768)
+
+    ref_data = load_tensor_from_file("loop_seattle_5T_forecast_ref.csv")
+    
+    assert torch.allclose(mean, ref_data), "Forecasted tensor has to match reference data."