tirex-mirror 2025.12.2__tar.gz → 2025.12.16__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {tirex_mirror-2025.12.2/src/tirex_mirror.egg-info → tirex_mirror-2025.12.16}/PKG-INFO +6 -3
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/pyproject.toml +6 -5
- tirex_mirror-2025.12.16/src/tirex/models/base/base_classifier.py +36 -0
- tirex_mirror-2025.12.16/src/tirex/models/base/base_regressor.py +23 -0
- tirex_mirror-2025.12.2/src/tirex/models/classification/heads/base_classifier.py → tirex_mirror-2025.12.16/src/tirex/models/base/base_tirex.py +42 -26
- tirex_mirror-2025.12.16/src/tirex/models/classification/__init__.py +8 -0
- {tirex_mirror-2025.12.2/src/tirex/models/classification/heads → tirex_mirror-2025.12.16/src/tirex/models/classification}/gbm_classifier.py +4 -26
- {tirex_mirror-2025.12.2/src/tirex/models/classification/heads → tirex_mirror-2025.12.16/src/tirex/models/classification}/linear_classifier.py +11 -12
- {tirex_mirror-2025.12.2/src/tirex/models/classification/heads → tirex_mirror-2025.12.16/src/tirex/models/classification}/rf_classifier.py +3 -3
- {tirex_mirror-2025.12.2/src/tirex/models/classification → tirex_mirror-2025.12.16/src/tirex/models}/embedding.py +2 -3
- tirex_mirror-2025.12.16/src/tirex/models/regression/__init__.py +8 -0
- tirex_mirror-2025.12.16/src/tirex/models/regression/gbm_regressor.py +181 -0
- tirex_mirror-2025.12.16/src/tirex/models/regression/linear_regressor.py +250 -0
- tirex_mirror-2025.12.16/src/tirex/models/regression/rf_regressor.py +130 -0
- {tirex_mirror-2025.12.2/src/tirex/models/classification → tirex_mirror-2025.12.16/src/tirex/models}/trainer.py +29 -10
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/src/tirex/util.py +82 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16/src/tirex_mirror.egg-info}/PKG-INFO +6 -3
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/src/tirex_mirror.egg-info/SOURCES.txt +15 -7
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/src/tirex_mirror.egg-info/requires.txt +5 -1
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/tests/test_embedding.py +1 -3
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/tests/test_gbm_classifier.py +35 -26
- tirex_mirror-2025.12.16/tests/test_gbm_regressor.py +175 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/tests/test_linear_classifier.py +71 -52
- tirex_mirror-2025.12.16/tests/test_linear_regressor.py +217 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/tests/test_rf_classifier.py +35 -27
- tirex_mirror-2025.12.16/tests/test_rf_regressor.py +163 -0
- tirex_mirror-2025.12.2/src/tirex/models/classification/__init__.py +0 -8
- tirex_mirror-2025.12.2/src/tirex/models/classification/utils.py +0 -84
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/LICENSE +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/LICENSE_MIRROR.txt +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/MANIFEST.in +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/NOTICE.txt +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/README.md +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/setup.cfg +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/src/tirex/__init__.py +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/src/tirex/api_adapter/__init__.py +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/src/tirex/api_adapter/forecast.py +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/src/tirex/api_adapter/gluon.py +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/src/tirex/api_adapter/hf_data.py +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/src/tirex/api_adapter/standard_adapter.py +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/src/tirex/base.py +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/src/tirex/models/__init__.py +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/src/tirex/models/patcher.py +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/src/tirex/models/slstm/block.py +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/src/tirex/models/slstm/cell.py +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/src/tirex/models/slstm/layer.py +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/src/tirex/models/tirex.py +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/src/tirex_mirror.egg-info/dependency_links.txt +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/src/tirex_mirror.egg-info/top_level.txt +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/tests/test_chronos_zs.py +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/tests/test_compile.py +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/tests/test_forecast.py +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/tests/test_forecast_adapter.py +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/tests/test_load_model.py +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/tests/test_patcher.py +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/tests/test_slstm_torch_vs_cuda.py +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/tests/test_standard_adapter.py +0 -0
- {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.16}/tests/test_util_freq.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: tirex-mirror
|
|
3
|
-
Version: 2025.12.
|
|
3
|
+
Version: 2025.12.16
|
|
4
4
|
Summary: Unofficial mirror of NX-AI/tirex for packaging
|
|
5
5
|
Author-email: Arpad Rozsas <rozsasarpi@gmail.com>
|
|
6
6
|
License: NXAI COMMUNITY LICENSE AGREEMENT
|
|
@@ -55,7 +55,7 @@ License: NXAI COMMUNITY LICENSE AGREEMENT
|
|
|
55
55
|
|
|
56
56
|
Project-URL: Repository, https://github.com/rozsasarpi/tirex-mirror
|
|
57
57
|
Project-URL: Issues, https://github.com/rozsasarpi/tirex-mirror/issues
|
|
58
|
-
Keywords: TiRex,xLSTM,
|
|
58
|
+
Keywords: TiRex,xLSTM,Timeseries,Zero-shot,Deep Learning,Timeseries-Forecasting,Timeseries-Classification,Timeseries-Regression
|
|
59
59
|
Classifier: Programming Language :: Python :: 3
|
|
60
60
|
Classifier: Operating System :: OS Independent
|
|
61
61
|
Requires-Python: >=3.11
|
|
@@ -66,6 +66,7 @@ License-File: NOTICE.txt
|
|
|
66
66
|
Requires-Dist: torch
|
|
67
67
|
Requires-Dist: huggingface-hub
|
|
68
68
|
Requires-Dist: numpy
|
|
69
|
+
Requires-Dist: scikit-learn
|
|
69
70
|
Provides-Extra: cuda
|
|
70
71
|
Requires-Dist: xlstm; extra == "cuda"
|
|
71
72
|
Requires-Dist: ninja; extra == "cuda"
|
|
@@ -84,9 +85,11 @@ Requires-Dist: datasets; extra == "hfdataset"
|
|
|
84
85
|
Provides-Extra: test
|
|
85
86
|
Requires-Dist: fev>=0.6.0; extra == "test"
|
|
86
87
|
Requires-Dist: pytest; extra == "test"
|
|
88
|
+
Requires-Dist: aeon; extra == "test"
|
|
87
89
|
Provides-Extra: classification
|
|
88
|
-
Requires-Dist: scikit-learn; extra == "classification"
|
|
89
90
|
Requires-Dist: lightgbm[scikit-learn]; extra == "classification"
|
|
91
|
+
Provides-Extra: regression
|
|
92
|
+
Requires-Dist: lightgbm[scikit-learn]; extra == "regression"
|
|
90
93
|
Provides-Extra: all
|
|
91
94
|
Requires-Dist: xlstm; extra == "all"
|
|
92
95
|
Requires-Dist: ninja; extra == "all"
|
|
@@ -1,12 +1,12 @@
|
|
|
1
1
|
[project]
|
|
2
2
|
name = "tirex-mirror"
|
|
3
|
-
version = "2025.12.
|
|
3
|
+
version = "2025.12.16"
|
|
4
4
|
description = "Unofficial mirror of NX-AI/tirex for packaging"
|
|
5
5
|
readme = "README.md"
|
|
6
6
|
requires-python = ">=3.11"
|
|
7
7
|
classifiers = [ "Programming Language :: Python :: 3", "Operating System :: OS Independent",]
|
|
8
|
-
keywords = [ "TiRex", "xLSTM", "
|
|
9
|
-
dependencies = [ "torch", "huggingface-hub", "numpy",]
|
|
8
|
+
keywords = [ "TiRex", "xLSTM", "Timeseries", "Zero-shot", "Deep Learning", "Timeseries-Forecasting", "Timeseries-Classification", "Timeseries-Regression",]
|
|
9
|
+
dependencies = [ "torch", "huggingface-hub", "numpy", "scikit-learn",]
|
|
10
10
|
[[project.authors]]
|
|
11
11
|
name = "Arpad Rozsas"
|
|
12
12
|
email = "rozsasarpi@gmail.com"
|
|
@@ -28,8 +28,9 @@ notebooks = [ "ipykernel", "matplotlib", "pandas", "python-dotenv",]
|
|
|
28
28
|
plotting = [ "matplotlib",]
|
|
29
29
|
gluonts = [ "gluonts", "pandas",]
|
|
30
30
|
hfdataset = [ "datasets",]
|
|
31
|
-
test = [ "fev>=0.6.0", "pytest",]
|
|
32
|
-
classification = [ "
|
|
31
|
+
test = [ "fev>=0.6.0", "pytest", "aeon",]
|
|
32
|
+
classification = [ "lightgbm[scikit-learn]",]
|
|
33
|
+
regression = [ "lightgbm[scikit-learn]",]
|
|
33
34
|
all = [ "xlstm", "ninja", "ipykernel", "matplotlib", "pandas", "python-dotenv", "gluonts", "datasets", "pytest", "fev>=0.6.0", "scikit-learn", "lightgbm[scikit-learn]",]
|
|
34
35
|
|
|
35
36
|
[tool.docformatter]
|
|
@@ -0,0 +1,36 @@
|
|
|
1
|
+
# Copyright (c) NXAI GmbH.
|
|
2
|
+
# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
|
|
3
|
+
|
|
4
|
+
import torch
|
|
5
|
+
|
|
6
|
+
from .base_tirex import BaseTirexEmbeddingModel
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class BaseTirexClassifier(BaseTirexEmbeddingModel):
|
|
10
|
+
"""Abstract base class for TiRex classification models."""
|
|
11
|
+
|
|
12
|
+
@torch.inference_mode()
|
|
13
|
+
def predict(self, x: torch.Tensor) -> torch.Tensor:
|
|
14
|
+
"""Predict class labels for input time series data.
|
|
15
|
+
|
|
16
|
+
Args:
|
|
17
|
+
x: Input time series data as torch.Tensor with shape
|
|
18
|
+
(batch_size, num_variates, seq_len).
|
|
19
|
+
Returns:
|
|
20
|
+
torch.Tensor: Predicted class labels with shape (batch_size,).
|
|
21
|
+
"""
|
|
22
|
+
emb = self._compute_embeddings(x)
|
|
23
|
+
return torch.from_numpy(self.head.predict(emb)).long()
|
|
24
|
+
|
|
25
|
+
@torch.inference_mode()
|
|
26
|
+
def predict_proba(self, x: torch.Tensor) -> torch.Tensor:
|
|
27
|
+
"""Predict class probabilities for input time series data.
|
|
28
|
+
|
|
29
|
+
Args:
|
|
30
|
+
x: Input time series data as torch.Tensor with shape
|
|
31
|
+
(batch_size, num_variates, seq_len).
|
|
32
|
+
Returns:
|
|
33
|
+
torch.Tensor: Class probabilities with shape (batch_size, num_classes).
|
|
34
|
+
"""
|
|
35
|
+
emb = self._compute_embeddings(x)
|
|
36
|
+
return torch.from_numpy(self.head.predict_proba(emb))
|
|
@@ -0,0 +1,23 @@
|
|
|
1
|
+
# Copyright (c) NXAI GmbH.
|
|
2
|
+
# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
|
|
3
|
+
|
|
4
|
+
import torch
|
|
5
|
+
|
|
6
|
+
from .base_tirex import BaseTirexEmbeddingModel
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class BaseTirexRegressor(BaseTirexEmbeddingModel):
|
|
10
|
+
"""Abstract base class for TiRex regression models."""
|
|
11
|
+
|
|
12
|
+
@torch.inference_mode()
|
|
13
|
+
def predict(self, x: torch.Tensor) -> torch.Tensor:
|
|
14
|
+
"""Predict values for input time series data.
|
|
15
|
+
|
|
16
|
+
Args:
|
|
17
|
+
x: Input time series data as torch.Tensor with shape
|
|
18
|
+
(batch_size, num_variates, seq_len).
|
|
19
|
+
Returns:
|
|
20
|
+
torch.Tensor: Predicted values with shape (batch_size,).
|
|
21
|
+
"""
|
|
22
|
+
emb = self._compute_embeddings(x)
|
|
23
|
+
return torch.from_numpy(self.head.predict(emb)).float()
|
|
@@ -3,15 +3,18 @@
|
|
|
3
3
|
|
|
4
4
|
from abc import ABC, abstractmethod
|
|
5
5
|
|
|
6
|
+
import numpy as np
|
|
6
7
|
import torch
|
|
7
8
|
|
|
9
|
+
from tirex.util import train_val_split
|
|
10
|
+
|
|
8
11
|
from ..embedding import TiRexEmbedding
|
|
9
12
|
|
|
10
13
|
|
|
11
|
-
class
|
|
12
|
-
"""Abstract base class for TiRex
|
|
14
|
+
class BaseTirexEmbeddingModel(ABC):
|
|
15
|
+
"""Abstract base class for TiRex models.
|
|
13
16
|
|
|
14
|
-
This base class provides common functionality for all TiRex
|
|
17
|
+
This base class provides common functionality for all TiRex classifier and regression models,
|
|
15
18
|
including embedding model initialization and a consistent interface.
|
|
16
19
|
|
|
17
20
|
"""
|
|
@@ -19,7 +22,10 @@ class BaseTirexClassifier(ABC):
|
|
|
19
22
|
def __init__(
|
|
20
23
|
self, data_augmentation: bool = False, device: str | None = None, compile: bool = False, batch_size: int = 512
|
|
21
24
|
) -> None:
|
|
22
|
-
"""Initializes a TiRex
|
|
25
|
+
"""Initializes a base TiRex model.
|
|
26
|
+
|
|
27
|
+
This base class initializes the embedding model and common configuration
|
|
28
|
+
used by both classification and regression models.
|
|
23
29
|
|
|
24
30
|
Args:
|
|
25
31
|
data_augmentation : bool
|
|
@@ -49,45 +55,55 @@ class BaseTirexClassifier(ABC):
|
|
|
49
55
|
|
|
50
56
|
@abstractmethod
|
|
51
57
|
def fit(self, train_data: tuple[torch.Tensor, torch.Tensor]) -> None:
|
|
58
|
+
"""Abstract method for model training"""
|
|
52
59
|
pass
|
|
53
60
|
|
|
54
|
-
|
|
55
|
-
|
|
56
|
-
"""Predict class labels for input time series data.
|
|
61
|
+
def _compute_embeddings(self, x: torch.Tensor) -> np.ndarray:
|
|
62
|
+
"""Compute embeddings for input time series data.
|
|
57
63
|
|
|
58
64
|
Args:
|
|
59
65
|
x: Input time series data as torch.Tensor with shape
|
|
60
66
|
(batch_size, num_variates, seq_len).
|
|
61
|
-
Returns:
|
|
62
|
-
torch.Tensor: Predicted class labels with shape (batch_size,).
|
|
63
|
-
"""
|
|
64
|
-
self.emb_model.eval()
|
|
65
|
-
x = x.to(self.device)
|
|
66
|
-
embeddings = self.emb_model(x).cpu().numpy()
|
|
67
|
-
return torch.from_numpy(self.head.predict(embeddings)).long()
|
|
68
|
-
|
|
69
|
-
@torch.inference_mode()
|
|
70
|
-
def predict_proba(self, x: torch.Tensor) -> torch.Tensor:
|
|
71
|
-
"""Predict class probabilities for input time series data.
|
|
72
67
|
|
|
73
|
-
Args:
|
|
74
|
-
x: Input time series data as torch.Tensor with shape
|
|
75
|
-
(batch_size, num_variates, seq_len).
|
|
76
68
|
Returns:
|
|
77
|
-
|
|
69
|
+
np.ndarray: Embeddings with shape (batch_size, embedding_dim).
|
|
78
70
|
"""
|
|
79
71
|
self.emb_model.eval()
|
|
80
72
|
x = x.to(self.device)
|
|
81
|
-
|
|
82
|
-
|
|
73
|
+
return self.emb_model(x).cpu().numpy()
|
|
74
|
+
|
|
75
|
+
def _create_train_val_datasets(
|
|
76
|
+
self,
|
|
77
|
+
train_data: tuple[torch.Tensor, torch.Tensor],
|
|
78
|
+
val_data: tuple[torch.Tensor, torch.Tensor] | None = None,
|
|
79
|
+
val_split_ratio: float = 0.2,
|
|
80
|
+
stratify: bool = False,
|
|
81
|
+
seed: int | None = None,
|
|
82
|
+
) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]:
|
|
83
|
+
if val_data is None:
|
|
84
|
+
train_data, val_data = train_val_split(
|
|
85
|
+
train_data=train_data, val_split_ratio=val_split_ratio, stratify=stratify, seed=seed
|
|
86
|
+
)
|
|
87
|
+
return train_data, val_data
|
|
83
88
|
|
|
84
89
|
@abstractmethod
|
|
85
90
|
def save_model(self, path: str) -> None:
|
|
86
|
-
"""
|
|
91
|
+
"""Save model to file.
|
|
92
|
+
|
|
93
|
+
Args:
|
|
94
|
+
path: File path where the model should be saved.
|
|
95
|
+
"""
|
|
87
96
|
pass
|
|
88
97
|
|
|
89
98
|
@classmethod
|
|
90
99
|
@abstractmethod
|
|
91
100
|
def load_model(cls, path: str):
|
|
92
|
-
"""
|
|
101
|
+
"""Load model from file.
|
|
102
|
+
|
|
103
|
+
Args:
|
|
104
|
+
path: File path to the saved model checkpoint.
|
|
105
|
+
|
|
106
|
+
Returns:
|
|
107
|
+
Instance of the model class with loaded weights and configuration.
|
|
108
|
+
"""
|
|
93
109
|
pass
|
|
@@ -0,0 +1,8 @@
|
|
|
1
|
+
# Copyright (c) NXAI GmbH.
|
|
2
|
+
# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
|
|
3
|
+
|
|
4
|
+
from .gbm_classifier import TirexGBMClassifier
|
|
5
|
+
from .linear_classifier import TirexLinearClassifier
|
|
6
|
+
from .rf_classifier import TirexRFClassifier
|
|
7
|
+
|
|
8
|
+
__all__ = ["TirexLinearClassifier", "TirexRFClassifier", "TirexGBMClassifier"]
|
|
@@ -5,8 +5,7 @@ import joblib
|
|
|
5
5
|
import torch
|
|
6
6
|
from lightgbm import LGBMClassifier, early_stopping
|
|
7
7
|
|
|
8
|
-
from ..
|
|
9
|
-
from .base_classifier import BaseTirexClassifier
|
|
8
|
+
from ..base.base_classifier import BaseTirexClassifier
|
|
10
9
|
|
|
11
10
|
|
|
12
11
|
class TirexGBMClassifier(BaseTirexClassifier):
|
|
@@ -118,11 +117,11 @@ class TirexGBMClassifier(BaseTirexClassifier):
|
|
|
118
117
|
seed=self.random_state,
|
|
119
118
|
)
|
|
120
119
|
|
|
121
|
-
self.emb_model.eval()
|
|
122
120
|
X_train = X_train.to(self.device)
|
|
123
121
|
X_val = X_val.to(self.device)
|
|
124
|
-
|
|
125
|
-
|
|
122
|
+
|
|
123
|
+
embeddings_train = self._compute_embeddings(X_train)
|
|
124
|
+
embeddings_val = self._compute_embeddings(X_val)
|
|
126
125
|
|
|
127
126
|
y_train = y_train.detach().cpu().numpy() if isinstance(y_train, torch.Tensor) else y_train
|
|
128
127
|
y_val = y_val.detach().cpu().numpy() if isinstance(y_val, torch.Tensor) else y_val
|
|
@@ -187,24 +186,3 @@ class TirexGBMClassifier(BaseTirexClassifier):
|
|
|
187
186
|
model.random_state = getattr(model.head, "random_state", None)
|
|
188
187
|
|
|
189
188
|
return model
|
|
190
|
-
|
|
191
|
-
def _create_train_val_datasets(
|
|
192
|
-
self,
|
|
193
|
-
train_data: tuple[torch.Tensor, torch.Tensor],
|
|
194
|
-
val_data: tuple[torch.Tensor, torch.Tensor] | None = None,
|
|
195
|
-
val_split_ratio: float = 0.2,
|
|
196
|
-
stratify: bool = True,
|
|
197
|
-
seed: int | None = None,
|
|
198
|
-
) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]:
|
|
199
|
-
X_train, y_train = train_data
|
|
200
|
-
|
|
201
|
-
if val_data is None:
|
|
202
|
-
train_data, val_data = train_val_split(
|
|
203
|
-
train_data=train_data, val_split_ratio=val_split_ratio, stratify=stratify, seed=seed
|
|
204
|
-
)
|
|
205
|
-
X_train, y_train = train_data
|
|
206
|
-
X_val, y_val = val_data
|
|
207
|
-
else:
|
|
208
|
-
X_val, y_val = val_data
|
|
209
|
-
|
|
210
|
-
return (X_train, y_train), (X_val, y_val)
|
|
@@ -5,11 +5,11 @@ from dataclasses import asdict
|
|
|
5
5
|
|
|
6
6
|
import torch
|
|
7
7
|
|
|
8
|
-
from ..
|
|
9
|
-
from
|
|
8
|
+
from ..base.base_classifier import BaseTirexClassifier
|
|
9
|
+
from ..trainer import TrainConfig, Trainer, TrainingMetrics
|
|
10
10
|
|
|
11
11
|
|
|
12
|
-
class
|
|
12
|
+
class TirexLinearClassifier(BaseTirexClassifier, torch.nn.Module):
|
|
13
13
|
"""
|
|
14
14
|
A PyTorch classifier that combines time series embeddings with a linear classification head.
|
|
15
15
|
|
|
@@ -19,10 +19,10 @@ class TirexClassifierTorch(BaseTirexClassifier, torch.nn.Module):
|
|
|
19
19
|
|
|
20
20
|
Example:
|
|
21
21
|
>>> import torch
|
|
22
|
-
>>> from tirex.models.classification import
|
|
22
|
+
>>> from tirex.models.classification import TirexLinearClassifier
|
|
23
23
|
>>>
|
|
24
24
|
>>> # Create model with TiRex embeddings
|
|
25
|
-
>>> model =
|
|
25
|
+
>>> model = TirexLinearClassifier(
|
|
26
26
|
... data_augmentation=True,
|
|
27
27
|
... max_epochs=2,
|
|
28
28
|
... lr=1e-4,
|
|
@@ -72,7 +72,7 @@ class TirexClassifierTorch(BaseTirexClassifier, torch.nn.Module):
|
|
|
72
72
|
compile: bool
|
|
73
73
|
Whether to compile the frozen embedding model. Default: False
|
|
74
74
|
max_epochs : int
|
|
75
|
-
Maximum number of training epochs. Default:
|
|
75
|
+
Maximum number of training epochs. Default: 10
|
|
76
76
|
lr : float
|
|
77
77
|
Learning rate for the optimizer. Default: 1e-4
|
|
78
78
|
weight_decay : float
|
|
@@ -115,6 +115,7 @@ class TirexClassifierTorch(BaseTirexClassifier, torch.nn.Module):
|
|
|
115
115
|
lr=lr,
|
|
116
116
|
weight_decay=weight_decay,
|
|
117
117
|
class_weights=class_weights,
|
|
118
|
+
task_type="classification",
|
|
118
119
|
batch_size=batch_size,
|
|
119
120
|
val_split_ratio=val_split_ratio,
|
|
120
121
|
stratify=stratify,
|
|
@@ -132,9 +133,7 @@ class TirexClassifierTorch(BaseTirexClassifier, torch.nn.Module):
|
|
|
132
133
|
|
|
133
134
|
@torch.inference_mode()
|
|
134
135
|
def _identify_head_dims(self, x: torch.Tensor, y: torch.Tensor) -> None:
|
|
135
|
-
self.
|
|
136
|
-
sample_emb = self.emb_model(x[:1])
|
|
137
|
-
self.emb_dim = sample_emb.shape[-1]
|
|
136
|
+
self.emb_dim = self._compute_embeddings(x[:1]).shape[-1]
|
|
138
137
|
self.num_classes = len(torch.unique(y))
|
|
139
138
|
|
|
140
139
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
@@ -155,7 +154,7 @@ class TirexClassifierTorch(BaseTirexClassifier, torch.nn.Module):
|
|
|
155
154
|
|
|
156
155
|
def fit(
|
|
157
156
|
self, train_data: tuple[torch.Tensor, torch.Tensor], val_data: tuple[torch.Tensor, torch.Tensor] | None = None
|
|
158
|
-
) ->
|
|
157
|
+
) -> TrainingMetrics:
|
|
159
158
|
"""Train the classification head on the provided data.
|
|
160
159
|
|
|
161
160
|
This method initializes the classification head based on the data dimensions,
|
|
@@ -231,7 +230,7 @@ class TirexClassifierTorch(BaseTirexClassifier, torch.nn.Module):
|
|
|
231
230
|
)
|
|
232
231
|
|
|
233
232
|
@classmethod
|
|
234
|
-
def load_model(cls, path: str) -> "
|
|
233
|
+
def load_model(cls, path: str) -> "TirexLinearClassifier":
|
|
235
234
|
"""Load a saved model from file.
|
|
236
235
|
|
|
237
236
|
This reconstructs the model architecture and loads the trained weights from
|
|
@@ -240,7 +239,7 @@ class TirexClassifierTorch(BaseTirexClassifier, torch.nn.Module):
|
|
|
240
239
|
Args:
|
|
241
240
|
path: File path to the saved model checkpoint.
|
|
242
241
|
Returns:
|
|
243
|
-
|
|
242
|
+
TirexLinearClassifier: The loaded model with trained weights, ready for inference.
|
|
244
243
|
"""
|
|
245
244
|
checkpoint = torch.load(path)
|
|
246
245
|
|
|
@@ -5,7 +5,7 @@ import joblib
|
|
|
5
5
|
import torch
|
|
6
6
|
from sklearn.ensemble import RandomForestClassifier
|
|
7
7
|
|
|
8
|
-
from .base_classifier import BaseTirexClassifier
|
|
8
|
+
from ..base.base_classifier import BaseTirexClassifier
|
|
9
9
|
|
|
10
10
|
|
|
11
11
|
class TirexRFClassifier(BaseTirexClassifier):
|
|
@@ -84,9 +84,9 @@ class TirexRFClassifier(BaseTirexClassifier):
|
|
|
84
84
|
if isinstance(y_train, torch.Tensor):
|
|
85
85
|
y_train = y_train.detach().cpu().numpy()
|
|
86
86
|
|
|
87
|
-
self.emb_model.eval()
|
|
88
87
|
X_train = X_train.to(self.device)
|
|
89
|
-
embeddings = self.
|
|
88
|
+
embeddings = self._compute_embeddings(X_train)
|
|
89
|
+
|
|
90
90
|
self.head.fit(embeddings, y_train)
|
|
91
91
|
|
|
92
92
|
def save_model(self, path: str) -> None:
|
|
@@ -6,8 +6,7 @@ import torch.nn as nn
|
|
|
6
6
|
import torch.nn.functional as F
|
|
7
7
|
|
|
8
8
|
from tirex import load_model
|
|
9
|
-
|
|
10
|
-
from .utils import nanmax, nanmin, nanstd
|
|
9
|
+
from tirex.util import nanmax, nanmin, nanstd
|
|
11
10
|
|
|
12
11
|
|
|
13
12
|
class TiRexEmbedding(nn.Module):
|
|
@@ -20,7 +19,7 @@ class TiRexEmbedding(nn.Module):
|
|
|
20
19
|
self.batch_size = batch_size
|
|
21
20
|
|
|
22
21
|
if device is None:
|
|
23
|
-
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
22
|
+
device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
|
24
23
|
self.device = device
|
|
25
24
|
self._compile = compile
|
|
26
25
|
|
|
@@ -0,0 +1,8 @@
|
|
|
1
|
+
# Copyright (c) NXAI GmbH.
|
|
2
|
+
# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
|
|
3
|
+
|
|
4
|
+
from .gbm_regressor import TirexGBMRegressor
|
|
5
|
+
from .linear_regressor import TirexLinearRegressor
|
|
6
|
+
from .rf_regressor import TirexRFRegressor
|
|
7
|
+
|
|
8
|
+
__all__ = ["TirexLinearRegressor", "TirexRFRegressor", "TirexGBMRegressor"]
|
|
@@ -0,0 +1,181 @@
|
|
|
1
|
+
# Copyright (c) NXAI GmbH.
|
|
2
|
+
# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
|
|
3
|
+
|
|
4
|
+
import joblib
|
|
5
|
+
import torch
|
|
6
|
+
from lightgbm import LGBMRegressor, early_stopping
|
|
7
|
+
|
|
8
|
+
from ..base.base_regressor import BaseTirexRegressor
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class TirexGBMRegressor(BaseTirexRegressor):
|
|
12
|
+
"""
|
|
13
|
+
A Gradient Boosting regressor that uses time series embeddings as features.
|
|
14
|
+
|
|
15
|
+
This regressor combines a pre-trained embedding model for feature extraction with a
|
|
16
|
+
Gradient Boosting regressor.
|
|
17
|
+
|
|
18
|
+
Example:
|
|
19
|
+
>>> import torch
|
|
20
|
+
>>> from tirex.models.regression import TirexGBMRegressor
|
|
21
|
+
>>>
|
|
22
|
+
>>> # Create model with custom LightGBM parameters
|
|
23
|
+
>>> model = TirexGBMRegressor(
|
|
24
|
+
... data_augmentation=True,
|
|
25
|
+
... n_estimators=50,
|
|
26
|
+
... random_state=42
|
|
27
|
+
... )
|
|
28
|
+
>>>
|
|
29
|
+
>>> # Prepare data (can use NumPy arrays or PyTorch tensors)
|
|
30
|
+
>>> X_train = torch.randn(100, 1, 128) # 100 samples, 1 number of variates, 128 sequence length
|
|
31
|
+
>>> y_train = torch.randn(100,) # target values
|
|
32
|
+
>>>
|
|
33
|
+
>>> # Train the model
|
|
34
|
+
>>> model.fit((X_train, y_train))
|
|
35
|
+
>>>
|
|
36
|
+
>>> # Make predictions
|
|
37
|
+
>>> X_test = torch.randn(20, 1, 128)
|
|
38
|
+
>>> predictions = model.predict(X_test)
|
|
39
|
+
"""
|
|
40
|
+
|
|
41
|
+
def __init__(
|
|
42
|
+
self,
|
|
43
|
+
data_augmentation: bool = False,
|
|
44
|
+
device: str | None = None,
|
|
45
|
+
compile: bool = False,
|
|
46
|
+
batch_size: int = 512,
|
|
47
|
+
early_stopping_rounds: int | None = 10,
|
|
48
|
+
min_delta: float = 0.0,
|
|
49
|
+
val_split_ratio: float = 0.2,
|
|
50
|
+
# LightGBM kwargs
|
|
51
|
+
**lgbm_kwargs,
|
|
52
|
+
) -> None:
|
|
53
|
+
"""Initializes Embedding Based Gradient Boosting Regression model.
|
|
54
|
+
|
|
55
|
+
Args:
|
|
56
|
+
data_augmentation : bool
|
|
57
|
+
Whether to use data_augmentation for embeddings (sample statistics and first-order differences of the original data). Default: False
|
|
58
|
+
device : str | None
|
|
59
|
+
Device to run the embedding model on. If None, uses CUDA if available, else CPU. Default: None
|
|
60
|
+
compile: bool
|
|
61
|
+
Whether to compile the frozen embedding model. Default: False
|
|
62
|
+
batch_size : int
|
|
63
|
+
Batch size for embedding calculations. Default: 512
|
|
64
|
+
early_stopping_rounds: int | None
|
|
65
|
+
Number of rounds without improvement of all metrics for Early Stopping. Default: 10
|
|
66
|
+
min_delta: float
|
|
67
|
+
Minimum improvement in score to keep training. Default 0.0
|
|
68
|
+
val_split_ratio : float
|
|
69
|
+
Proportion of training data to use for validation, if validation data are not provided. Default: 0.2
|
|
70
|
+
**lgbm_kwargs
|
|
71
|
+
Additional keyword arguments to pass to LightGBM's LGBMRegressor.
|
|
72
|
+
Common options include n_estimators, max_depth, learning_rate, random_state, etc.
|
|
73
|
+
"""
|
|
74
|
+
super().__init__(data_augmentation=data_augmentation, device=device, compile=compile, batch_size=batch_size)
|
|
75
|
+
|
|
76
|
+
# Early Stopping callback
|
|
77
|
+
self.early_stopping_rounds = early_stopping_rounds
|
|
78
|
+
self.min_delta = min_delta
|
|
79
|
+
|
|
80
|
+
# Data split parameters:
|
|
81
|
+
self.val_split_ratio = val_split_ratio
|
|
82
|
+
|
|
83
|
+
# Extract random_state for train_val_split if provided
|
|
84
|
+
self.random_state = lgbm_kwargs.get("random_state", None)
|
|
85
|
+
|
|
86
|
+
self.head = LGBMRegressor(**lgbm_kwargs)
|
|
87
|
+
|
|
88
|
+
@torch.inference_mode()
|
|
89
|
+
def fit(
|
|
90
|
+
self,
|
|
91
|
+
train_data: tuple[torch.Tensor, torch.Tensor],
|
|
92
|
+
val_data: tuple[torch.Tensor, torch.Tensor] | None = None,
|
|
93
|
+
) -> None:
|
|
94
|
+
"""Train the LightGBM regressor on embedded time series data.
|
|
95
|
+
|
|
96
|
+
This method generates embeddings for the training data using the embedding
|
|
97
|
+
model, then trains the LightGBM regressor on these embeddings.
|
|
98
|
+
|
|
99
|
+
Args:
|
|
100
|
+
train_data: Tuple of (X_train, y_train) where X_train is the input time
|
|
101
|
+
series data (torch.Tensor) and y_train is a torch.Tensor
|
|
102
|
+
of target values.
|
|
103
|
+
val_data: Optional tuple of (X_val, y_val) for validation where X_val is the input time
|
|
104
|
+
series data (torch.Tensor) and y_val is a torch.Tensor
|
|
105
|
+
of target values. If None, validation data will be automatically split from train_data (20% split).
|
|
106
|
+
"""
|
|
107
|
+
|
|
108
|
+
(X_train, y_train), (X_val, y_val) = self._create_train_val_datasets(
|
|
109
|
+
train_data=train_data,
|
|
110
|
+
val_data=val_data,
|
|
111
|
+
val_split_ratio=self.val_split_ratio,
|
|
112
|
+
seed=self.random_state,
|
|
113
|
+
)
|
|
114
|
+
|
|
115
|
+
X_train = X_train.to(self.device)
|
|
116
|
+
X_val = X_val.to(self.device)
|
|
117
|
+
|
|
118
|
+
embeddings_train = self._compute_embeddings(X_train)
|
|
119
|
+
embeddings_val = self._compute_embeddings(X_val)
|
|
120
|
+
|
|
121
|
+
y_train = y_train.detach().cpu().numpy() if isinstance(y_train, torch.Tensor) else y_train
|
|
122
|
+
y_val = y_val.detach().cpu().numpy() if isinstance(y_val, torch.Tensor) else y_val
|
|
123
|
+
|
|
124
|
+
self.head.fit(
|
|
125
|
+
embeddings_train,
|
|
126
|
+
y_train,
|
|
127
|
+
eval_set=[(embeddings_val, y_val)],
|
|
128
|
+
callbacks=[early_stopping(stopping_rounds=self.early_stopping_rounds, min_delta=self.min_delta)]
|
|
129
|
+
if self.early_stopping_rounds is not None
|
|
130
|
+
else None,
|
|
131
|
+
)
|
|
132
|
+
|
|
133
|
+
def save_model(self, path: str) -> None:
|
|
134
|
+
"""This method saves the trained LightGBM regressor head (joblib format) and embedding information.
|
|
135
|
+
|
|
136
|
+
Args:
|
|
137
|
+
path: File path where the model should be saved (e.g., 'model.joblib').
|
|
138
|
+
"""
|
|
139
|
+
|
|
140
|
+
payload = {
|
|
141
|
+
"data_augmentation": self.data_augmentation,
|
|
142
|
+
"compile": self._compile,
|
|
143
|
+
"batch_size": self.batch_size,
|
|
144
|
+
"early_stopping_rounds": self.early_stopping_rounds,
|
|
145
|
+
"min_delta": self.min_delta,
|
|
146
|
+
"val_split_ratio": self.val_split_ratio,
|
|
147
|
+
"head": self.head,
|
|
148
|
+
}
|
|
149
|
+
joblib.dump(payload, path)
|
|
150
|
+
|
|
151
|
+
@classmethod
|
|
152
|
+
def load_model(cls, path: str) -> "TirexGBMRegressor":
|
|
153
|
+
"""Load a saved model from file.
|
|
154
|
+
|
|
155
|
+
This reconstructs the model with the embedding configuration and loads
|
|
156
|
+
the trained LightGBM regressor from a checkpoint file created by save_model().
|
|
157
|
+
|
|
158
|
+
Args:
|
|
159
|
+
path: File path to the saved model checkpoint.
|
|
160
|
+
Returns:
|
|
161
|
+
TirexGBMRegressor: The loaded model with trained Gradient Boosting regressor, ready for inference.
|
|
162
|
+
"""
|
|
163
|
+
checkpoint = joblib.load(path)
|
|
164
|
+
|
|
165
|
+
# Create new instance with saved configuration
|
|
166
|
+
model = cls(
|
|
167
|
+
data_augmentation=checkpoint["data_augmentation"],
|
|
168
|
+
compile=checkpoint["compile"],
|
|
169
|
+
batch_size=checkpoint["batch_size"],
|
|
170
|
+
early_stopping_rounds=checkpoint["early_stopping_rounds"],
|
|
171
|
+
min_delta=checkpoint["min_delta"],
|
|
172
|
+
val_split_ratio=checkpoint["val_split_ratio"],
|
|
173
|
+
)
|
|
174
|
+
|
|
175
|
+
# Load the trained LightGBM head
|
|
176
|
+
model.head = checkpoint["head"]
|
|
177
|
+
|
|
178
|
+
# Extract random_state from the loaded head if available
|
|
179
|
+
model.random_state = getattr(model.head, "random_state", None)
|
|
180
|
+
|
|
181
|
+
return model
|