tirex-mirror 2025.12.2__tar.gz → 2025.12.23__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.
Files changed (58) hide show
  1. {tirex_mirror-2025.12.2/src/tirex_mirror.egg-info → tirex_mirror-2025.12.23}/PKG-INFO +6 -3
  2. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/pyproject.toml +6 -5
  3. tirex_mirror-2025.12.23/src/tirex/models/base/base_classifier.py +36 -0
  4. tirex_mirror-2025.12.23/src/tirex/models/base/base_regressor.py +23 -0
  5. tirex_mirror-2025.12.2/src/tirex/models/classification/heads/base_classifier.py → tirex_mirror-2025.12.23/src/tirex/models/base/base_tirex.py +42 -26
  6. tirex_mirror-2025.12.23/src/tirex/models/classification/__init__.py +8 -0
  7. {tirex_mirror-2025.12.2/src/tirex/models/classification/heads → tirex_mirror-2025.12.23/src/tirex/models/classification}/gbm_classifier.py +4 -26
  8. {tirex_mirror-2025.12.2/src/tirex/models/classification/heads → tirex_mirror-2025.12.23/src/tirex/models/classification}/linear_classifier.py +11 -12
  9. {tirex_mirror-2025.12.2/src/tirex/models/classification/heads → tirex_mirror-2025.12.23/src/tirex/models/classification}/rf_classifier.py +3 -3
  10. {tirex_mirror-2025.12.2/src/tirex/models/classification → tirex_mirror-2025.12.23/src/tirex/models}/embedding.py +2 -3
  11. tirex_mirror-2025.12.23/src/tirex/models/regression/__init__.py +8 -0
  12. tirex_mirror-2025.12.23/src/tirex/models/regression/gbm_regressor.py +181 -0
  13. tirex_mirror-2025.12.23/src/tirex/models/regression/linear_regressor.py +250 -0
  14. tirex_mirror-2025.12.23/src/tirex/models/regression/rf_regressor.py +130 -0
  15. {tirex_mirror-2025.12.2/src/tirex/models/classification → tirex_mirror-2025.12.23/src/tirex/models}/trainer.py +29 -10
  16. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/src/tirex/util.py +82 -0
  17. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23/src/tirex_mirror.egg-info}/PKG-INFO +6 -3
  18. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/src/tirex_mirror.egg-info/SOURCES.txt +15 -7
  19. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/src/tirex_mirror.egg-info/requires.txt +5 -1
  20. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/tests/test_embedding.py +1 -3
  21. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/tests/test_gbm_classifier.py +35 -26
  22. tirex_mirror-2025.12.23/tests/test_gbm_regressor.py +175 -0
  23. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/tests/test_linear_classifier.py +71 -52
  24. tirex_mirror-2025.12.23/tests/test_linear_regressor.py +217 -0
  25. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/tests/test_rf_classifier.py +35 -27
  26. tirex_mirror-2025.12.23/tests/test_rf_regressor.py +163 -0
  27. tirex_mirror-2025.12.2/src/tirex/models/classification/__init__.py +0 -8
  28. tirex_mirror-2025.12.2/src/tirex/models/classification/utils.py +0 -84
  29. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/LICENSE +0 -0
  30. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/LICENSE_MIRROR.txt +0 -0
  31. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/MANIFEST.in +0 -0
  32. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/NOTICE.txt +0 -0
  33. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/README.md +0 -0
  34. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/setup.cfg +0 -0
  35. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/src/tirex/__init__.py +0 -0
  36. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/src/tirex/api_adapter/__init__.py +0 -0
  37. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/src/tirex/api_adapter/forecast.py +0 -0
  38. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/src/tirex/api_adapter/gluon.py +0 -0
  39. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/src/tirex/api_adapter/hf_data.py +0 -0
  40. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/src/tirex/api_adapter/standard_adapter.py +0 -0
  41. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/src/tirex/base.py +0 -0
  42. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/src/tirex/models/__init__.py +0 -0
  43. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/src/tirex/models/patcher.py +0 -0
  44. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/src/tirex/models/slstm/block.py +0 -0
  45. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/src/tirex/models/slstm/cell.py +0 -0
  46. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/src/tirex/models/slstm/layer.py +0 -0
  47. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/src/tirex/models/tirex.py +0 -0
  48. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/src/tirex_mirror.egg-info/dependency_links.txt +0 -0
  49. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/src/tirex_mirror.egg-info/top_level.txt +0 -0
  50. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/tests/test_chronos_zs.py +0 -0
  51. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/tests/test_compile.py +0 -0
  52. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/tests/test_forecast.py +0 -0
  53. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/tests/test_forecast_adapter.py +0 -0
  54. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/tests/test_load_model.py +0 -0
  55. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/tests/test_patcher.py +0 -0
  56. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/tests/test_slstm_torch_vs_cuda.py +0 -0
  57. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/tests/test_standard_adapter.py +0 -0
  58. {tirex_mirror-2025.12.2 → tirex_mirror-2025.12.23}/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.2
3
+ Version: 2025.12.23
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,Time Series,Zero-shot,Deep Learning,Classification,Timeseries-Classification
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.02"
3
+ version = "2025.12.23"
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", "Time Series", "Zero-shot", "Deep Learning", "Classification", "Timeseries-Classification",]
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 = [ "scikit-learn", "lightgbm[scikit-learn]",]
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 BaseTirexClassifier(ABC):
12
- """Abstract base class for TiRex classification models.
14
+ class BaseTirexEmbeddingModel(ABC):
15
+ """Abstract base class for TiRex models.
13
16
 
14
- This base class provides common functionality for all TiRex classifiers,
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 classification model.
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
- @torch.inference_mode()
55
- def predict(self, x: torch.Tensor) -> torch.Tensor:
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
- torch.Tensor: Class probabilities with shape (batch_size, num_classes).
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
- embeddings = self.emb_model(x).cpu().numpy()
82
- return torch.from_numpy(self.head.predict_proba(embeddings))
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
- """Saving model abstract method"""
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
- """Loading model abstract method"""
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 ..utils import train_val_split
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
- embeddings_train = self.emb_model(X_train).cpu().numpy()
125
- embeddings_val = self.emb_model(X_val).cpu().numpy()
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 ..trainer import TrainConfig, Trainer
9
- from .base_classifier import BaseTirexClassifier
8
+ from ..base.base_classifier import BaseTirexClassifier
9
+ from ..trainer import TrainConfig, Trainer, TrainingMetrics
10
10
 
11
11
 
12
- class TirexClassifierTorch(BaseTirexClassifier, torch.nn.Module):
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 TirexClassifierTorch
22
+ >>> from tirex.models.classification import TirexLinearClassifier
23
23
  >>>
24
24
  >>> # Create model with TiRex embeddings
25
- >>> model = TirexClassifierTorch(
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: 50
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.emb_model.eval()
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
- ) -> dict[str, float]:
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) -> "TirexClassifierTorch":
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
- TirexClassifierTorch: The loaded model with trained weights, ready for inference.
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.emb_model(X_train).cpu().numpy()
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