tirex-mirror 2025.10.3__tar.gz → 2025.10.7__tar.gz

{tirex_mirror-2025.10.3/src/tirex_mirror.egg-info → tirex_mirror-2025.10.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tirex-mirror
-Version: 2025.10.3
+Version: 2025.10.7
 Summary: Unofficial mirror of NX-AI/tirex for packaging
 Author-email: Arpad Rozsas <rozsasarpi@gmail.com>
 License: NXAI COMMUNITY LICENSE AGREEMENT

{tirex_mirror-2025.10.3 → tirex_mirror-2025.10.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tirex-mirror"
-version = "2025.10.03"
+version = "2025.10.07"
 description = "Unofficial mirror of NX-AI/tirex for packaging"
 readme = "README.md"
 requires-python = ">=3.11"

{tirex_mirror-2025.10.3 → tirex_mirror-2025.10.7}/src/tirex/api_adapter/forecast.py

@@ -2,15 +2,22 @@
 # This software may be used and distributed according to the terms of the NXAI Community License Agreement.
 
 from abc import ABC, abstractmethod
-from typing import Literal
+from functools import partial
+from math import ceil
+from typing import Literal, Optional
 
 import torch
 
+from tirex.util import frequency_resample
+
 from .standard_adapter import ContextType, get_batches
 
 DEF_TARGET_COLUMN = "target"
 DEF_META_COLUMNS = ("start", "item_id")
 
+# Allowed resampling strategies (extend as new strategies are implemented)
+RESAMPLE_STRATEGIES: list[str] = ["frequency"]
+
 
 def _format_output(
     quantiles: torch.Tensor,
@@ -33,6 +40,27 @@ def _format_output(
         raise ValueError(f"Invalid output type: {output_type}")
 
 
+def _pad_time_series_batch(
+    batch_series: list[torch.Tensor],
+    max_length: int,
+) -> torch.Tensor:
+    if not batch_series:
+        return torch.empty((0, max_length))
+
+    first = batch_series[0]
+    dtype = first.dtype if first.is_floating_point() else torch.float32
+    device = first.device
+
+    padded = torch.full((len(batch_series), max_length), float("nan"), dtype=dtype, device=device)
+
+    for idx, series in enumerate(batch_series):
+        series = series.to(padded.dtype)
+        series_len = series.shape[0]
+        padded[idx, max_length - series_len :] = series
+
+    return padded
+
+
 def _as_generator(batches, fc_func, quantile_levels, output_type, **predict_kwargs):
     for batch_ctx, batch_meta in batches:
         quantiles, mean = fc_func(batch_ctx, **predict_kwargs)
@@ -45,7 +73,105 @@ def _as_generator(batches, fc_func, quantile_levels, output_type, **predict_kwar
         )
 
 
-def _gen_forecast(fc_func, batches, output_type, quantile_levels, yield_per_batch, **predict_kwargs):
+def _call_fc_with_padding(fc_func, batch_series: list[torch.Tensor], **predict_kwargs):
+    if not batch_series:
+        raise ValueError("Received empty batch for forecasting")
+
+    max_len = max(series.shape[0] for series in batch_series)
+    padded_ts = _pad_time_series_batch(batch_series, max_len)
+
+    return fc_func(padded_ts, **predict_kwargs)
+
+
+def _resample_fc_func_wrapper(
+    fc_func,
+    batch,
+    resample_strategy: str,
+    max_context: int = 2016,
+    **predict_kwargs,
+):
+    # downsample the time series based on the dominant frequencies, if enabled
+    max_period = (max_context // 1000) * 500
+    prediction_length = predict_kwargs.get("prediction_length", 100)
+    batch_resampled_ts: list[torch.Tensor] = []
+    fc_resample_fns = []
+    scaling_factors = []
+
+    # select the function doing the resampling
+    ctx_resample_fn = lambda x: (x, 1.0, (lambda y: y))
+    match resample_strategy:
+        case "frequency":
+            ctx_resample_fn = frequency_resample
+        case _:
+            raise RuntimeError("This shouldn't happen.")
+
+    for series in batch:
+        resampled_ts, _sample_factor, fc_resample_fn = ctx_resample_fn(
+            series,
+            prediction_length=prediction_length,
+            max_period=max_period,
+        )
+
+        batch_resampled_ts.append(resampled_ts)
+        fc_resample_fns.append(fc_resample_fn)
+        scaling_factors.append(_sample_factor)
+
+    # Compute per-item required horizons (in downsampled domain)
+    per_item_pred_lens = [int(ceil(prediction_length * sf)) for sf in scaling_factors]
+    max_pred_len = max(per_item_pred_lens) if per_item_pred_lens else int(prediction_length)
+    predict_kwargs.update(prediction_length=max_pred_len)
+
+    max_ts_length = max(ts.shape[0] for ts in batch_resampled_ts)
+    padded_ts = _pad_time_series_batch(batch_resampled_ts, max_ts_length)
+    print(f"Average sample batch factor: {sum(scaling_factors) / len(scaling_factors)}")
+
+    # generate prediction
+    fc_quantiles, fc_mean = fc_func(padded_ts, **predict_kwargs)
+
+    batch_prediction_q = []
+    batch_prediction_m = []
+    for el_q, el_m, fc_resample_fn, item_pred_len in zip(fc_quantiles, fc_mean, fc_resample_fns, per_item_pred_lens):
+        # truncate the forecasts to their individual sample factor adjusted prediction lengths
+        el_q = el_q[:item_pred_len, ...]  # [T, Q]
+        el_m = el_m[:item_pred_len]  # [T]
+
+        # upsample prediction
+        quantiles = fc_resample_fn(el_q.squeeze(0).transpose(0, 1)).transpose(0, 1)  # [T, Q]
+        mean = fc_resample_fn(el_m.squeeze(0))
+
+        quantiles = quantiles[:prediction_length, ...]
+        mean = mean[:prediction_length]
+
+        batch_prediction_q.append(quantiles)
+        batch_prediction_m.append(mean)
+
+    return torch.stack(batch_prediction_q, dim=0), torch.stack(batch_prediction_m, dim=0)
+
+
+def _gen_forecast(
+    fc_func,
+    batches,
+    output_type,
+    quantile_levels,
+    yield_per_batch,
+    resample_strategy: str | None = None,
+    max_context: int = 2016,
+    **predict_kwargs,
+):
+    base_fc_func = fc_func
+
+    if resample_strategy is not None:
+        if resample_strategy not in RESAMPLE_STRATEGIES:
+            raise ValueError(f"Invalid resample strategy: {resample_strategy}. Allowed: {RESAMPLE_STRATEGIES}")
+        fc_func = partial(
+            _resample_fc_func_wrapper,
+            base_fc_func,
+            resample_strategy=resample_strategy,
+            max_context=max_context,
+        )
+    else:
+        fc_func = partial(_call_fc_with_padding, base_fc_func)
+
     if yield_per_batch:
         return _as_generator(batches, fc_func, quantile_levels, output_type, **predict_kwargs)
 
@@ -92,6 +218,9 @@ def _common_forecast_doc():
                                               forecasts batch by batch as they are computed.
                                               Defaults to `False`.
 
+            resample_strategy (Optional[str], optional): Choose a resampling strategy. Allowed values: {RESAMPLE_STRATEGIES}.
+                                                If `None`, no resampling is applied. Currently only "frequency" is supported.
+
             **predict_kwargs: Additional keyword arguments that are passed directly to the underlying
                               prediction mechanism of the pre-trained model. Refer to the model's
                               internal prediction method documentation for available options.
@@ -113,6 +242,11 @@ class ForecastModel(ABC):
     def _forecast_quantiles(self, batch, **predict_kwargs):
         pass
 
+    @property
+    def max_context_length(self) -> int:
+        # retrieve the max_context attribute of the model configuration if present
+        return getattr(getattr(self, "config", None), "max_context", 2016)
+
     def forecast(
         self,
         context: ContextType,
@@ -120,6 +254,7 @@ class ForecastModel(ABC):
         batch_size: int = 512,
         quantile_levels: list[float] = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9),
         yield_per_batch: bool = False,
+        resample_strategy: Literal["frequency"] | None = None,
         **predict_kwargs,
     ):
         f"""
@@ -134,7 +269,14 @@ class ForecastModel(ABC):
         assert batch_size >= 1, "Batch size must be >= 1"
         batches = get_batches(context, batch_size)
         return _gen_forecast(
-            self._forecast_quantiles, batches, output_type, quantile_levels, yield_per_batch, **predict_kwargs
+            self._forecast_quantiles,
+            batches,
+            output_type,
+            quantile_levels,
+            yield_per_batch,
+            resample_strategy=resample_strategy,
+            max_context=self.max_context_length,
+            **predict_kwargs,
         )
 
     def forecast_gluon(
@@ -144,6 +286,7 @@ class ForecastModel(ABC):
         batch_size: int = 512,
         quantile_levels: list[float] = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9),
         yield_per_batch: bool = False,
+        resample_strategy: Literal["frequency"] | None = None,
         data_kwargs: dict = {},
         **predict_kwargs,
     ):
@@ -165,7 +308,14 @@ class ForecastModel(ABC):
 
         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
+            self._forecast_quantiles,
+            batches,
+            output_type,
+            quantile_levels,
+            yield_per_batch,
+            resample_strategy=resample_strategy,
+            max_context=self.max_context_length,
+            **predict_kwargs,
         )
 
     def forecast_hfdata(
@@ -175,6 +325,7 @@ class ForecastModel(ABC):
         batch_size: int = 512,
         quantile_levels: list[float] = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9),
         yield_per_batch: bool = False,
+        resample_strategy: Literal["frequency"] | None = None,
         data_kwargs: dict = {},
         **predict_kwargs,
     ):
@@ -198,5 +349,12 @@ class ForecastModel(ABC):
 
         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
+            self._forecast_quantiles,
+            batches,
+            output_type,
+            quantile_levels,
+            yield_per_batch,
+            resample_strategy=resample_strategy,
+            max_context=self.max_context_length,
+            **predict_kwargs,
         )

{tirex_mirror-2025.10.3 → tirex_mirror-2025.10.7}/src/tirex/api_adapter/gluon.py

@@ -7,7 +7,7 @@ from gluonts.dataset.common import Dataset
 from gluonts.dataset.field_names import FieldName
 from gluonts.model.forecast import QuantileForecast
 
-from .standard_adapter import _batch_pad_iterable
+from .standard_adapter import _batch_iterable
 
 DEF_TARGET_COLUMN = FieldName.TARGET  # target
 DEF_META_COLUMNS = (FieldName.START, FieldName.ITEM_ID)
@@ -27,7 +27,7 @@ def _get_gluon_ts_map(**gluon_kwargs):
 
 
 def get_gluon_batches(gluonDataset: Dataset, batch_size: int, **gluon_kwargs):
-    return _batch_pad_iterable(map(_get_gluon_ts_map(**gluon_kwargs), gluonDataset), batch_size)
+    return _batch_iterable(map(_get_gluon_ts_map(**gluon_kwargs), gluonDataset), batch_size)
 
 
 def format_gluonts_output(quantile_forecasts: torch.Tensor, mean_forecasts, meta: list[dict], quantile_levels):

{tirex_mirror-2025.10.3 → tirex_mirror-2025.10.7}/src/tirex/api_adapter/hf_data.py

@@ -4,7 +4,7 @@
 import datasets
 import torch
 
-from .standard_adapter import _batch_pad_iterable
+from .standard_adapter import _batch_iterable
 
 DEF_TARGET_COLUMN = "target"
 
@@ -35,4 +35,4 @@ def _get_hf_map(dataset: datasets.Dataset, **hf_kwargs):
 
 def get_hfdata_batches(hf_dataset: datasets.Dataset, batch_size: int, **hf_kwargs):
     dataset, map_func = _get_hf_map(hf_dataset, **hf_kwargs)
-    return _batch_pad_iterable(map(map_func, dataset), batch_size)
+    return _batch_iterable(map(map_func, dataset), batch_size)

tirex_mirror-2025.10.7/src/tirex/api_adapter/standard_adapter.py

@@ -0,0 +1,90 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+import itertools
+from collections.abc import Iterable, Iterator, Sequence
+from typing import Union
+
+import numpy as np
+import torch
+
+ContextType = Union[
+    torch.Tensor,
+    np.ndarray,
+    list[torch.Tensor],
+    list[np.ndarray],
+]
+
+
+def _ensure_1d_tensor(sample) -> torch.Tensor:
+    if isinstance(sample, torch.Tensor):
+        tensor = sample
+    else:
+        tensor = torch.as_tensor(sample)
+
+    if tensor.ndim > 1:
+        tensor = tensor.squeeze()
+
+    assert tensor.ndim == 1, "Each sample must be one-dimensional"
+    return tensor
+
+
+def _batched_slice(
+    full_batch,
+    full_meta: list[dict] | None,
+    batch_size: int,
+) -> Iterator[tuple[list[torch.Tensor], list[dict]]]:
+    total = len(full_batch)
+    for start in range(0, total, batch_size):
+        batch = full_batch[start : start + batch_size]
+        meta = full_meta[start : start + batch_size] if full_meta is not None else [{} for _ in range(len(batch))]
+
+        batch_series = []
+        for idx in range(len(batch)):
+            sample = batch[idx]
+            tensor = _ensure_1d_tensor(sample)
+            batch_series.append(tensor)
+
+        yield batch_series, meta
+
+
+def _batched(iterable: Iterable, n: int):
+    it = iter(iterable)
+    while batch := tuple(itertools.islice(it, n)):
+        yield batch
+
+
+def _batch_iterable(
+    iterable: Iterable[tuple[torch.Tensor, dict | None]],
+    batch_size: int,
+) -> Iterator[tuple[list[torch.Tensor], list[dict]]]:
+    for batch in _batched(iterable, batch_size):
+        series_list: list[torch.Tensor] = []
+        meta_list: list[dict] = []
+
+        for sample, meta in batch:
+            tensor = _ensure_1d_tensor(sample)
+            assert len(tensor) > 0, "Each sample needs to have a length > 0"
+            series_list.append(tensor)
+            meta_list.append(meta if meta is not None else {})
+
+        yield series_list, meta_list
+
+
+def get_batches(context: ContextType, batch_size: int):
+    batches = None
+    if isinstance(context, torch.Tensor):
+        if context.ndim == 1:
+            context = context.unsqueeze(0)
+        assert context.ndim == 2
+        batches = _batched_slice(context, None, batch_size)
+    elif isinstance(context, np.ndarray):
+        if context.ndim == 1:
+            context = np.expand_dims(context, axis=0)
+        assert context.ndim == 2
+        batches = _batched_slice(context, None, batch_size)
+    elif isinstance(context, (list, Iterable)):
+        batches = _batch_iterable(map(lambda x: (torch.Tensor(x), None), context), batch_size)
+    if batches is None:
+        raise ValueError(f"Context type {type(context)} not supported! Supported Types: {ContextType}")
+    return batches

{tirex_mirror-2025.10.3 → tirex_mirror-2025.10.7}/src/tirex/base.py

@@ -12,6 +12,7 @@ from huggingface_hub import hf_hub_download
 from tirex.models.slstm.cell import sLSTMCellTorch
 
 T = TypeVar("T", bound="PretrainedModel")
+VERSION_DELIMITER = "-"
 
 
 def skip_cuda():
@@ -32,6 +33,17 @@ def parse_hf_repo_id(path):
     return "/".join(parts[0:2])
 
 
+def parse_model_string(model_string):
+    if VERSION_DELIMITER in model_string:
+        parts = model_string.split(VERSION_DELIMITER)
+        model_id, version = parts[0], parts[0]
+    else:
+        model_id = model_string
+        version = None
+
+    return model_id, version
+
+
 class PretrainedModel(ABC):
     REGISTRY: dict[str, "PretrainedModel"] = {}
 
@@ -105,7 +117,8 @@ def load_model(
         backend = "torch" if skip_cuda() or not xlstm_available() else "cuda"
 
     try:
-        _, model_id = parse_hf_repo_id(path).split("/")
+        _, model_string = parse_hf_repo_id(path).split("/")
+        model_id, version = parse_model_string(model_string)
     except:
         raise ValueError(f"Invalid model path {path}")
     model_cls = PretrainedModel.REGISTRY.get(model_id, None)