tirex-mirror 2025.11.29__tar.gz → 2025.12.16__tar.gz

{tirex_mirror-2025.11.29 → tirex_mirror-2025.12.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tirex-mirror
-Version: 2025.11.29
+Version: 2025.12.16
 Summary: Unofficial mirror of NX-AI/tirex for packaging
 Author-email: Arpad Rozsas <rozsasarpi@gmail.com>
 License: NXAI COMMUNITY LICENSE AGREEMENT
@@ -55,7 +55,7 @@ License: NXAI COMMUNITY LICENSE AGREEMENT
         
 Project-URL: Repository, https://github.com/rozsasarpi/tirex-mirror
 Project-URL: Issues, https://github.com/rozsasarpi/tirex-mirror/issues
-Keywords: TiRex,xLSTM,Time Series,Zero-shot,Deep Learning
+Keywords: TiRex,xLSTM,Timeseries,Zero-shot,Deep Learning,Timeseries-Forecasting,Timeseries-Classification,Timeseries-Regression
 Classifier: Programming Language :: Python :: 3
 Classifier: Operating System :: OS Independent
 Requires-Python: >=3.11
@@ -66,6 +66,7 @@ License-File: NOTICE.txt
 Requires-Dist: torch
 Requires-Dist: huggingface-hub
 Requires-Dist: numpy
+Requires-Dist: scikit-learn
 Provides-Extra: cuda
 Requires-Dist: xlstm; extra == "cuda"
 Requires-Dist: ninja; extra == "cuda"
@@ -84,8 +85,11 @@ Requires-Dist: datasets; extra == "hfdataset"
 Provides-Extra: test
 Requires-Dist: fev>=0.6.0; extra == "test"
 Requires-Dist: pytest; extra == "test"
+Requires-Dist: aeon; extra == "test"
 Provides-Extra: classification
-Requires-Dist: scikit-learn; extra == "classification"
+Requires-Dist: lightgbm[scikit-learn]; extra == "classification"
+Provides-Extra: regression
+Requires-Dist: lightgbm[scikit-learn]; extra == "regression"
 Provides-Extra: all
 Requires-Dist: xlstm; extra == "all"
 Requires-Dist: ninja; extra == "all"
@@ -98,6 +102,7 @@ Requires-Dist: datasets; extra == "all"
 Requires-Dist: pytest; extra == "all"
 Requires-Dist: fev>=0.6.0; extra == "all"
 Requires-Dist: scikit-learn; extra == "all"
+Requires-Dist: lightgbm[scikit-learn]; extra == "all"
 Dynamic: license-file
 
 # tirex-mirror

{tirex_mirror-2025.11.29 → tirex_mirror-2025.12.16}/pyproject.toml

@@ -1,12 +1,12 @@
 [project]
 name = "tirex-mirror"
-version = "2025.11.29"
+version = "2025.12.16"
 description = "Unofficial mirror of NX-AI/tirex for packaging"
 readme = "README.md"
 requires-python = ">=3.11"
 classifiers = [ "Programming Language :: Python :: 3", "Operating System :: OS Independent",]
-keywords = [ "TiRex", "xLSTM", "Time Series", "Zero-shot", "Deep Learning",]
-dependencies = [ "torch", "huggingface-hub", "numpy",]
+keywords = [ "TiRex", "xLSTM", "Timeseries", "Zero-shot", "Deep Learning", "Timeseries-Forecasting", "Timeseries-Classification", "Timeseries-Regression",]
+dependencies = [ "torch", "huggingface-hub", "numpy", "scikit-learn",]
 [[project.authors]]
 name = "Arpad Rozsas"
 email = "rozsasarpi@gmail.com"
@@ -28,9 +28,10 @@ notebooks = [ "ipykernel", "matplotlib", "pandas", "python-dotenv",]
 plotting = [ "matplotlib",]
 gluonts = [ "gluonts", "pandas",]
 hfdataset = [ "datasets",]
-test = [ "fev>=0.6.0", "pytest",]
-classification = [ "scikit-learn",]
-all = [ "xlstm", "ninja", "ipykernel", "matplotlib", "pandas", "python-dotenv", "gluonts", "datasets", "pytest", "fev>=0.6.0", "scikit-learn",]
+test = [ "fev>=0.6.0", "pytest", "aeon",]
+classification = [ "lightgbm[scikit-learn]",]
+regression = [ "lightgbm[scikit-learn]",]
+all = [ "xlstm", "ninja", "ipykernel", "matplotlib", "pandas", "python-dotenv", "gluonts", "datasets", "pytest", "fev>=0.6.0", "scikit-learn", "lightgbm[scikit-learn]",]
 
 [tool.docformatter]
 diff = false

tirex_mirror-2025.12.16/src/tirex/models/base/base_classifier.py

@@ -0,0 +1,36 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+import torch
+
+from .base_tirex import BaseTirexEmbeddingModel
+
+
+class BaseTirexClassifier(BaseTirexEmbeddingModel):
+    """Abstract base class for TiRex classification models."""
+
+    @torch.inference_mode()
+    def predict(self, x: torch.Tensor) -> torch.Tensor:
+        """Predict class labels for input time series data.
+
+        Args:
+            x: Input time series data as torch.Tensor with shape
+                (batch_size, num_variates, seq_len).
+        Returns:
+            torch.Tensor: Predicted class labels with shape (batch_size,).
+        """
+        emb = self._compute_embeddings(x)
+        return torch.from_numpy(self.head.predict(emb)).long()
+
+    @torch.inference_mode()
+    def predict_proba(self, x: torch.Tensor) -> torch.Tensor:
+        """Predict class probabilities for input time series data.
+
+        Args:
+            x: Input time series data as torch.Tensor with shape
+                (batch_size, num_variates, seq_len).
+        Returns:
+            torch.Tensor: Class probabilities with shape (batch_size, num_classes).
+        """
+        emb = self._compute_embeddings(x)
+        return torch.from_numpy(self.head.predict_proba(emb))

tirex_mirror-2025.12.16/src/tirex/models/base/base_regressor.py

@@ -0,0 +1,23 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+import torch
+
+from .base_tirex import BaseTirexEmbeddingModel
+
+
+class BaseTirexRegressor(BaseTirexEmbeddingModel):
+    """Abstract base class for TiRex regression models."""
+
+    @torch.inference_mode()
+    def predict(self, x: torch.Tensor) -> torch.Tensor:
+        """Predict values for input time series data.
+
+        Args:
+            x: Input time series data as torch.Tensor with shape
+                (batch_size, num_variates, seq_len).
+        Returns:
+            torch.Tensor: Predicted values with shape (batch_size,).
+        """
+        emb = self._compute_embeddings(x)
+        return torch.from_numpy(self.head.predict(emb)).float()

tirex_mirror-2025.12.16/src/tirex/models/base/base_tirex.py

@@ -0,0 +1,109 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+from abc import ABC, abstractmethod
+
+import numpy as np
+import torch
+
+from tirex.util import train_val_split
+
+from ..embedding import TiRexEmbedding
+
+
+class BaseTirexEmbeddingModel(ABC):
+    """Abstract base class for TiRex models.
+
+    This base class provides common functionality for all TiRex classifier and regression models,
+    including embedding model initialization and a consistent interface.
+
+    """
+
+    def __init__(
+        self, data_augmentation: bool = False, device: str | None = None, compile: bool = False, batch_size: int = 512
+    ) -> None:
+        """Initializes a base TiRex model.
+
+        This base class initializes the embedding model and common configuration
+        used by both classification and regression models.
+
+        Args:
+            data_augmentation : bool
+                Whether to use data_augmentation for embeddings (sample statistics and first-order differences of the original data). Default: False
+            device : str | None
+                Device to run the embedding model on. If None, uses CUDA if available, else CPU. Default: None
+            compile: bool
+                Whether to compile the frozen embedding model. Default: False
+            batch_size : int
+                Batch size for embedding calculations. Default: 512
+        """
+
+        # Set device
+        if device is None:
+            device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        self.device = device
+        self._compile = compile
+
+        self.batch_size = batch_size
+        self.data_augmentation = data_augmentation
+        self.emb_model = TiRexEmbedding(
+            device=self.device,
+            data_augmentation=self.data_augmentation,
+            batch_size=self.batch_size,
+            compile=self._compile,
+        )
+
+    @abstractmethod
+    def fit(self, train_data: tuple[torch.Tensor, torch.Tensor]) -> None:
+        """Abstract method for model training"""
+        pass
+
+    def _compute_embeddings(self, x: torch.Tensor) -> np.ndarray:
+        """Compute embeddings for input time series data.
+
+        Args:
+            x: Input time series data as torch.Tensor with shape
+                (batch_size, num_variates, seq_len).
+
+        Returns:
+            np.ndarray: Embeddings with shape (batch_size, embedding_dim).
+        """
+        self.emb_model.eval()
+        x = x.to(self.device)
+        return self.emb_model(x).cpu().numpy()
+
+    def _create_train_val_datasets(
+        self,
+        train_data: tuple[torch.Tensor, torch.Tensor],
+        val_data: tuple[torch.Tensor, torch.Tensor] | None = None,
+        val_split_ratio: float = 0.2,
+        stratify: bool = False,
+        seed: int | None = None,
+    ) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]:
+        if val_data is None:
+            train_data, val_data = train_val_split(
+                train_data=train_data, val_split_ratio=val_split_ratio, stratify=stratify, seed=seed
+            )
+        return train_data, val_data
+
+    @abstractmethod
+    def save_model(self, path: str) -> None:
+        """Save model to file.
+
+        Args:
+            path: File path where the model should be saved.
+        """
+        pass
+
+    @classmethod
+    @abstractmethod
+    def load_model(cls, path: str):
+        """Load model from file.
+
+        Args:
+            path: File path to the saved model checkpoint.
+
+        Returns:
+            Instance of the model class with loaded weights and configuration.
+        """
+        pass

tirex_mirror-2025.12.16/src/tirex/models/classification/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+from .gbm_classifier import TirexGBMClassifier
+from .linear_classifier import TirexLinearClassifier
+from .rf_classifier import TirexRFClassifier
+
+__all__ = ["TirexLinearClassifier", "TirexRFClassifier", "TirexGBMClassifier"]

tirex_mirror-2025.12.16/src/tirex/models/classification/gbm_classifier.py

@@ -0,0 +1,188 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
+import joblib
+import torch
+from lightgbm import LGBMClassifier, early_stopping
+
+from ..base.base_classifier import BaseTirexClassifier
+
+
+class TirexGBMClassifier(BaseTirexClassifier):
+    """
+    A Gradient Boosting classifier that uses time series embeddings as features.
+
+    This classifier combines a pre-trained embedding model for feature extraction with a
+    Gradient Boosting classifier.
+
+    Example:
+        >>> from tirex.models.classification import TirexGBMClassifier
+        >>>
+        >>> # Create model with custom LightGBM parameters
+        >>> model = TirexGBMClassifier(
+        ...     data_augmentation=True,
+        ...     n_estimators=50,
+        ...     random_state=42
+        ... )
+        >>>
+        >>> # Prepare data (can use NumPy arrays or PyTorch tensors)
+        >>> X_train = torch.randn(100, 1, 128)  # 100 samples, 1 number of variates, 128 sequence length
+        >>> y_train = torch.randint(0, 3, (100,))  # 3 classes
+        >>>
+        >>> # Train the model
+        >>> model.fit((X_train, y_train))
+        >>>
+        >>> # Make predictions
+        >>> X_test = torch.randn(20, 1, 128)
+        >>> predictions = model.predict(X_test)
+        >>> probabilities = model.predict_proba(X_test)
+    """
+
+    def __init__(
+        self,
+        data_augmentation: bool = False,
+        device: str | None = None,
+        compile: bool = False,
+        batch_size: int = 512,
+        early_stopping_rounds: int | None = 10,
+        min_delta: float = 0.0,
+        val_split_ratio: float = 0.2,
+        stratify: bool = True,
+        # LightGBM kwargs
+        **lgbm_kwargs,
+    ) -> None:
+        """Initializes Embedding Based Gradient Boosting Classification model.
+
+        Args:
+            data_augmentation : bool
+                Whether to use data_augmentation for embeddings (sample statistics and first-order differences of the original data). Default: False
+            device : str | None
+                Device to run the embedding model on. If None, uses CUDA if available, else CPU. Default: None
+            compile: bool
+                Whether to compile the frozen embedding model. Default: False
+            batch_size : int
+                Batch size for embedding calculations. Default: 512
+            early_stopping_rounds: int | None
+                Number of rounds without improvement of all metrics for Early Stopping. Default: 10
+            min_delta: float
+                Minimum improvement in score to keep training. Default 0.0
+            val_split_ratio : float
+                Proportion of training data to use for validation, if validation data are not provided. Default: 0.2
+            stratify : bool
+                Whether to stratify the train/validation split by class labels. Default: True
+            **lgbm_kwargs
+                Additional keyword arguments to pass to LightGBM's LGBMClassifier.
+                Common options include n_estimators, max_depth, learning_rate, random_state, etc.
+        """
+        super().__init__(data_augmentation=data_augmentation, device=device, compile=compile, batch_size=batch_size)
+
+        # Early Stopping callback
+        self.early_stopping_rounds = early_stopping_rounds
+        self.min_delta = min_delta
+
+        # Data split parameters:
+        self.val_split_ratio = val_split_ratio
+        self.stratify = stratify
+
+        # Extract random_state for train_val_split if provided
+        self.random_state = lgbm_kwargs.get("random_state", None)
+
+        self.head = LGBMClassifier(**lgbm_kwargs)
+
+    @torch.inference_mode()
+    def fit(
+        self,
+        train_data: tuple[torch.Tensor, torch.Tensor],
+        val_data: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> None:
+        """Train the LightGBM classifier on embedded time series data.
+
+        This method generates embeddings for the training data using the embedding
+        model, then trains the LightGBM classifier on these embeddings.
+
+        Args:
+            train_data: Tuple of (X_train, y_train) where X_train is the input time
+                series data (torch.Tensor) and y_train is a torch.Tensor
+                of class labels.
+            val_data: Optional tuple of (X_val, y_val) for validation where X_train is the input time
+                series data (torch.Tensor) and y_train is a torch.Tensor
+                of class labels. If None, validation data will be automatically split from train_data (20% split).
+        """
+
+        (X_train, y_train), (X_val, y_val) = self._create_train_val_datasets(
+            train_data=train_data,
+            val_data=val_data,
+            val_split_ratio=self.val_split_ratio,
+            stratify=self.stratify,
+            seed=self.random_state,
+        )
+
+        X_train = X_train.to(self.device)
+        X_val = X_val.to(self.device)
+
+        embeddings_train = self._compute_embeddings(X_train)
+        embeddings_val = self._compute_embeddings(X_val)
+
+        y_train = y_train.detach().cpu().numpy() if isinstance(y_train, torch.Tensor) else y_train
+        y_val = y_val.detach().cpu().numpy() if isinstance(y_val, torch.Tensor) else y_val
+
+        self.head.fit(
+            embeddings_train,
+            y_train,
+            eval_set=[(embeddings_val, y_val)],
+            callbacks=[early_stopping(stopping_rounds=self.early_stopping_rounds, min_delta=self.min_delta)]
+            if self.early_stopping_rounds is not None
+            else None,
+        )
+
+    def save_model(self, path: str) -> None:
+        """This method saves the trained LightGBM classifier head (joblib format) and embedding information.
+
+        Args:
+            path: File path where the model should be saved (e.g., 'model.joblib').
+        """
+
+        payload = {
+            "data_augmentation": self.data_augmentation,
+            "compile": self._compile,
+            "batch_size": self.batch_size,
+            "early_stopping_rounds": self.early_stopping_rounds,
+            "min_delta": self.min_delta,
+            "val_split_ratio": self.val_split_ratio,
+            "stratify": self.stratify,
+            "head": self.head,
+        }
+        joblib.dump(payload, path)
+
+    @classmethod
+    def load_model(cls, path: str) -> "TirexGBMClassifier":
+        """Load a saved model from file.
+
+        This reconstructs the model with the embedding configuration and loads
+        the trained LightGBM classifier from a checkpoint file created by save_model().
+
+        Args:
+            path: File path to the saved model checkpoint.
+        Returns:
+            TirexGBMClassifier: The loaded model with trained Gradient Boosting, ready for inference.
+        """
+        checkpoint = joblib.load(path)
+
+        # Create new instance with saved configuration
+        model = cls(
+            data_augmentation=checkpoint["data_augmentation"],
+            compile=checkpoint["compile"],
+            batch_size=checkpoint["batch_size"],
+            early_stopping_rounds=checkpoint["early_stopping_rounds"],
+            min_delta=checkpoint["min_delta"],
+            val_split_ratio=checkpoint["val_split_ratio"],
+            stratify=checkpoint["stratify"],
+        )
+
+        # Load the trained LightGBM head
+        model.head = checkpoint["head"]
+
+        # Extract random_state from the loaded head if available
+        model.random_state = getattr(model.head, "random_state", None)
+
+        return model

{tirex_mirror-2025.11.29 → tirex_mirror-2025.12.16}/src/tirex/models/classification/linear_classifier.py

@@ -1,12 +1,15 @@
+# Copyright (c) NXAI GmbH.
+# This software may be used and distributed according to the terms of the NXAI Community License Agreement.
+
 from dataclasses import asdict
 
 import torch
 
-from .embedding import TiRexEmbedding
-from .trainer import TrainConfig, Trainer
+from ..base.base_classifier import BaseTirexClassifier
+from ..trainer import TrainConfig, Trainer, TrainingMetrics
 
 
-class TirexClassifierTorch(torch.nn.Module):
+class TirexLinearClassifier(BaseTirexClassifier, torch.nn.Module):
     """
     A PyTorch classifier that combines time series embeddings with a linear classification head.
 
@@ -16,10 +19,10 @@ class TirexClassifierTorch(torch.nn.Module):
 
     Example:
         >>> import torch
-        >>> from tirex.models.classification import TirexClassifierTorch
+        >>> from tirex.models.classification import TirexLinearClassifier
         >>>
-        >>> # Create model with TIREX embeddings
-        >>> model = TirexClassifierTorch(
+        >>> # Create model with TiRex embeddings
+        >>> model = TirexLinearClassifier(
         ...     data_augmentation=True,
         ...     max_epochs=2,
         ...     lr=1e-4,
@@ -43,8 +46,9 @@ class TirexClassifierTorch(torch.nn.Module):
         self,
         data_augmentation: bool = False,
         device: str | None = None,
+        compile: bool = False,
         # Training parameters
-        max_epochs: int = 50,
+        max_epochs: int = 10,
         lr: float = 1e-4,
         weight_decay: float = 0.01,
         batch_size: int = 512,
@@ -62,11 +66,13 @@ class TirexClassifierTorch(torch.nn.Module):
 
         Args:
             data_augmentation : bool | None
-                Whether to use data_augmentation for embeddings (stats and first-order differences of the original data). Default: False
+                Whether to use data_augmentation for embeddings (sample statistics and first-order differences of the original data). Default: False
             device : str | None
-                Device to run the model on. If None, uses CUDA if available, else CPU. Default: None
+                Device to run the embedding model on. If None, uses CUDA if available, else CPU. Default: None
+            compile: bool
+                Whether to compile the frozen embedding model. Default: False
             max_epochs : int
-                Maximum number of training epochs. Default: 50
+                Maximum number of training epochs. Default: 10
             lr : float
                 Learning rate for the optimizer. Default: 1e-4
             weight_decay : float
@@ -91,15 +97,9 @@ class TirexClassifierTorch(torch.nn.Module):
                 Dropout probability for the classification head. If None, no dropout is used. Default: None
         """
 
-        super().__init__()
-
-        if device is None:
-            device = "cuda:0" if torch.cuda.is_available() else "cpu"
-        self.device = device
+        torch.nn.Module.__init__(self)
 
-        # Create embedding model
-        self.emb_model = TiRexEmbedding(device=self.device, data_augmentation=data_augmentation, batch_size=batch_size)
-        self.data_augmentation = data_augmentation
+        super().__init__(data_augmentation=data_augmentation, device=device, compile=compile, batch_size=batch_size)
 
         # Head parameters
         self.dropout = dropout
@@ -115,6 +115,7 @@ class TirexClassifierTorch(torch.nn.Module):
             lr=lr,
             weight_decay=weight_decay,
             class_weights=class_weights,
+            task_type="classification",
             batch_size=batch_size,
             val_split_ratio=val_split_ratio,
             stratify=stratify,
@@ -132,9 +133,7 @@ class TirexClassifierTorch(torch.nn.Module):
 
     @torch.inference_mode()
     def _identify_head_dims(self, x: torch.Tensor, y: torch.Tensor) -> None:
-        self.emb_model.eval()
-        sample_emb = self.emb_model(x[:1])
-        self.emb_dim = sample_emb.shape[-1]
+        self.emb_dim = self._compute_embeddings(x[:1]).shape[-1]
         self.num_classes = len(torch.unique(y))
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -155,7 +154,7 @@ class TirexClassifierTorch(torch.nn.Module):
 
     def fit(
         self, train_data: tuple[torch.Tensor, torch.Tensor], val_data: tuple[torch.Tensor, torch.Tensor] | None = None
-    ) -> dict[str, float]:
+    ) -> TrainingMetrics:
         """Train the classification head on the provided data.
 
         This method initializes the classification head based on the data dimensions,
@@ -221,6 +220,7 @@ class TirexClassifierTorch(torch.nn.Module):
             {
                 "head_state_dict": self.head.state_dict(),  # need to save only head, embedding is frozen
                 "data_augmentation": self.data_augmentation,
+                "compile": self._compile,
                 "emb_dim": self.emb_dim,
                 "num_classes": self.num_classes,
                 "dropout": self.dropout,
@@ -230,7 +230,7 @@ class TirexClassifierTorch(torch.nn.Module):
         )
 
     @classmethod
-    def load_model(cls, path: str) -> "TirexClassifierTorch":
+    def load_model(cls, path: str) -> "TirexLinearClassifier":
         """Load a saved model from file.
 
         This reconstructs the model architecture and loads the trained weights from
@@ -239,7 +239,7 @@ class TirexClassifierTorch(torch.nn.Module):
         Args:
             path: File path to the saved model checkpoint.
         Returns:
-            TirexClassifierTorch: The loaded model with trained weights, ready for inference.
+            TirexLinearClassifier: The loaded model with trained weights, ready for inference.
         """
         checkpoint = torch.load(path)
 
@@ -248,6 +248,7 @@ class TirexClassifierTorch(torch.nn.Module):
 
         model = cls(
             data_augmentation=checkpoint["data_augmentation"],
+            compile=checkpoint["compile"],
             dropout=checkpoint["dropout"],
             max_epochs=train_config_dict.get("max_epochs", 50),
             lr=train_config_dict.get("lr", 1e-4),