tirex-mirror 2025.11.26__py3-none-any.whl → 2025.11.29__py3-none-any.whl

tirex/models/classification/__init__.py

@@ -0,0 +1,8 @@
+# classification/__init__.py
+from .linear_classifier import TirexClassifierTorch
+from .rf_classifier import TirexRFClassifier
+
+__all__ = [
+    "TirexClassifierTorch",
+    "TirexRFClassifier",
+]

tirex/models/classification/embedding.py

@@ -0,0 +1,125 @@
+import inspect
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from tirex import load_model
+
+from .utils import nanmax, nanmin, nanstd
+
+
+class TiRexEmbedding(nn.Module):
+    def __init__(self, device: str | None = None, data_augmentation: bool = False, batch_size: int = 512) -> None:
+        super().__init__()
+        self.data_augmentation = data_augmentation
+        self.number_of_patches = 8
+        self.batch_size = batch_size
+
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = device
+
+        self.model = load_model(path="NX-AI/TiRex", device=self.device)
+
+    def _gen_emb_batched(self, data: torch.Tensor) -> torch.Tensor:
+        batches = list(torch.split(data, self.batch_size))
+        embedding_list = []
+        for batch in batches:
+            embedding = self.model._embed_context(batch)
+            embedding_list.append(embedding.cpu())
+        return torch.cat(embedding_list, dim=0)
+
+    def _calculate_n_patches(self, data: torch.Tensor) -> int:
+        _, _, n_steps = data.shape
+        n_patches = -(-n_steps // self.model.config.input_patch_size)
+        return n_patches
+
+    def forward(self, data: torch.Tensor) -> torch.Tensor:
+        n_patches = self._calculate_n_patches(data)
+
+        embedding = torch.stack(
+            [self._gen_emb_batched(var_slice) for var_slice in torch.unbind(data, dim=1)], dim=1
+        )  # Stack in case of multivar
+        embedding = self.process_embedding(embedding, n_patches)
+
+        if self.data_augmentation:
+            # Difference Embedding
+            diff_data = torch.diff(data, dim=-1, prepend=data[..., :0])
+            n_patches = self._calculate_n_patches(diff_data)
+
+            diff_embedding = torch.stack(
+                [self._gen_emb_batched(var_slice) for var_slice in torch.unbind(diff_data, dim=1)], dim=1
+            )
+            diff_embedding = self.process_embedding(diff_embedding, n_patches)
+            embedding = torch.cat((diff_embedding, embedding), dim=-1)
+
+            # Stats Embedding
+            stat_features = self._generate_stats_features(data)
+            normalized_stats = self._normalize_stats(stat_features)
+            normalized_stats = normalized_stats.to(embedding.device)
+
+            # Concat all together
+            embedding = torch.cat((embedding, normalized_stats), dim=-1)
+
+        return embedding
+
+    def process_embedding(self, embedding: torch.Tensor, n_patches: int) -> torch.Tensor:
+        # embedding shape: (bs, var_dim, n_patches, n_layer, emb_dim)
+        embedding = embedding[:, :, -n_patches:, :, :]
+        embedding = torch.mean(embedding, dim=2)  # sequence
+        embedding = torch.nn.functional.normalize(embedding, p=2, dim=-1)
+        embedding = torch.transpose(embedding, 1, -2).flatten(start_dim=-2)  # var
+        embedding = torch.transpose(embedding, 1, -2).flatten(start_dim=-2)  # layer
+        embedding = F.layer_norm(embedding, (embedding.shape[-1],))
+        return embedding
+
+    def _normalize_stats(self, stat_features: torch.Tensor) -> torch.Tensor:
+        dataset_mean = torch.nanmean(stat_features, dim=0, keepdim=True)
+        dataset_std = nanstd(stat_features, dim=0, keepdim=True)
+        stat_features = (stat_features - dataset_mean) / (dataset_std + 1e-8)
+        stat_features = torch.nan_to_num(stat_features, nan=0.0)
+
+        stat_features = (stat_features - stat_features.nanmean(dim=-1, keepdim=True)) / (
+            stat_features.std(dim=-1, keepdim=True) + 1e-8
+        )
+        return stat_features
+
+    def _generate_stats_features(self, data: torch.Tensor) -> torch.Tensor:
+        bs, variates, n_steps = data.shape
+
+        patch_size = max(1, n_steps // self.number_of_patches)
+        n_full_patches = n_steps // patch_size
+        n_remain = n_steps % patch_size
+
+        # [batch, variates, n_patches, patch_size]
+        patches = data[..., : n_full_patches * patch_size].unfold(-1, patch_size, patch_size)
+
+        # Stats for full patches
+        patch_means = torch.nanmean(patches, dim=-1)
+        patch_stds = nanstd(patches, dim=-1)
+        patch_maxes = nanmax(patches, dim=-1)
+        patch_mins = nanmin(patches, dim=-1)
+
+        stats = [patch_means, patch_stds, patch_maxes, patch_mins]
+
+        # Handle last smaller patch if needed
+        if n_remain > 0:
+            self._handle_remaining_patch(data, stats, n_full_patches * patch_size)
+
+        stats = torch.stack(stats, dim=-1)  # [batch, variates, n_patches(+1), 4]
+        return stats.flatten(start_dim=1)  # [batch, variates * n_patches * 4]
+
+    def _handle_remaining_patch(self, data: torch.Tensor, stats: list[torch.Tensor], full_patch_length: int) -> None:
+        last_patch = data[..., full_patch_length:]
+
+        mean_last = last_patch.mean(dim=-1, keepdim=True)
+        std_last = last_patch.std(dim=-1, keepdim=True)
+        max_last = last_patch.max(dim=-1, keepdim=True)
+        min_last = last_patch.min(dim=-1, keepdim=True)
+
+        stats[0] = torch.cat([stats[0], mean_last], dim=-1)
+        stats[1] = torch.cat([stats[1], std_last], dim=-1)
+        stats[2] = torch.cat([stats[2], max_last], dim=-1)
+        stats[3] = torch.cat([stats[3], min_last], dim=-1)

tirex/models/classification/linear_classifier.py

@@ -0,0 +1,274 @@
+from dataclasses import asdict
+
+import torch
+
+from .embedding import TiRexEmbedding
+from .trainer import TrainConfig, Trainer
+
+
+class TirexClassifierTorch(torch.nn.Module):
+    """
+    A PyTorch classifier that combines time series embeddings with a linear classification head.
+
+    This model uses a pre-trained TiRex embedding model to generate feature representations from time series
+    data, followed by a linear layer (with optional dropout) for classification. The embedding backbone
+    is frozen during training, and only the classification head is trained.
+
+    Example:
+        >>> import torch
+        >>> from tirex.models.classification import TirexClassifierTorch
+        >>>
+        >>> # Create model with TIREX embeddings
+        >>> model = TirexClassifierTorch(
+        ...     data_augmentation=True,
+        ...     max_epochs=2,
+        ...     lr=1e-4,
+        ...     batch_size=32
+        ... )
+        >>>
+        >>> # Prepare data
+        >>> 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
+        >>> metrics = model.fit((X_train, y_train)) # doctest: +ELLIPSIS
+        Epoch 1, Train Loss: ...
+        >>> # 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,
+        # Training parameters
+        max_epochs: int = 50,
+        lr: float = 1e-4,
+        weight_decay: float = 0.01,
+        batch_size: int = 512,
+        val_split_ratio: float = 0.2,
+        stratify: bool = True,
+        patience: int = 7,
+        delta: float = 0.001,
+        log_every_n_steps: int = 5,
+        seed: int | None = None,
+        class_weights: torch.Tensor | None = None,
+        # Head parameters
+        dropout: float | None = None,
+    ) -> None:
+        """Initializes Embedding Based Linear Classification model.
+
+        Args:
+            data_augmentation : bool | None
+                Whether to use data_augmentation for embeddings (stats 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
+            max_epochs : int
+                Maximum number of training epochs. Default: 50
+            lr : float
+                Learning rate for the optimizer. Default: 1e-4
+            weight_decay : float
+                Weight decay coefficient. Default: 0.01
+            batch_size : int
+                Batch size for training and embedding calculations. Default: 512
+            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
+            patience : int
+                Number of epochs to wait for improvement before early stopping. Default: 7
+            delta : float
+                Minimum change in validation loss to qualify as an improvement. Default: 0.001
+            log_every_n_steps : int
+                Frequency of logging during training. Default: 5
+            seed : int | None
+                Random seed for reproducibility. If None, no seed is set. Default: None
+            class_weights: torch.Tensor | None
+                Weight classes according to given values, has to be a Tensor of size number of categories. Default: None
+            dropout : float | None
+                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
+
+        # Create embedding model
+        self.emb_model = TiRexEmbedding(device=self.device, data_augmentation=data_augmentation, batch_size=batch_size)
+        self.data_augmentation = data_augmentation
+
+        # Head parameters
+        self.dropout = dropout
+        self.head = None
+        self.emb_dim = None
+        self.num_classes = None
+
+        # Train config
+        train_config = TrainConfig(
+            max_epochs=max_epochs,
+            log_every_n_steps=log_every_n_steps,
+            device=self.device,
+            lr=lr,
+            weight_decay=weight_decay,
+            class_weights=class_weights,
+            batch_size=batch_size,
+            val_split_ratio=val_split_ratio,
+            stratify=stratify,
+            patience=patience,
+            delta=delta,
+            seed=seed,
+        )
+        self.trainer = Trainer(self, train_config=train_config)
+
+    def _init_classifier(self, emb_dim: int, num_classes: int, dropout: float | None) -> torch.nn.Module:
+        if dropout:
+            return torch.nn.Sequential(torch.nn.Dropout(p=dropout), torch.nn.Linear(emb_dim, num_classes))
+        else:
+            return torch.nn.Linear(emb_dim, num_classes)
+
+    @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.num_classes = len(torch.unique(y))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass through the embedding model and classification head.
+
+        Args:
+            x: Input tensor of time series data with shape (batch_size, num_variates, seq_len).
+        Returns:
+            torch.Tensor: Logits for each class with shape (batch_size, num_classes).
+        Raises:
+            RuntimeError: If the classification head has not been initialized via fit().
+        """
+        if self.head is None:
+            raise RuntimeError("Head not initialized. Call fit() first to automatically build the head.")
+
+        embedding = self.emb_model(x).to(self.device)
+        return self.head(embedding)
+
+    def fit(
+        self, train_data: tuple[torch.Tensor, torch.Tensor], val_data: tuple[torch.Tensor, torch.Tensor] | None = None
+    ) -> dict[str, float]:
+        """Train the classification head on the provided data.
+
+        This method initializes the classification head based on the data dimensions,
+        then trains it on provided data. The embedding model remains frozen.
+
+        Args:
+            train_data: Tuple of (X_train, y_train) where X_train is the input time series
+                data and y_train are the corresponding class labels.
+            val_data: Optional tuple of (X_val, y_val) for validation. If None and
+                val_split_ratio > 0, validation data will be split from train_data.
+
+        Returns:
+            dict[str, float]: Dictionary containing final training and validation losses.
+        """
+        X_train, y_train = train_data
+
+        self._identify_head_dims(X_train, y_train)
+        self.head = self._init_classifier(self.emb_dim, self.num_classes, self.dropout)
+        self.head = self.head.to(self.trainer.device)
+
+        return self.trainer.fit(train_data, val_data=val_data)
+
+    @torch.inference_mode()
+    def predict(self, x: torch.Tensor) -> torch.Tensor:
+        """Predict class labels for input time series data.
+
+        Args:
+            x: Input tensor of time series data with shape (batch_size, num_variates, seq_len).
+        Returns:
+            torch.Tensor: Predicted class labels with shape (batch_size,).
+        """
+        self.eval()
+        x = x.to(self.device)
+        logits = self.forward(x)
+        return torch.argmax(logits, dim=1)
+
+    @torch.inference_mode()
+    def predict_proba(self, x: torch.Tensor) -> torch.Tensor:
+        """Predict class probabilities for input time series data.
+
+        Args:
+            x: Input tensor of time series data with shape (batch_size, num_variates, seq_len).
+        Returns:
+            torch.Tensor: Class probabilities with shape (batch_size, num_classes).
+        """
+        self.eval()
+        x = x.to(self.device)
+        logits = self.forward(x)
+        return torch.softmax(logits, dim=1)
+
+    def save_model(self, path: str) -> None:
+        """Save the trained classification head.
+
+        This function saves the trained classification head weights (.pt format), embedding configuration,
+        model dimensions, and device information. The embedding model itself is not
+        saved as it uses a pre-trained backbone that can be reloaded.
+
+        Args:
+            path: File path where the model should be saved (e.g., 'model.pt').
+        """
+        train_config_dict = asdict(self.trainer.train_config)
+        torch.save(
+            {
+                "head_state_dict": self.head.state_dict(),  # need to save only head, embedding is frozen
+                "data_augmentation": self.data_augmentation,
+                "emb_dim": self.emb_dim,
+                "num_classes": self.num_classes,
+                "dropout": self.dropout,
+                "train_config": train_config_dict,
+            },
+            path,
+        )
+
+    @classmethod
+    def load_model(cls, path: str) -> "TirexClassifierTorch":
+        """Load a saved model from file.
+
+        This reconstructs the model architecture and loads the trained weights from
+        a checkpoint file created by save_model().
+
+        Args:
+            path: File path to the saved model checkpoint.
+        Returns:
+            TirexClassifierTorch: The loaded model with trained weights, ready for inference.
+        """
+        checkpoint = torch.load(path)
+
+        # Extract train_config if available, otherwise use defaults
+        train_config_dict = checkpoint.get("train_config", {})
+
+        model = cls(
+            data_augmentation=checkpoint["data_augmentation"],
+            dropout=checkpoint["dropout"],
+            max_epochs=train_config_dict.get("max_epochs", 50),
+            lr=train_config_dict.get("lr", 1e-4),
+            weight_decay=train_config_dict.get("weight_decay", 0.01),
+            batch_size=train_config_dict.get("batch_size", 512),
+            val_split_ratio=train_config_dict.get("val_split_ratio", 0.2),
+            stratify=train_config_dict.get("stratify", True),
+            patience=train_config_dict.get("patience", 7),
+            delta=train_config_dict.get("delta", 0.001),
+            log_every_n_steps=train_config_dict.get("log_every_n_steps", 5),
+            seed=train_config_dict.get("seed", None),
+            class_weights=train_config_dict.get("class_weights", None),
+        )
+
+        # Initialize head with dimensions
+        model.emb_dim = checkpoint["emb_dim"]
+        model.num_classes = checkpoint["num_classes"]
+        model.head = model._init_classifier(model.emb_dim, model.num_classes, model.dropout)
+
+        # Load the trained weights
+        model.head.load_state_dict(checkpoint["head_state_dict"])
+        model.to(model.device)
+
+        return model

tirex/models/classification/rf_classifier.py

@@ -0,0 +1,155 @@
+import joblib
+import numpy as np
+import torch
+from sklearn.ensemble import RandomForestClassifier
+
+from .embedding import TiRexEmbedding
+
+
+class TirexRFClassifier:
+    """
+    A Random Forest classifier that uses time series embeddings as features.
+
+    This classifier combines a pre-trained embedding model for feature extraction with a scikit-learn
+    Random Forest classifier. The embedding model generates fixed-size feature vectors from variable-length
+    time series, which are then used to train the Random Forest.
+
+    Example:
+        >>> import numpy as np
+        >>> from tirex.models.classification import TirexRFClassifier
+        >>>
+        >>> # Create model with custom Random Forest parameters
+        >>> model = TirexRFClassifier(
+        ...     data_augmentation=True,
+        ...     n_estimators=50,
+        ...     max_depth=10,
+        ...     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,
+        batch_size: int = 512,
+        # Random Forest parameters
+        **rf_kwargs,
+    ) -> None:
+        """Initializes Embedding Based Random Forest Classification model.
+
+        Args:
+            data_augmentation : bool
+                Whether to use data_augmentation for embeddings (stats 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
+            batch_size : int
+                Batch size for embedding calculations. Default: 512
+            **rf_kwargs
+                Additional keyword arguments to pass to sklearn's RandomForestClassifier.
+                Common options include n_estimators, max_depth, min_samples_split, random_state, etc.
+        """
+
+        # Set device
+        if device is None:
+            device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        self.device = device
+
+        self.emb_model = TiRexEmbedding(device=self.device, data_augmentation=data_augmentation, batch_size=batch_size)
+        self.data_augmentation = data_augmentation
+
+        self.head = RandomForestClassifier(**rf_kwargs)
+
+    @torch.inference_mode()
+    def fit(self, train_data: tuple[torch.Tensor, torch.Tensor]) -> None:
+        """Train the Random Forest classifier on embedded time series data.
+
+        This method generates embeddings for the training data using the embedding
+        model, then trains the Random Forest 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 numpy array
+                of class labels.
+        """
+        X_train, y_train = train_data
+
+        if isinstance(y_train, torch.Tensor):
+            y_train = y_train.detach().cpu().numpy()
+
+        embeddings = self.emb_model(X_train).cpu().numpy()
+        self.head.fit(embeddings, y_train)
+
+    @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 or np.ndarray with shape
+                (batch_size, num_variates, seq_len).
+        Returns:
+            torch.Tensor: Predicted class labels with shape (batch_size,).
+        """
+
+        embeddings = self.emb_model(x).cpu().numpy()
+        return torch.from_numpy(self.head.predict(embeddings)).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 or np.ndarray with shape
+                (batch_size, num_variates, seq_len).
+        Returns:
+            torch.Tensor: Class probabilities with shape (batch_size, num_classes).
+        """
+        embeddings = self.emb_model(x).cpu().numpy()
+        return torch.from_numpy(self.head.predict_proba(embeddings))
+
+    def save_model(self, path: str) -> None:
+        """This method saves the trained Random Forest classifier head and embedding information in joblib format
+
+        Args:
+            path: File path where the model should be saved (e.g., 'model.joblib').
+        """
+        payload = {
+            "data_augmentation": self.data_augmentation,
+            "head": self.head,
+        }
+        joblib.dump(payload, path)
+
+    @classmethod
+    def load_model(cls, path: str) -> "TirexRFClassifier":
+        """Load a saved model from file.
+
+        This reconstructs the model with the embedding configuration and loads
+        the trained Random Forest classifier from a checkpoint file created by save_model().
+
+        Args:
+            path: File path to the saved model checkpoint.
+        Returns:
+            TirexRFClassifier: The loaded model with trained Random Forest, ready for inference.
+        """
+        checkpoint = joblib.load(path)
+
+        # Create new instance with saved configuration
+        model = cls(
+            data_augmentation=checkpoint["data_augmentation"],
+        )
+
+        # Load the trained Random Forest head
+        model.head = checkpoint["head"]
+
+        return model

tirex/models/classification/trainer.py

@@ -0,0 +1,171 @@
+from dataclasses import dataclass
+
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+
+from .utils import EarlyStopping, set_seed, train_val_split
+
+
+@dataclass
+class TrainConfig:
+    # Training loop parameters
+    max_epochs: int
+    log_every_n_steps: int
+    device: str
+
+    # Optimizer parameters
+    lr: float
+    weight_decay: float
+
+    # Loss parameters
+    class_weights: torch.Tensor | None
+
+    # Data loading parameters
+    batch_size: int
+    val_split_ratio: float
+    stratify: bool
+
+    # Earlystopping parameters
+    patience: int
+    delta: float
+
+    # Reproducability
+    seed: int | None
+
+    def __post_init__(self) -> None:
+        if self.max_epochs <= 0:
+            raise ValueError(f"max_epochs must be positive, got {self.max_epochs}")
+
+        if self.log_every_n_steps <= 0:
+            raise ValueError(f"log_every_n_steps must be positive, got {self.log_every_n_steps}")
+
+        if self.lr <= 0:
+            raise ValueError(f"lr (learning rate) must be positive, got {self.lr}")
+
+        if self.weight_decay < 0:
+            raise ValueError(f"weight_decay must be non-negative, got {self.weight_decay}")
+
+        if self.batch_size <= 0:
+            raise ValueError(f"batch_size must be positive, got {self.batch_size}")
+
+        if not (0 < self.val_split_ratio < 1):
+            raise ValueError(f"val_split_ratio must be in (0, 1), got {self.val_split_ratio}")
+
+        if self.patience <= 0:
+            raise ValueError(f"patience must be positive, got {self.patience}")
+
+        if self.delta < 0:
+            raise ValueError(f"delta must be non-negative, got {self.delta}")
+
+
+class Trainer:
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        train_config: TrainConfig,
+    ) -> None:
+        self.device = train_config.device
+        self.train_config = train_config
+
+        self.model = model.to(self.device)
+        class_weights = (
+            self.train_config.class_weights.to(self.device) if self.train_config.class_weights is not None else None
+        )
+        self.loss_fn = torch.nn.CrossEntropyLoss(weight=class_weights).to(self.device)
+
+        self.optimizer = None
+        self.early_stopper = EarlyStopping(patience=self.train_config.patience, delta=self.train_config.delta)
+
+    def fit(
+        self, train_data: tuple[torch.Tensor, torch.Tensor], val_data: tuple[torch.Tensor, torch.Tensor] | None = None
+    ) -> dict[str, float]:
+        if self.train_config.seed is not None:
+            set_seed(self.train_config.seed)
+
+        self._freeze_embedding()
+
+        if self.optimizer is None:
+            self.optimizer = torch.optim.AdamW(
+                self.model.parameters(), lr=self.train_config.lr, weight_decay=self.train_config.weight_decay
+            )
+
+        train_loader, val_loader = self._create_data_loaders(train_data, val_data)
+
+        for epoch in range(self.train_config.max_epochs):
+            train_loss = self._train_epoch(train_loader)
+            val_loss = self._validate_epoch(val_loader)
+
+            self._log_epoch_metrics(epoch, train_loss, val_loss)
+
+            stop_training = self.early_stopper(epoch=epoch + 1, val_loss=val_loss)
+            if stop_training:
+                break
+
+        return {"train_loss": train_loss, "val_loss": val_loss}
+
+    def _freeze_embedding(self) -> None:
+        if hasattr(self.model, "emb_model"):
+            for param in self.model.emb_model.parameters():
+                param.requires_grad = False
+
+    def _train_epoch(self, train_loader: DataLoader) -> float:
+        train_loss = []
+        self.model.train()
+        for batch in train_loader:
+            x, y = batch
+            x = x.to(self.device)
+            y = y.to(self.device)
+
+            self.optimizer.zero_grad()
+            y_hat = self.model.head(x)  # Only classification head is involved, embeddings are precomputed
+            loss = self.loss_fn(y_hat, y)
+
+            loss.backward()
+            self.optimizer.step()
+
+            train_loss.append(loss.detach())
+
+        return torch.stack(train_loss).mean().item()
+
+    @torch.inference_mode()
+    def _validate_epoch(self, val_loader: DataLoader) -> float:
+        self.model.eval()
+        val_loss = []
+        for batch in val_loader:
+            x, y = batch
+            x = x.to(self.device)
+            y = y.to(self.device)
+
+            y_hat = self.model.head(x)  # Only classification head is involved, embeddings are precomputed
+            loss = self.loss_fn(y_hat, y)
+
+            val_loss.append(loss.detach())
+        return torch.stack(val_loss).mean().item()
+
+    def _create_data_loaders(
+        self, train_data: tuple[torch.Tensor, torch.Tensor], val_data: tuple[torch.Tensor, torch.Tensor] | None
+    ) -> tuple[DataLoader, DataLoader]:
+        if val_data is None:
+            train_data, val_data = train_val_split(
+                train_data, self.train_config.val_split_ratio, self.train_config.stratify, self.train_config.seed
+            )
+
+        train_embeddings = self.model.emb_model(train_data[0])
+        val_embeddings = self.model.emb_model(val_data[0])
+
+        train_loader = DataLoader(
+            TensorDataset(train_embeddings, train_data[1]),
+            batch_size=self.train_config.batch_size,
+            shuffle=True,
+        )
+
+        val_loader = DataLoader(
+            TensorDataset(val_embeddings, val_data[1]),
+            batch_size=self.train_config.batch_size,
+            shuffle=False,
+        )
+        return train_loader, val_loader
+
+    def _log_epoch_metrics(self, epoch: int, train_loss: float, val_loss: float) -> None:
+        if epoch % self.train_config.log_every_n_steps == 0:
+            print(f"Epoch {epoch + 1}, Train Loss: {train_loss:.6f}, Val Loss: {val_loss:.6f}")

tirex/models/classification/utils.py

@@ -0,0 +1,81 @@
+import numpy as np
+import torch
+from sklearn.model_selection import train_test_split
+
+
+# Remove after Issue will be solved: https://github.com/pytorch/pytorch/issues/61474
+def nanmax(tensor: torch.Tensor, dim: int | None = None, keepdim: bool = False) -> torch.Tensor:
+    min_value = torch.finfo(tensor.dtype).min
+    output = tensor.nan_to_num(min_value).max(dim=dim, keepdim=keepdim)
+    return output.values
+
+
+def nanmin(tensor: torch.Tensor, dim: int | None = None, keepdim: bool = False) -> torch.Tensor:
+    max_value = torch.finfo(tensor.dtype).max
+    output = tensor.nan_to_num(max_value).min(dim=dim, keepdim=keepdim)
+    return output.values
+
+
+def nanvar(tensor: torch.Tensor, dim: int | None = None, keepdim: bool = False) -> torch.Tensor:
+    tensor_mean = tensor.nanmean(dim=dim, keepdim=True)
+    output = (tensor - tensor_mean).square().nanmean(dim=dim, keepdim=keepdim)
+    return output
+
+
+def nanstd(tensor: torch.Tensor, dim: int | None = None, keepdim: bool = False) -> torch.Tensor:
+    output = nanvar(tensor, dim=dim, keepdim=keepdim)
+    output = output.sqrt()
+    return output
+
+
+def train_val_split(
+    train_data: tuple[torch.Tensor, torch.Tensor],
+    val_split_ratio: float,
+    stratify: bool,
+    seed: int | None,
+) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]:
+    idx_train, idx_val = train_test_split(
+        np.arange(len(train_data[0])),
+        test_size=val_split_ratio,
+        random_state=seed,
+        shuffle=True,
+        stratify=train_data[1] if stratify else None,
+    )
+
+    return (
+        (train_data[0][idx_train], train_data[1][idx_train]),
+        (train_data[0][idx_val], train_data[1][idx_val]),
+    )
+
+
+def set_seed(seed: int) -> None:
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+    np.random.seed(seed)
+
+
+class EarlyStopping:
+    def __init__(
+        self,
+        patience: int = 7,
+        delta: float = 0.0001,
+    ) -> None:
+        self.patience: int = patience
+        self.delta: float = delta
+
+        self.best: float = np.inf
+        self.wait_count: int = 0
+        self.early_stop: bool = False
+
+    def __call__(self, epoch: int, val_loss: float) -> bool:
+        improved = val_loss < (self.best - self.delta)
+        if improved:
+            self.best = val_loss
+            self.wait_count = 0
+        else:
+            self.wait_count += 1
+            if self.wait_count >= self.patience:
+                self.early_stop = True
+                print(f"Early stopping triggered at epoch {epoch}.")
+        return self.early_stop

tirex/models/tirex.py

@@ -58,6 +58,25 @@ class TiRexZero(nn.Module, PretrainedModel, ForecastModel):
     def register_name(cls):
         return "TiRex"
 
+    def _adjust_context_length(self, context: torch.Tensor, min_context: int, max_context: int):
+        pad_len = 0
+        if context.shape[-1] > max_context:
+            context = context[..., -max_context:]
+        if context.shape[-1] < min_context:
+            pad_len = min_context - context.shape[-1]
+            pad = torch.full(
+                (context.shape[0], pad_len),
+                fill_value=torch.nan,
+                device=context.device,
+                dtype=context.dtype,
+            )
+            context = torch.concat((pad, context), dim=1)
+        return context, pad_len
+
+    # ===============================
+    #       Forecasting Functions
+    # ===============================
+
     @torch.inference_mode()
     def _forecast_quantiles(
         self,
@@ -106,18 +125,7 @@ class TiRexZero(nn.Module, PretrainedModel, ForecastModel):
 
     def _forecast_single_step(self, context: torch.Tensor, new_patch_count: int = 1) -> torch.Tensor:
         max_context, min_context = self.config.train_ctx_len, self.config.train_ctx_len
-
-        if context.shape[-1] > max_context:
-            context = context[..., -max_context:]
-        if context.shape[-1] < min_context:
-            pad = torch.full(
-                (context.shape[0], min_context - context.shape[-1]),
-                fill_value=torch.nan,
-                device=context.device,
-                dtype=context.dtype,
-            )
-            context = torch.concat((pad, context), dim=1)
-
+        context, _ = self._adjust_context_length(context, max_context, min_context)
         input_token, tokenizer_state = self.tokenizer.input_transform(context)
         prediction = self._forward_model_tokenized(input_token=input_token, new_patch_count=new_patch_count)
         predicted_token = prediction[:, :, -new_patch_count:, :].to(input_token)  # predicted token
@@ -161,16 +169,58 @@ class TiRexZero(nn.Module, PretrainedModel, ForecastModel):
         # quantile_preds: [batch_size, num_quantiles, num_token, output_patch_size]
         return quantile_preds
 
-    def _forward_model(self, input: torch.Tensor) -> torch.Tensor:
+    def _forward_model(self, input: torch.Tensor, return_all_hidden: bool = False) -> torch.Tensor:
         hidden_states = self.input_patch_embedding(input)
+        all_hidden_states = []
 
         for block in self.blocks:
             hidden_states = block(hidden_states)
-
+            if return_all_hidden:
+                all_hidden_states.append(hidden_states)
         hidden_states = self.out_norm(hidden_states)
 
+        if return_all_hidden:
+            return self.output_patch_embedding(hidden_states), torch.stack(all_hidden_states, dim=-2)
+
         return self.output_patch_embedding(hidden_states)
 
+    # ===============================
+    #   Context Embedding Functions
+    # ===============================
+    @torch.inference_mode()
+    def _embed_context(
+        self,
+        context: torch.Tensor,
+        max_context: int | None = None,
+    ) -> torch.Tensor:
+        input_embeds, padded_token = self._prepare_context_for_embedding(context, max_context)
+        _, hidden_states = self._forward_model(input_embeds, return_all_hidden=True)
+        # Shape: [batch_size, num_tokens, num_layers, hidden_dim]
+        return hidden_states[:, padded_token:, :, :]
+
+    def _prepare_context_for_embedding(
+        self, context: torch.Tensor, max_context: int | None
+    ) -> tuple[torch.Tensor, int]:
+        max_context = self.config.train_ctx_len if max_context is None else max_context
+        min_context = max(self.config.train_ctx_len, max_context)
+
+        device = self.input_patch_embedding.hidden_layer.weight.device
+        context = context.to(
+            device=device,
+            dtype=torch.float32,
+        )
+
+        context, pad_len = self._adjust_context_length(context, min_context, max_context)
+
+        padded_token = pad_len // self.tokenizer.patch_size
+        input_token, _ = self.tokenizer.input_transform(context)
+
+        input_mask = torch.isnan(input_token).logical_not().to(input_token.dtype)
+        input_token = torch.nan_to_num(input_token, nan=self.config.nan_mask_value)
+        input_embeds = torch.cat((input_token, input_mask), dim=2)
+
+        return input_embeds, padded_token
+
     def on_load_checkpoint(self, checkpoint: dict) -> None:
         # rename keys of state_dict, because the block_stack was moved directly into the tirex model
         checkpoint["state_dict"] = {k.replace("block_stack.", ""): v for k, v in checkpoint["state_dict"].items()}

{tirex_mirror-2025.11.26.dist-info → tirex_mirror-2025.11.29.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tirex-mirror
-Version: 2025.11.26
+Version: 2025.11.29
 Summary: Unofficial mirror of NX-AI/tirex for packaging
 Author-email: Arpad Rozsas <rozsasarpi@gmail.com>
 License: NXAI COMMUNITY LICENSE AGREEMENT
@@ -84,6 +84,8 @@ Requires-Dist: datasets; extra == "hfdataset"
 Provides-Extra: test
 Requires-Dist: fev>=0.6.0; extra == "test"
 Requires-Dist: pytest; extra == "test"
+Provides-Extra: classification
+Requires-Dist: scikit-learn; extra == "classification"
 Provides-Extra: all
 Requires-Dist: xlstm; extra == "all"
 Requires-Dist: ninja; extra == "all"
@@ -95,6 +97,7 @@ Requires-Dist: gluonts; extra == "all"
 Requires-Dist: datasets; extra == "all"
 Requires-Dist: pytest; extra == "all"
 Requires-Dist: fev>=0.6.0; extra == "all"
+Requires-Dist: scikit-learn; extra == "all"
 Dynamic: license-file
 
 # tirex-mirror

{tirex_mirror-2025.11.26.dist-info → tirex_mirror-2025.11.29.dist-info}/RECORD

@@ -8,14 +8,20 @@ tirex/api_adapter/hf_data.py,sha256=TRyys2xKIGZS0Yhq2Eb61lWCMg5CWWn1yRlLIN1mU7o,
 tirex/api_adapter/standard_adapter.py,sha256=vdlxNs8mTUtPgK_5WMqYqNdMj8W44igqWsAgtggt_xk,2809
 tirex/models/__init__.py,sha256=YnTtPf5jGqvhfqoX8Ku7Yd0xohy0MmocE2ryrXVnQ1Q,135
 tirex/models/patcher.py,sha256=8T4c3PZnOAsEpahhrjtt7S7405WUjN6g3cV33E55PD4,1911
-tirex/models/tirex.py,sha256=URt-MClXu0zdUHACQ96Zu3Ytdb52vbeG_SXnj5C4tI8,7522
+tirex/models/tirex.py,sha256=wtdjrdE1TuoueIvhlKf-deLH3oKyQuGmd2k15My7SWA,9710
+tirex/models/classification/__init__.py,sha256=qVn84uBosWaHm9wr0FoYAXNvmajyyB3_OmpeHNzDH4g,194
+tirex/models/classification/embedding.py,sha256=TJlchUaBhz8Pf1mLpTdDVqmkGVOiaCI55sSoqk2tSXE,5259
+tirex/models/classification/linear_classifier.py,sha256=yE2sekKw2LCLGiITTsTktOJw28RjjOWgsJE2PakAjN8,11142
+tirex/models/classification/rf_classifier.py,sha256=wzMiF5TELGl3V-i94MPOEpNn2yZgrCdPfuzpcl8A18M,5816
+tirex/models/classification/trainer.py,sha256=JzM1XtRoRI3fn6Sbu7V-9IuiKVy454O73uNrMNgCREs,5759
+tirex/models/classification/utils.py,sha256=db9056u6uIVhm0qmHDoOo3K5f-ZiuDytDLcLOg-zFb0,2599
 tirex/models/slstm/block.py,sha256=V91Amgz8WAOOHo4fK1UZxd4Dgbx4-X6kUBS6X4m0tKQ,2006
 tirex/models/slstm/cell.py,sha256=Otyil_AjpJbUckkINWGHxlqP14J5epm_J_zdWPzvD2g,7290
 tirex/models/slstm/layer.py,sha256=hrDydQJIAHf5W0A0Rt0hXG4yKXrOSY-HPL0UbigR6Q8,2867
-tirex_mirror-2025.11.26.dist-info/licenses/LICENSE,sha256=HlwHKnGTlE2oNm6734V-Vy62zlkWohnuZpYXSdkqDk4,7362
-tirex_mirror-2025.11.26.dist-info/licenses/LICENSE_MIRROR.txt,sha256=ulPZMcOZdN7JvISjiID3KUwovTjrPwiMv5ku9dM7nls,496
-tirex_mirror-2025.11.26.dist-info/licenses/NOTICE.txt,sha256=rcgDscFHb-uuZO3L0_vIxYhTYl-a2Rm0lBpp3_kKdFQ,147
-tirex_mirror-2025.11.26.dist-info/METADATA,sha256=t-Qg4PeJHaY2ZGN6CFBSiwz8dig4Ts1kT_qiLlk4xZE,11494
-tirex_mirror-2025.11.26.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tirex_mirror-2025.11.26.dist-info/top_level.txt,sha256=AOLDhfv0F_7nn3pFq0Kapg6Ky_28I_cGDXzQX3w9eO4,6
-tirex_mirror-2025.11.26.dist-info/RECORD,,
+tirex_mirror-2025.11.29.dist-info/licenses/LICENSE,sha256=HlwHKnGTlE2oNm6734V-Vy62zlkWohnuZpYXSdkqDk4,7362
+tirex_mirror-2025.11.29.dist-info/licenses/LICENSE_MIRROR.txt,sha256=ulPZMcOZdN7JvISjiID3KUwovTjrPwiMv5ku9dM7nls,496
+tirex_mirror-2025.11.29.dist-info/licenses/NOTICE.txt,sha256=rcgDscFHb-uuZO3L0_vIxYhTYl-a2Rm0lBpp3_kKdFQ,147
+tirex_mirror-2025.11.29.dist-info/METADATA,sha256=6MeKMufnn5yIl1BKxzQKW22jeR7ZU1U9BEuFd5fBJAg,11624
+tirex_mirror-2025.11.29.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tirex_mirror-2025.11.29.dist-info/top_level.txt,sha256=AOLDhfv0F_7nn3pFq0Kapg6Ky_28I_cGDXzQX3w9eO4,6
+tirex_mirror-2025.11.29.dist-info/RECORD,,