-
-
-
-
-
-
-
-
-import lightning as pl
-import torch
-import torch.nn as nn
-import torchmetrics
-from ..utils.mamba_arch import Mamba
-from ..utils.mlp_utils import MLP
-from ..utils.normalization_layers import (
- RMSNorm,
- LayerNorm,
- LearnableLayerScaling,
- BatchNorm,
- InstanceNorm,
- GroupNorm,
-)
-from ..utils.configs import DefaultMambularConfig
-
-
-[docs]class BaseMambularClassifier(pl.LightningModule):
- """
- A base class for building classification models using the Mambular architecture within the PyTorch Lightning framework.
-
- This class integrates various components such as embeddings for categorical and numerical features, the Mambular model
- for processing sequences of embeddings, and a classification head for prediction. It supports multi-class and binary classification tasks.
-
- Parameters
- ----------
- num_classes : int
- number of classes for classification.
- cat_feature_info : dict
- Dictionary containing information about categorical features.
- num_feature_info : dict
- Dictionary containing information about numerical features.
- config : DefaultMambularConfig, optional
- Configuration object containing default hyperparameters for the model (default is DefaultMambularConfig()).
- **kwargs : dict
- Additional keyword arguments.
-
-
- Attributes
- ----------
- lr : float
- Learning rate.
- lr_patience : int
- Patience for learning rate scheduler.
- weight_decay : float
- Weight decay for optimizer.
- lr_factor : float
- Factor by which the learning rate will be reduced.
- pooling_method : str
- Method to pool the features.
- cat_feature_info : dict
- Dictionary containing information about categorical features.
- num_feature_info : dict
- Dictionary containing information about numerical features.
- embedding_activation : callable
- Activation function for embeddings.
- mamba : Mamba
- Mamba architecture component.
- norm_f : nn.Module
- Normalization layer.
- num_embeddings : nn.ModuleList
- Module list for numerical feature embeddings.
- cat_embeddings : nn.ModuleList
- Module list for categorical feature embeddings.
- tabular_head : MLP
- Multi-layer perceptron head for tabular data.
- cls_token : nn.Parameter
- Class token parameter.
- embedding_norm : nn.Module, optional
- Layer normalization applied after embedding if specified.
- loss_fct : nn.Module
- The loss function used for training the model, configured based on the number of classes.
- acc : torchmetrics.Accuracy
- A metric for computing the accuracy of predictions.
- auroc : torchmetrics.AUROC
- A metric for computing the Area Under the Receiver Operating Characteristic curve.
- precision : torchmetrics.Precision
- A metric for computing the precision of predictions.
-
- """
-
- def __init__(
- self,
- num_classes,
- cat_feature_info,
- num_feature_info,
- config: DefaultMambularConfig = DefaultMambularConfig(),
- **kwargs,
- ):
- super().__init__()
-
- self.num_classes = 1 if num_classes == 2 else num_classes
- # Save all hyperparameters
- self.save_hyperparameters(ignore=["cat_feature_info", "num_feature_info"])
-
- # Assigning values from hyperparameters
- self.lr = self.hparams.get("lr", config.lr)
- self.lr_patience = self.hparams.get("lr_patience", config.lr_patience)
- self.weight_decay = self.hparams.get("weight_decay", config.weight_decay)
- self.lr_factor = self.hparams.get("lr_factor", config.lr_factor)
- self.pooling_method = self.hparams.get("pooling_method", config.pooling_method)
- self.cat_feature_info = cat_feature_info
- self.num_feature_info = num_feature_info
-
- self.embedding_activation = self.hparams.get(
- "num_embedding_activation", config.num_embedding_activation
- )
-
- # Additional layers and components initialization based on hyperparameters
- self.mamba = Mamba(
- d_model=self.hparams.get("d_model", config.d_model),
- n_layers=self.hparams.get("n_layers", config.n_layers),
- expand_factor=self.hparams.get("expand_factor", config.expand_factor),
- bias=self.hparams.get("bias", config.bias),
- d_conv=self.hparams.get("d_conv", config.d_conv),
- conv_bias=self.hparams.get("conv_bias", config.conv_bias),
- dropout=self.hparams.get("dropout", config.dropout),
- dt_rank=self.hparams.get("dt_rank", config.dt_rank),
- d_state=self.hparams.get("d_state", config.d_state),
- dt_scale=self.hparams.get("dt_scale", config.dt_scale),
- dt_init=self.hparams.get("dt_init", config.dt_init),
- dt_max=self.hparams.get("dt_max", config.dt_max),
- dt_min=self.hparams.get("dt_min", config.dt_min),
- dt_init_floor=self.hparams.get("dt_init_floor", config.dt_init_floor),
- norm=globals()[self.hparams.get("norm", config.norm)],
- activation=self.hparams.get("activation", config.activation),
- )
-
- # Set the normalization layer dynamically
- norm_layer = self.hparams.get("norm", config.norm)
- if norm_layer == "RMSNorm":
- self.norm_f = RMSNorm(self.hparams.get("d_model", config.d_model))
- elif norm_layer == "LayerNorm":
- self.norm_f = LayerNorm(self.hparams.get("d_model", config.d_model))
- elif norm_layer == "BatchNorm":
- self.norm_f = BatchNorm(self.hparams.get("d_model", config.d_model))
- elif norm_layer == "InstanceNorm":
- self.norm_f = InstanceNorm(self.hparams.get("d_model", config.d_model))
- elif norm_layer == "GroupNorm":
- self.norm_f = GroupNorm(1, self.hparams.get("d_model", config.d_model))
- elif norm_layer == "LearnableLayerScaling":
- self.norm_f = LearnableLayerScaling(
- self.hparams.get("d_model", config.d_model)
- )
- else:
- raise ValueError(f"Unsupported normalization layer: {norm_layer}")
-
- self.num_embeddings = nn.ModuleList(
- [
- nn.Sequential(
- nn.Linear(
- input_shape,
- self.hparams.get("d_model", config.d_model),
- bias=False,
- ),
- self.embedding_activation,
- )
- for feature_name, input_shape in num_feature_info.items()
- ]
- )
-
- self.cat_embeddings = nn.ModuleList(
- [
- nn.Embedding(
- num_categories + 1, self.hparams.get("d_model", config.d_model)
- )
- for feature_name, num_categories in cat_feature_info.items()
- ]
- )
-
- head_activation = self.hparams.get("head_activation", config.head_activation)
-
- self.tabular_head = MLP(
- self.hparams.get("d_model", config.d_model),
- hidden_units_list=self.hparams.get(
- "head_layer_sizes", config.head_layer_sizes
- ),
- dropout_rate=self.hparams.get("head_dropout", config.head_dropout),
- use_skip_layers=self.hparams.get(
- "head_skip_layers", config.head_skip_layers
- ),
- activation_fn=head_activation,
- use_batch_norm=self.hparams.get(
- "head_use_batch_norm", config.head_use_batch_norm
- ),
- n_output_units=self.num_classes,
- )
-
- self.cls_token = nn.Parameter(
- torch.zeros(1, 1, self.hparams.get("d_model", config.d_model))
- )
-
- self.loss_fct = nn.MSELoss()
-
- if self.hparams.get("layer_norm_after_embedding"):
- self.embedding_norm = nn.LayerNorm(
- self.hparams.get("d_model", config.d_model)
- )
-
- if self.num_classes > 2:
- self.loss_fct = nn.CrossEntropyLoss()
- self.acc = torchmetrics.Accuracy(
- task="multiclass", num_classes=self.num_classes
- )
- self.auroc = torchmetrics.AUROC(
- task="multiclass", num_classes=self.num_classes
- )
- self.precision = torchmetrics.Precision(
- task="multiclass", num_classes=self.num_classes
- )
- else:
- self.loss_fct = torch.nn.BCEWithLogitsLoss()
- self.acc = torchmetrics.Accuracy(task="binary")
- self.auroc = torchmetrics.AUROC(task="binary")
- self.precision = torchmetrics.Precision(task="binary")
-
-[docs] def forward(self, num_features, cat_features):
- """
- Defines the forward pass of the classifier.
-
- Parameters
- ----------
- cat_features : Tensor
- Tensor containing the categorical features.
- num_features : Tensor
- Tensor containing the numerical features.
-
-
- Returns
- -------
- Tensor
- The output predictions of the model.
- """
- batch_size = (
- cat_features[0].size(0) if cat_features != [] else num_features[0].size(0)
- )
- cls_tokens = self.cls_token.expand(batch_size, -1, -1)
-
- # Process categorical features if present
- if len(self.cat_embeddings) > 0 and cat_features:
- cat_embeddings = [
- emb(cat_features[i]) for i, emb in enumerate(self.cat_embeddings)
- ]
- cat_embeddings = torch.stack(cat_embeddings, dim=1)
- cat_embeddings = torch.squeeze(cat_embeddings, dim=2)
- if self.hparams.get("layer_norm_after_embedding"):
- cat_embeddings = self.embedding_norm(cat_embeddings)
- else:
- cat_embeddings = None
-
- # Process numerical features if present
- if len(self.num_embeddings) > 0 and num_features:
- num_embeddings = [
- emb(num_features[i]) for i, emb in enumerate(self.num_embeddings)
- ]
- num_embeddings = torch.stack(num_embeddings, dim=1)
- if self.hparams.get("layer_norm_after_embedding"):
- num_embeddings = self.embedding_norm(num_embeddings)
- else:
- num_embeddings = None
-
- # Combine embeddings if both types are present, otherwise use whichever is available
- if cat_embeddings is not None and num_embeddings is not None:
- x = torch.cat([cls_tokens, cat_embeddings, num_embeddings], dim=1)
- elif cat_embeddings is not None:
- x = torch.cat([cls_tokens, cat_embeddings], dim=1)
- elif num_embeddings is not None:
- x = torch.cat([cls_tokens, num_embeddings], dim=1)
- else:
- raise ValueError("No features provided to the model.")
-
- x = self.mamba(x)
-
- # Apply pooling based on the specified method
- if self.pooling_method == "avg":
- x = torch.mean(x, dim=1)
- elif self.pooling_method == "max":
- x, _ = torch.max(x, dim=1)
- elif self.pooling_method == "sum":
- x = torch.sum(x, dim=1)
- elif self.pooling_method == "cls":
- x = x[:, 0]
- else:
- raise ValueError(f"Invalid pooling method: {self.pooling_method}")
-
- x = self.norm_f(x)
- preds = self.tabular_head(x)
- return preds
-
-[docs] def training_step(self, batch, batch_idx):
- """
- Processes a single batch during training, computes the loss and logs training metrics.
-
- Parameters
- ----------
- batch : tuple
- A batch of data from the DataLoader, containing numerical features, categorical features, and labels.
- batch_idx : int
- The index of the batch within the epoch.
- """
-
- cat_features, num_features, labels = batch
- preds = self(num_features=num_features, cat_features=cat_features)
-
- if self.num_classes == 1:
- labels = torch.float(
- labels.unsqueeze(1), dtype=torch.float32
- ) # Reshape for binary classification loss calculation
-
- loss = self.loss_fct(preds, labels)
- self.log("train_loss", loss)
- # Calculate and log training accuracy
-
- acc = self.acc(preds, labels.int())
- self.log(
- "train_acc",
- acc,
- on_step=False,
- on_epoch=True,
- prog_bar=True,
- logger=True,
- )
-
- # Calculate and log AUROC
- auroc = self.auroc(preds, labels.int())
- self.log(
- "train_auroc",
- auroc,
- on_step=False,
- on_epoch=True,
- prog_bar=True,
- logger=True,
- )
-
- # Calculate and log precision
- precision = self.precision(preds, labels.int())
- self.log(
- "train_precision",
- precision,
- on_step=False,
- on_epoch=True,
- prog_bar=True,
- logger=True,
- )
-
- return loss
-
-[docs] def validation_step(self, batch, batch_idx):
- """
- Processes a single batch during validation, computes the loss and logs validation metrics.
-
- Parameters
- ----------
- batch : tuple
- A batch of data from the DataLoader, containing numerical features, categorical features, and labels.
- batch_idx : int
- The index of the batch within the epoch.
- """
- cat_features, num_features, labels = batch
- preds = self(num_features=num_features, cat_features=cat_features)
-
- if self.num_classes == 1:
- labels = labels.unsqueeze(
- 1
- ).float() # Reshape for binary classification loss calculation
-
- loss = self.loss_fct(preds, labels)
- self.log("val_loss", loss)
- # Calculate and log training accuracy
-
- acc = self.acc(preds, labels.int())
- self.log(
- "val_acc",
- acc,
- on_step=False,
- on_epoch=True,
- prog_bar=True,
- logger=True,
- )
-
- auroc = self.auroc(preds, labels.int())
- self.log(
- "val_auroc", auroc, on_step=False, on_epoch=True, prog_bar=True, logger=True
- )
-
- # Calculate and log precision
- precision = self.precision(preds, labels.int())
- self.log(
- "val_precision",
- precision,
- on_step=False,
- on_epoch=True,
- prog_bar=True,
- logger=True,
- )
-
-[docs] def configure_optimizers(self):
- """
- Sets up the model's optimizer and learning rate scheduler based on the configurations provided.
-
- Returns
- -------
- dict
- A dictionary containing the optimizer and lr_scheduler configurations.
- """
- optimizer = torch.optim.Adam(
- self.parameters(), lr=self.lr, weight_decay=self.weight_decay
- )
- scheduler = {
- "scheduler": torch.optim.lr_scheduler.ReduceLROnPlateau(
- optimizer,
- mode="min",
- factor=self.lr_factor,
- patience=self.lr_patience,
- verbose=True,
- ),
- "monitor": "val_loss", # Name of the metric to monitor
- "interval": "epoch",
- "frequency": 1,
- }
-
- return {"optimizer": optimizer, "lr_scheduler": scheduler}
-
-
-
-
-
-import lightning as pl
-import torch
-import torch.nn as nn
-
-from ..utils.configs import DefaultMambularConfig
-from ..utils.mlp_utils import MLP
-from ..utils.distributions import (
- BetaDistribution,
- CategoricalDistribution,
- DirichletDistribution,
- GammaDistribution,
- InverseGammaDistribution,
- NegativeBinomialDistribution,
- NormalDistribution,
- PoissonDistribution,
- StudentTDistribution,
-)
-from ..utils.normalization_layers import (
- RMSNorm,
- LayerNorm,
- LearnableLayerScaling,
- BatchNorm,
- InstanceNorm,
- GroupNorm,
-)
-from ..utils.mamba_arch import Mamba
-
-
-[docs]class BaseMambularLSS(pl.LightningModule):
- """
- A base module for likelihood-based statistical learning (LSS) models built on PyTorch Lightning,
- integrating the Mamba architecture for tabular data. This module is designed to accommodate various
- statistical distribution families for different types of regression and classification tasks.
-
-
- Parameters
- ----------
- family : str
- The name of the statistical distribution family to be used for modeling. Supported families include
- 'normal', 'poisson', 'gamma', 'beta', 'dirichlet', 'studentt', 'negativebinom', 'inversegamma', and 'categorical'.
- cat_feature_info : dict
- Dictionary containing information about categorical features.
- num_feature_info : dict
- Dictionary containing information about numerical features.
- config : DefaultMambularConfig, optional
- Configuration object containing default hyperparameters for the model (default is DefaultMambularConfig()).
- **kwargs : dict
- Additional keyword arguments.
-
-
- Attributes
- ----------
- lr : float
- Learning rate.
- lr_patience : int
- Patience for learning rate scheduler.
- weight_decay : float
- Weight decay for optimizer.
- lr_factor : float
- Factor by which the learning rate will be reduced.
- pooling_method : str
- Method to pool the features.
- cat_feature_info : dict
- Dictionary containing information about categorical features.
- num_feature_info : dict
- Dictionary containing information about numerical features.
- embedding_activation : callable
- Activation function for embeddings.
- mamba : Mamba
- Mamba architecture component.
- norm_f : nn.Module
- Normalization layer.
- num_embeddings : nn.ModuleList
- Module list for numerical feature embeddings.
- cat_embeddings : nn.ModuleList
- Module list for categorical feature embeddings.
- tabular_head : MLP
- Multi-layer perceptron head for tabular data.
- cls_token : nn.Parameter
- Class token parameter.
- embedding_norm : nn.Module, optional
- Layer normalization applied after embedding if specified.
- """
-
- def __init__(
- self,
- family,
- cat_feature_info,
- num_feature_info,
- config: DefaultMambularConfig = DefaultMambularConfig(),
- distributional_kwargs=None,
- **kwargs,
- ):
- super().__init__()
-
- # Save all hyperparameters
- self.save_hyperparameters(ignore=["cat_feature_info", "num_feature_info"])
-
- # Assigning values from hyperparameters
- self.lr = self.hparams.get("lr", config.lr)
- self.lr_patience = self.hparams.get("lr_patience", config.lr_patience)
- self.weight_decay = self.hparams.get("weight_decay", config.weight_decay)
- self.lr_factor = self.hparams.get("lr_factor", config.lr_factor)
- self.pooling_method = self.hparams.get("pooling_method", config.pooling_method)
- self.cat_feature_info = cat_feature_info
- self.num_feature_info = num_feature_info
-
- self.embedding_activation = self.hparams.get(
- "num_embedding_activation", config.num_embedding_activation
- )
-
- distribution_classes = {
- "normal": NormalDistribution,
- "poisson": PoissonDistribution,
- "gamma": GammaDistribution,
- "beta": BetaDistribution,
- "dirichlet": DirichletDistribution,
- "studentt": StudentTDistribution,
- "negativebinom": NegativeBinomialDistribution,
- "inversegamma": InverseGammaDistribution,
- "categorical": CategoricalDistribution,
- }
-
- if distributional_kwargs is None:
- distributional_kwargs = {}
-
- if family in distribution_classes:
- self.family = distribution_classes[family](**distributional_kwargs)
- else:
- raise ValueError("Unsupported family: {}".format(family))
-
- # Set the normalization layer dynamically
- norm_layer = self.hparams.get("norm", config.norm)
- if norm_layer == "RMSNorm":
- self.norm_f = RMSNorm(self.hparams.get("d_model", config.d_model))
- elif norm_layer == "LayerNorm":
- self.norm_f = LayerNorm(self.hparams.get("d_model", config.d_model))
- elif norm_layer == "BatchNorm":
- self.norm_f = BatchNorm(self.hparams.get("d_model", config.d_model))
- elif norm_layer == "InstanceNorm":
- self.norm_f = InstanceNorm(self.hparams.get("d_model", config.d_model))
- elif norm_layer == "GroupNorm":
- self.norm_f = GroupNorm(1, self.hparams.get("d_model", config.d_model))
- elif norm_layer == "LearnableLayerScaling":
- self.norm_f = LearnableLayerScaling(
- self.hparams.get("d_model", config.d_model)
- )
- else:
- raise ValueError(f"Unsupported normalization layer: {norm_layer}")
-
- # Additional layers and components initialization based on hyperparameters
- self.mamba = Mamba(
- d_model=self.hparams.get("d_model", config.d_model),
- n_layers=self.hparams.get("n_layers", config.n_layers),
- expand_factor=self.hparams.get("expand_factor", config.expand_factor),
- bias=self.hparams.get("bias", config.bias),
- d_conv=self.hparams.get("d_conv", config.d_conv),
- conv_bias=self.hparams.get("conv_bias", config.conv_bias),
- dropout=self.hparams.get("dropout", config.dropout),
- dt_rank=self.hparams.get("dt_rank", config.dt_rank),
- d_state=self.hparams.get("d_state", config.d_state),
- dt_scale=self.hparams.get("dt_scale", config.dt_scale),
- dt_init=self.hparams.get("dt_init", config.dt_init),
- dt_max=self.hparams.get("dt_max", config.dt_max),
- dt_min=self.hparams.get("dt_min", config.dt_min),
- dt_init_floor=self.hparams.get("dt_init_floor", config.dt_init_floor),
- norm=globals()[self.hparams.get("norm", config.norm)],
- activation=self.hparams.get("activation", config.activation),
- )
-
- self.num_embeddings = nn.ModuleList(
- [
- nn.Sequential(
- nn.Linear(
- input_shape,
- self.hparams.get("d_model", config.d_model),
- bias=False,
- ),
- self.embedding_activation,
- )
- for feature_name, input_shape in num_feature_info.items()
- ]
- )
-
- self.cat_embeddings = nn.ModuleList(
- [
- nn.Embedding(
- num_categories + 1, self.hparams.get("d_model", config.d_model)
- )
- for feature_name, num_categories in cat_feature_info.items()
- ]
- )
-
- head_activation = self.hparams.get("head_activation", config.head_activation)
-
- self.tabular_head = MLP(
- self.hparams.get("d_model", config.d_model),
- hidden_units_list=self.hparams.get(
- "head_layer_sizes", config.head_layer_sizes
- ),
- dropout_rate=self.hparams.get("head_dropout", config.head_dropout),
- use_skip_layers=self.hparams.get(
- "head_skip_layers", config.head_skip_layers
- ),
- activation_fn=head_activation,
- use_batch_norm=self.hparams.get(
- "head_use_batch_norm", config.head_use_batch_norm
- ),
- n_output_units=self.family.param_count,
- )
-
- self.cls_token = nn.Parameter(
- torch.zeros(1, 1, self.hparams.get("d_model", config.d_model))
- )
-
- self.loss_fct = nn.MSELoss()
-
- if self.hparams.get("layer_norm_after_embedding"):
- self.embedding_norm = nn.LayerNorm(
- self.hparams.get("d_model", config.d_model)
- )
-
- def compute_loss(self, predictions, y_true):
- return self.family.compute_loss(predictions, y_true)
-
-[docs] def forward(self, cat_features, num_features):
- """
- Defines the forward pass of the model, processing both categorical and numerical features,
- and returning predictions based on the configured statistical distribution.
-
- Parameters
- ----------
- cat_features : Tensor
- Tensor containing the categorical features.
- num_features : Tensor
- Tensor containing the numerical features.
-
-
- Returns
- -------
- Tensor
- The predictions of the model, typically the parameters of the chosen statistical distribution.
- """
-
- batch_size = (
- cat_features[0].size(0) if cat_features != [] else num_features[0].size(0)
- )
- cls_tokens = self.cls_token.expand(batch_size, -1, -1)
-
- # Process categorical features if present
- if len(self.cat_embeddings) > 0 and cat_features:
- cat_embeddings = [
- emb(cat_features[i]) for i, emb in enumerate(self.cat_embeddings)
- ]
- cat_embeddings = torch.stack(cat_embeddings, dim=1)
- cat_embeddings = torch.squeeze(cat_embeddings, dim=2)
- if self.hparams.get("layer_norm_after_embedding"):
- cat_embeddings = self.embedding_norm(cat_embeddings)
- else:
- cat_embeddings = None
-
- # Process numerical features if present
- if len(self.num_embeddings) > 0 and num_features:
- num_embeddings = [
- emb(num_features[i]) for i, emb in enumerate(self.num_embeddings)
- ]
- num_embeddings = torch.stack(num_embeddings, dim=1)
- if self.hparams.get("layer_norm_after_embedding"):
- num_embeddings = self.embedding_norm(num_embeddings)
- else:
- num_embeddings = None
-
- # Combine embeddings if both types are present, otherwise use whichever is available
- if cat_embeddings is not None and num_embeddings is not None:
- x = torch.cat([cls_tokens, cat_embeddings, num_embeddings], dim=1)
- elif cat_embeddings is not None:
- x = torch.cat([cls_tokens, cat_embeddings], dim=1)
- elif num_embeddings is not None:
- x = torch.cat([cls_tokens, num_embeddings], dim=1)
- else:
- raise ValueError("No features provided to the model.")
-
- x = self.mamba(x)
-
- # Apply pooling based on the specified method
- if self.pooling_method == "avg":
- x = torch.mean(x, dim=1)
- elif self.pooling_method == "max":
- x, _ = torch.max(x, dim=1)
- elif self.pooling_method == "sum":
- x = torch.sum(x, dim=1)
- elif self.pooling_method == "cls":
- x = x[:, 0]
- else:
- raise ValueError(f"Invalid pooling method: {self.pooling_method}")
-
- x = self.norm_f(x)
- preds = self.tabular_head(x)
-
- return preds
-
-[docs] def training_step(self, batch, batch_idx):
- """
- Processes a single batch during training, computes the loss using the distribution-specific loss function,
- and logs training metrics.
-
- Parameters
- ----------
- batch : tuple
- A batch of data from the DataLoader, containing numerical features, categorical features, and labels.
- batch_idx : int
- The index of the batch within the epoch.
-
- Returns
- -------
- Tensor
- The computed loss for the batch.
- """
- cat_features, num_features, labels = batch
- preds = self(num_features=num_features, cat_features=cat_features)
-
- loss = self.compute_loss(preds, labels)
- self.log(
- "train_loss",
- loss,
- on_epoch=True,
- prog_bar=True,
- logger=True,
- )
- return loss
-
-[docs] def validation_step(self, batch, batch_idx):
- """
- Processes a single batch during validation, computes the loss using the distribution-specific loss function,
- and logs validation metrics.
-
- Parameters
- ----------
- batch : tuple
- A batch of data from the DataLoader, containing numerical features, categorical features, and labels.
- batch_idx : int
- The index of the batch within the epoch.
- """
-
- cat_features, num_features, labels = batch
- preds = self(num_features=num_features, cat_features=cat_features)
-
- loss = self.compute_loss(preds, labels)
- self.log(
- "val_loss",
- loss,
- on_epoch=True,
- prog_bar=True,
- logger=True,
- )
- return loss
-
-[docs] def configure_optimizers(self):
- """
- Sets up the model's optimizer and learning rate scheduler based on the configurations provided.
-
- Returns
- -------
- dict
- A dictionary containing the optimizer and lr_scheduler configurations.
- """
- optimizer = torch.optim.Adam(
- self.parameters(), lr=self.lr, weight_decay=self.weight_decay
- )
- scheduler = {
- "scheduler": torch.optim.lr_scheduler.ReduceLROnPlateau(
- optimizer,
- mode="min",
- factor=self.lr_factor,
- patience=self.lr_patience,
- verbose=True,
- ),
- "monitor": "val_loss", # Name of the metric to monitor
- "interval": "epoch",
- "frequency": 1,
- }
-
- return {"optimizer": optimizer, "lr_scheduler": scheduler}
-
-
-
-
-
-import lightning as pl
-import torch
-import torch.nn as nn
-import torchmetrics
-
-from ..utils.mamba_arch import Mamba
-from ..utils.mlp_utils import MLP
-from ..utils.normalization_layers import (
- RMSNorm,
- LayerNorm,
- LearnableLayerScaling,
- BatchNorm,
- InstanceNorm,
- GroupNorm,
-)
-from ..utils.configs import DefaultMambularConfig
-
-
-[docs]class BaseEmbeddingMambularClassifier(pl.LightningModule):
- """
- A specialized classification module for protein data, built on PyTorch Lightning and integrating the Mamba architecture.
- It supports embeddings for categorical features and can process raw or embedded numerical features, making it suitable
- for complex protein sequence data.
-
- Parameters
- ----------
- num_classes : int
- number of classes for classification.
- cat_feature_info : dict
- Dictionary containing information about categorical features.
- num_feature_info : dict
- Dictionary containing information about numerical features.
- config : DefaultMambularConfig, optional
- Configuration object containing default hyperparameters for the model (default is DefaultMambularConfig()).
- seq_size : int, optional
- Size of sequence chunks for processing numerical features. Relevant when `raw_embeddings` is False.
- raw_embeddings : bool, optional
- Indicates whether to use raw numerical features directly or to process them into embeddings. Defaults to False.
-
-
- Attributes
- ----------
- lr : float
- Learning rate.
- lr_patience : int
- Patience for learning rate scheduler.
- weight_decay : float
- Weight decay for optimizer.
- lr_factor : float
- Factor by which the learning rate will be reduced.
- pooling_method : str
- Method to pool the features.
- cat_feature_info : dict
- Dictionary containing information about categorical features.
- num_feature_info : dict
- Dictionary containing information about numerical features.
- embedding_activation : callable
- Activation function for embeddings.
- mamba : Mamba
- Mamba architecture component.
- norm_f : nn.Module
- Normalization layer.
- num_embeddings : nn.ModuleList
- Module list for numerical feature embeddings.
- cat_embeddings : nn.ModuleList
- Module list for categorical feature embeddings.
- tabular_head : MLP
- Multi-layer perceptron head for tabular data.
- cls_token : nn.Parameter
- Class token parameter.
- embedding_norm : nn.Module, optional
- Layer normalization applied after embedding if specified.
- loss_fct : nn.Module
- The loss function used for training the model, configured based on the number of classes.
- acc : torchmetrics.Accuracy
- A metric for computing the accuracy of predictions.
- auroc : torchmetrics.AUROC
- A metric for computing the Area Under the Receiver Operating Characteristic curve.
- precision : torchmetrics.Precision
- A metric for computing the precision of predictions.
-
- """
-
- def __init__(
- self,
- num_classes,
- cat_feature_info,
- num_feature_info,
- config: DefaultMambularConfig = DefaultMambularConfig(),
- seq_size: int = 20,
- raw_embeddings=False,
- **kwargs,
- ):
- super().__init__()
-
- self.num_classes = 1 if num_classes == 2 else num_classes
- # Save all hyperparameters
- self.save_hyperparameters(ignore=["cat_feature_info", "num_feature_info"])
-
- # Assigning values from hyperparameters
- self.lr = self.hparams.get("lr", config.lr)
- self.lr_patience = self.hparams.get("lr_patience", config.lr_patience)
- self.weight_decay = self.hparams.get("weight_decay", config.weight_decay)
- self.lr_factor = self.hparams.get("lr_factor", config.lr_factor)
- self.pooling_method = self.hparams.get("pooling_method", config.pooling_method)
- self.cat_feature_info = cat_feature_info
- self.num_feature_info = num_feature_info
-
- self.embedding_activation = self.hparams.get(
- "num_embedding_activation", config.num_embedding_activation
- )
- self.seq_size = seq_size
- self.raw_embeddings = raw_embeddings
-
- # Additional layers and components initialization based on hyperparameters
- self.mamba = Mamba(
- d_model=self.hparams.get("d_model", config.d_model),
- n_layers=self.hparams.get("n_layers", config.n_layers),
- expand_factor=self.hparams.get("expand_factor", config.expand_factor),
- bias=self.hparams.get("bias", config.bias),
- d_conv=self.hparams.get("d_conv", config.d_conv),
- conv_bias=self.hparams.get("conv_bias", config.conv_bias),
- dropout=self.hparams.get("dropout", config.dropout),
- dt_rank=self.hparams.get("dt_rank", config.dt_rank),
- d_state=self.hparams.get("d_state", config.d_state),
- dt_scale=self.hparams.get("dt_scale", config.dt_scale),
- dt_init=self.hparams.get("dt_init", config.dt_init),
- dt_max=self.hparams.get("dt_max", config.dt_max),
- dt_min=self.hparams.get("dt_min", config.dt_min),
- dt_init_floor=self.hparams.get("dt_init_floor", config.dt_init_floor),
- norm=globals()[self.hparams.get("norm", config.norm)],
- activation=self.hparams.get("activation", config.activation),
- )
-
- # Set the normalization layer dynamically
- norm_layer = self.hparams.get("norm", config.norm)
- if norm_layer == "RMSNorm":
- self.norm_f = RMSNorm(self.hparams.get("d_model", config.d_model))
- elif norm_layer == "LayerNorm":
- self.norm_f = LayerNorm(self.hparams.get("d_model", config.d_model))
- elif norm_layer == "BatchNorm":
- self.norm_f = BatchNorm(self.hparams.get("d_model", config.d_model))
- elif norm_layer == "InstanceNorm":
- self.norm_f = InstanceNorm(self.hparams.get("d_model", config.d_model))
- elif norm_layer == "GroupNorm":
- self.norm_f = GroupNorm(1, self.hparams.get("d_model", config.d_model))
- elif norm_layer == "LearnableLayerScaling":
- self.norm_f = LearnableLayerScaling(
- self.hparams.get("d_model", config.d_model)
- )
- else:
- raise ValueError(f"Unsupported normalization layer: {norm_layer}")
-
- if not self.raw_embeddings:
- data_size = len(num_feature_info.items())
- num_embedding_modules = data_size // self.seq_size
- self.num_embeddings = nn.ModuleList(
- [
- nn.Sequential(
- nn.Linear(
- self.seq_size,
- self.hparams.get("d_model", config.d_model),
- bias=False,
- ),
- # Example using ReLU as the activation function, change as needed
- self.embedding_activation,
- )
- for _ in range(num_embedding_modules)
- ]
- )
- else:
- data_size = len(num_feature_info.items())
- num_embedding_modules = data_size
- self.num_embeddings = nn.ModuleList(
- [
- nn.Sequential(
- nn.Linear(
- input_shape,
- self.hparams.get("d_model", config.d_model),
- bias=False,
- ),
- # Example using ReLU as the activation function, change as needed
- self.embedding_activation,
- )
- for feature_name, input_shape in num_feature_info.items()
- ]
- )
-
- self.cat_embeddings = nn.ModuleList(
- [
- nn.Embedding(
- num_categories + 1, self.hparams.get("d_model", config.d_model)
- )
- for feature_name, num_categories in cat_feature_info.items()
- ]
- )
-
- head_activation = self.hparams.get("head_activation", config.head_activation)
-
- self.tabular_head = MLP(
- self.hparams.get("d_model", config.d_model),
- hidden_units_list=self.hparams.get(
- "head_layer_sizes", config.head_layer_sizes
- ),
- dropout_rate=self.hparams.get("head_dropout", config.head_dropout),
- use_skip_layers=self.hparams.get(
- "head_skip_layers", config.head_skip_layers
- ),
- activation_fn=head_activation,
- use_batch_norm=self.hparams.get(
- "head_use_batch_norm", config.head_use_batch_norm
- ),
- n_output_units=self.num_classes,
- )
-
- self.cls_token = nn.Parameter(
- torch.zeros(1, 1, self.hparams.get("d_model", config.d_model))
- )
-
- self.loss_fct = nn.MSELoss()
-
- if self.hparams.get("layer_norm_after_embedding"):
- self.embedding_norm = nn.LayerNorm(
- self.hparams.get("d_model", config.d_model)
- )
-
- if self.num_classes > 2:
- self.loss_fct = nn.CrossEntropyLoss()
- self.acc = torchmetrics.Accuracy(
- task="multiclass", num_classes=self.num_classes
- )
- self.auroc = torchmetrics.AUROC(
- task="multiclass", num_classes=self.num_classes
- )
- self.precision = torchmetrics.Precision(
- task="multiclass", num_classes=self.num_classes
- )
- else:
- self.loss_fct = torch.nn.BCEWithLogitsLoss()
- self.acc = torchmetrics.Accuracy(task="binary")
- self.auroc = torchmetrics.AUROC(task="binary")
- self.precision = torchmetrics.Precision(task="binary")
-
-[docs] def forward(self, cat_features, num_features):
- """
- Defines the forward pass of the model, processing both categorical and numerical features,
- and returning regression predictions.
-
- Parameters
- ----------
- cat_features : Tensor
- Tensor containing the categorical features.
- num_features : Tensor
- Tensor containing the numerical features or raw sequence data, depending on `raw_embeddings`.
-
-
- Returns
- -------
- Tensor
- The output predictions of the model for regression tasks.
- """
- batch_size = (
- cat_features[0].size(0) if cat_features != [] else num_features[0].size(0)
- )
- cls_tokens = self.cls_token.expand(batch_size, -1, -1)
- # Process categorical features if present
- if not self.raw_embeddings:
- if len(self.cat_embeddings) > 0 and cat_features:
- cat_embeddings = [
- emb(cat_features[i]) for i, emb in enumerate(self.cat_embeddings)
- ]
- cat_embeddings = torch.stack(cat_embeddings, dim=1)
- cat_embeddings = torch.squeeze(cat_embeddings, dim=2)
- else:
- cat_embeddings = None
-
- if len(self.num_embeddings) > 0 and num_features:
- num_embeddings = []
- # Iterate through the num_embeddings, taking slices of num_features for each
- for i, emb in enumerate(self.num_embeddings):
- # Calculate start and end indices for slicing the list
- start_idx = i * self.seq_size
- end_idx = start_idx + self.seq_size
-
- # Slice the num_features list to get the current chunk
- current_chunk = num_features[start_idx:end_idx]
-
- # If the current_chunk is not empty, process it
- if current_chunk:
- # Concatenate tensors in the current chunk along dimension 1
- chunk_tensor = torch.cat(current_chunk, dim=1)
- # Apply the embedding layer to the chunk_tensor
- num_embeddings.append(emb(chunk_tensor))
-
- # Stack the resulting embeddings along the second dimension if num_embeddings is not empty
- if num_embeddings:
- num_embeddings = torch.stack(num_embeddings, dim=1)
- else:
- num_embeddings = None
-
- else:
- # Process categorical features if present
- if len(self.cat_embeddings) > 0 and cat_features:
- cat_embeddings = [
- emb(cat_features[i]) for i, emb in enumerate(self.cat_embeddings)
- ]
- cat_embeddings = torch.stack(cat_embeddings, dim=1)
- cat_embeddings = torch.squeeze(cat_embeddings, dim=2)
- if self.config.layer_norm_after_embedding:
- cat_embeddings = self.embedding_norm(cat_embeddings)
- else:
- cat_embeddings = None
-
- # Process numerical features if present
- if len(self.num_embeddings) > 0 and num_features:
- num_embeddings = [
- emb(num_features[i]) for i, emb in enumerate(self.num_embeddings)
- ]
- num_embeddings = torch.stack(num_embeddings, dim=1)
- if self.config.layer_norm_after_embedding:
- num_embeddings = self.embedding_norm(num_embeddings)
- else:
- num_embeddings = None
-
- # Combine embeddings if both types are present, otherwise use whichever is available
- if cat_embeddings is not None and num_embeddings is not None:
- x = torch.cat([cls_tokens, cat_embeddings, num_embeddings], dim=1)
- elif cat_embeddings is not None:
- x = torch.cat([cls_tokens, cat_embeddings], dim=1)
- elif num_embeddings is not None:
- x = torch.cat([cls_tokens, num_embeddings], dim=1)
- else:
- raise ValueError("No features provided to the model.")
-
- x = self.mamba(x)
-
- # Apply pooling based on the specified method
- if self.pooling_method == "avg":
- x = torch.mean(x, dim=1)
- elif self.pooling_method == "max":
- x, _ = torch.max(x, dim=1)
- elif self.pooling_method == "sum":
- x = torch.sum(x, dim=1)
- elif self.pooling_method == "cls_token":
- x = x[:, 0]
- else:
- raise ValueError(f"Invalid pooling method: {self.pooling_method}")
-
- x = self.norm_f(x)
- preds = self.tabular_head(x)
-
- return preds
-
-[docs] def training_step(self, batch, batch_idx):
- """
- Processes a single batch during training, computes the loss, and logs training metrics.
-
- Parameters
- ----------
- batch : tuple
- A batch of data from the DataLoader, containing numerical features, categorical features, and labels.
- batch_idx : int
- The index of the batch within the epoch.
-
-
- Returns
- -------
- Tensor
- The computed loss for the batch.
- """
- cat_features, num_features, labels = batch
- preds = self(num_features=num_features, cat_features=cat_features)
-
- if self.num_classes == 1:
- labels = torch.float(
- labels.unsqueeze(1), dtype=torch.float32
- ) # Reshape for binary classification loss calculation
-
- loss = self.loss_fct(preds, labels)
- self.log("train_loss", loss)
- # Calculate and log training accuracy
-
- acc = self.acc(preds, labels.int())
- self.log(
- "train_acc",
- acc,
- on_step=False,
- on_epoch=True,
- prog_bar=True,
- logger=True,
- )
-
- # Calculate and log AUROC
- auroc = self.auroc(preds, labels.int())
- self.log(
- "train_auroc",
- auroc,
- on_step=False,
- on_epoch=True,
- prog_bar=True,
- logger=True,
- )
-
- # Calculate and log precision
- precision = self.precision(preds, labels.int())
- self.log(
- "train_precision",
- precision,
- on_step=False,
- on_epoch=True,
- prog_bar=True,
- logger=True,
- )
-
- return loss
-
-[docs] def validation_step(self, batch, batch_idx):
- """
- Processes a single batch during validation, computes the loss, and logs validation metrics.
-
- Parameters
- ----------
- batch : tuple
- A batch of data from the DataLoader, containing numerical features, categorical features, and labels.
- batch_idx : int
- The index of the batch within the epoch.
- """
- cat_features, num_features, labels = batch
- preds = self(num_features=num_features, cat_features=cat_features)
-
- if self.num_classes == 1:
- labels = labels.unsqueeze(
- 1
- ).float() # Reshape for binary classification loss calculation
-
- loss = self.loss_fct(preds, labels)
- self.log("val_loss", loss)
- # Calculate and log training accuracy
-
- acc = self.acc(preds, labels.int())
- self.log(
- "val_acc",
- acc,
- on_step=False,
- on_epoch=True,
- prog_bar=True,
- logger=True,
- )
-
- auroc = self.auroc(preds, labels.int())
- self.log(
- "val_auroc", auroc, on_step=False, on_epoch=True, prog_bar=True, logger=True
- )
-
- # Calculate and log precision
- precision = self.precision(preds, labels.int())
- self.log(
- "val_precision",
- precision,
- on_step=False,
- on_epoch=True,
- prog_bar=True,
- logger=True,
- )
-
-[docs] def configure_optimizers(self):
- """
- Sets up the model's optimizer and learning rate scheduler based on the configurations provided.
-
- Returns
- -------
- dict
- A dictionary containing the optimizer and lr_scheduler configurations.
- """
- optimizer = torch.optim.Adam(
- self.parameters(), lr=self.lr, weight_decay=self.weight_decay
- )
- scheduler = {
- "scheduler": torch.optim.lr_scheduler.ReduceLROnPlateau(
- optimizer,
- mode="min",
- factor=self.lr_factor,
- patience=self.lr_patience,
- verbose=True,
- ),
- "monitor": "val_loss", # Name of the metric to monitor
- "interval": "epoch",
- "frequency": 1,
- }
-
- return {"optimizer": optimizer, "lr_scheduler": scheduler}
-
-
-
-
-
-import lightning as pl
-import torch
-import torch.nn as nn
-from ..utils.normalization_layers import (
- RMSNorm,
- LayerNorm,
- LearnableLayerScaling,
- BatchNorm,
- InstanceNorm,
- GroupNorm,
-)
-
-from ..utils.configs import DefaultMambularConfig
-from ..utils.mamba_arch import Mamba
-from ..utils.mlp_utils import MLP
-
-
-[docs]class BaseEmbeddingMambularRegressor(pl.LightningModule):
- """
- A specialized regression module for protein data, built on PyTorch Lightning and integrating the Mamba architecture.
- It supports embeddings for categorical features and can process raw or embedded numerical features, making it suitable
- for complex protein sequence data.
-
- Parameters
- ----------
- cat_feature_info : dict
- Dictionary containing information about categorical features.
- num_feature_info : dict
- Dictionary containing information about numerical features.
- config : DefaultMambularConfig, optional
- Configuration object containing default hyperparameters for the model (default is DefaultMambularConfig()).
- seq_size : int, optional
- Size of sequence chunks for processing numerical features. Relevant when `raw_embeddings` is False.
- raw_embeddings : bool, optional
- Indicates whether to use raw numerical features directly or to process them into embeddings. Defaults to False.
-
-
- Attributes
- ----------
- lr : float
- Learning rate.
- lr_patience : int
- Patience for learning rate scheduler.
- weight_decay : float
- Weight decay for optimizer.
- lr_factor : float
- Factor by which the learning rate will be reduced.
- pooling_method : str
- Method to pool the features.
- cat_feature_info : dict
- Dictionary containing information about categorical features.
- num_feature_info : dict
- Dictionary containing information about numerical features.
- embedding_activation : callable
- Activation function for embeddings.
- mamba : Mamba
- Mamba architecture component.
- norm_f : nn.Module
- Normalization layer.
- num_embeddings : nn.ModuleList
- Module list for numerical feature embeddings.
- cat_embeddings : nn.ModuleList
- Module list for categorical feature embeddings.
- tabular_head : MLP
- Multi-layer perceptron head for tabular data.
- cls_token : nn.Parameter
- Class token parameter.
- embedding_norm : nn.Module, optional
- Layer normalization applied after embedding if specified.
- loss_fct : nn.Module
- The loss function used for training the model, MSE loss.
-
- """
-
- def __init__(
- self,
- cat_feature_info,
- num_feature_info,
- config: DefaultMambularConfig = DefaultMambularConfig(),
- seq_size: int = 20,
- raw_embeddings=False,
- **kwargs,
- ):
- super().__init__()
-
- # Save all hyperparameters
- self.save_hyperparameters(ignore=["cat_feature_info", "num_feature_info"])
-
- # Assigning values from hyperparameters
- self.lr = self.hparams.get("lr", config.lr)
- self.lr_patience = self.hparams.get("lr_patience", config.lr_patience)
- self.weight_decay = self.hparams.get("weight_decay", config.weight_decay)
- self.lr_factor = self.hparams.get("lr_factor", config.lr_factor)
- self.pooling_method = self.hparams.get("pooling_method", config.pooling_method)
- self.cat_feature_info = cat_feature_info
- self.num_feature_info = num_feature_info
-
- self.embedding_activation = self.hparams.get(
- "num_embedding_activation", config.num_embedding_activation
- )
- self.seq_size = seq_size
- self.raw_embeddings = raw_embeddings
-
- # Additional layers and components initialization based on hyperparameters
- self.mamba = Mamba(
- d_model=self.hparams.get("d_model", config.d_model),
- n_layers=self.hparams.get("n_layers", config.n_layers),
- expand_factor=self.hparams.get("expand_factor", config.expand_factor),
- bias=self.hparams.get("bias", config.bias),
- d_conv=self.hparams.get("d_conv", config.d_conv),
- conv_bias=self.hparams.get("conv_bias", config.conv_bias),
- dropout=self.hparams.get("dropout", config.dropout),
- dt_rank=self.hparams.get("dt_rank", config.dt_rank),
- d_state=self.hparams.get("d_state", config.d_state),
- dt_scale=self.hparams.get("dt_scale", config.dt_scale),
- dt_init=self.hparams.get("dt_init", config.dt_init),
- dt_max=self.hparams.get("dt_max", config.dt_max),
- dt_min=self.hparams.get("dt_min", config.dt_min),
- dt_init_floor=self.hparams.get("dt_init_floor", config.dt_init_floor),
- norm=globals()[self.hparams.get("norm", config.norm)],
- activation=self.hparams.get("activation", config.activation),
- )
-
- # Set the normalization layer dynamically
- norm_layer = self.hparams.get("norm", config.norm)
- if norm_layer == "RMSNorm":
- self.norm_f = RMSNorm(self.hparams.get("d_model", config.d_model))
- elif norm_layer == "LayerNorm":
- self.norm_f = LayerNorm(self.hparams.get("d_model", config.d_model))
- elif norm_layer == "BatchNorm":
- self.norm_f = BatchNorm(self.hparams.get("d_model", config.d_model))
- elif norm_layer == "InstanceNorm":
- self.norm_f = InstanceNorm(self.hparams.get("d_model", config.d_model))
- elif norm_layer == "GroupNorm":
- self.norm_f = GroupNorm(1, self.hparams.get("d_model", config.d_model))
- elif norm_layer == "LearnableLayerScaling":
- self.norm_f = LearnableLayerScaling(
- self.hparams.get("d_model", config.d_model)
- )
- else:
- raise ValueError(f"Unsupported normalization layer: {norm_layer}")
-
- if not self.raw_embeddings:
- data_size = len(num_feature_info.items())
- num_embedding_modules = data_size // self.seq_size
- self.num_embeddings = nn.ModuleList(
- [
- nn.Sequential(
- nn.Linear(
- self.seq_size,
- self.hparams.get("d_model", config.d_model),
- bias=False,
- ),
- # Example using ReLU as the activation function, change as needed
- self.embedding_activation,
- )
- for _ in range(num_embedding_modules)
- ]
- )
- else:
- data_size = len(num_feature_info.items())
- num_embedding_modules = data_size
- self.num_embeddings = nn.ModuleList(
- [
- nn.Sequential(
- nn.Linear(
- input_shape,
- self.hparams.get("d_model", config.d_model),
- bias=False,
- ),
- # Example using ReLU as the activation function, change as needed
- self.embedding_activation,
- )
- for feature_name, input_shape in num_feature_info.items()
- ]
- )
-
- self.cat_embeddings = nn.ModuleList(
- [
- nn.Embedding(
- num_categories + 1, self.hparams.get("d_model", config.d_model)
- )
- for feature_name, num_categories in cat_feature_info.items()
- ]
- )
-
- head_activation = self.hparams.get("head_activation", config.head_activation)
-
- self.tabular_head = MLP(
- self.hparams.get("d_model", config.d_model),
- hidden_units_list=self.hparams.get(
- "head_layer_sizes", config.head_layer_sizes
- ),
- dropout_rate=self.hparams.get("head_dropout", config.head_dropout),
- use_skip_layers=self.hparams.get(
- "head_skip_layers", config.head_skip_layers
- ),
- activation_fn=head_activation,
- use_batch_norm=self.hparams.get(
- "head_use_batch_norm", config.head_use_batch_norm
- ),
- )
-
- self.cls_token = nn.Parameter(
- torch.zeros(1, 1, self.hparams.get("d_model", config.d_model))
- )
-
- self.loss_fct = nn.MSELoss()
-
- if self.hparams.get("layer_norm_after_embedding"):
- self.embedding_norm = nn.LayerNorm(
- self.hparams.get("d_model", config.d_model)
- )
-
-[docs] def forward(self, cat_features, num_features):
- """
- Defines the forward pass of the model, processing both categorical and numerical features,
- and returning regression predictions.
-
- Parameters
- ----------
- cat_features : Tensor
- Tensor containing the categorical features.
- num_features : Tensor
- Tensor containing the numerical features or raw sequence data, depending on `raw_embeddings`.
-
-
- Returns
- -------
- Tensor
- The output predictions of the model for regression tasks.
- """
- batch_size = (
- cat_features[0].size(0) if cat_features != [] else num_features[0].size(0)
- )
- cls_tokens = self.cls_token.expand(batch_size, -1, -1)
- # Process categorical features if present
- if not self.raw_embeddings:
- if len(self.cat_embeddings) > 0 and cat_features:
- cat_embeddings = [
- emb(cat_features[i]) for i, emb in enumerate(self.cat_embeddings)
- ]
- cat_embeddings = torch.stack(cat_embeddings, dim=1)
- cat_embeddings = torch.squeeze(cat_embeddings, dim=2)
- else:
- cat_embeddings = None
-
- if len(self.num_embeddings) > 0 and num_features:
- num_embeddings = []
- # Iterate through the num_embeddings, taking slices of num_features for each
- for i, emb in enumerate(self.num_embeddings):
- # Calculate start and end indices for slicing the list
- start_idx = i * self.seq_size
- end_idx = start_idx + self.seq_size
-
- # Slice the num_features list to get the current chunk
- current_chunk = num_features[start_idx:end_idx]
-
- # If the current_chunk is not empty, process it
- if current_chunk:
- # Concatenate tensors in the current chunk along dimension 1
- chunk_tensor = torch.cat(current_chunk, dim=1)
- # Apply the embedding layer to the chunk_tensor
- num_embeddings.append(emb(chunk_tensor))
-
- # Stack the resulting embeddings along the second dimension if num_embeddings is not empty
- if num_embeddings:
- num_embeddings = torch.stack(num_embeddings, dim=1)
- else:
- num_embeddings = None
-
- else:
- # Process categorical features if present
- if len(self.cat_embeddings) > 0 and cat_features:
- cat_embeddings = [
- emb(cat_features[i]) for i, emb in enumerate(self.cat_embeddings)
- ]
- cat_embeddings = torch.stack(cat_embeddings, dim=1)
- cat_embeddings = torch.squeeze(cat_embeddings, dim=2)
- if self.hparams.get("layer_norm_after_embedding"):
- cat_embeddings = self.embedding_norm(cat_embeddings)
- else:
- cat_embeddings = None
-
- # Process numerical features if present
- if len(self.num_embeddings) > 0 and num_features:
- num_embeddings = [
- emb(num_features[i]) for i, emb in enumerate(self.num_embeddings)
- ]
- num_embeddings = torch.stack(num_embeddings, dim=1)
- if self.hparams.get("layer_norm_after_embedding"):
- num_embeddings = self.embedding_norm(num_embeddings)
- else:
- num_embeddings = None
-
- # Combine embeddings if both types are present, otherwise use whichever is available
- if cat_embeddings is not None and num_embeddings is not None:
- x = torch.cat([cls_tokens, cat_embeddings, num_embeddings], dim=1)
- elif cat_embeddings is not None:
- x = torch.cat([cls_tokens, cat_embeddings], dim=1)
- elif num_embeddings is not None:
- x = torch.cat([cls_tokens, num_embeddings], dim=1)
- else:
- raise ValueError("No features provided to the model.")
-
- x = self.mamba(x)
-
- # Apply pooling based on the specified method
- if self.pooling_method == "avg":
- x = torch.mean(x, dim=1)
- elif self.pooling_method == "max":
- x, _ = torch.max(x, dim=1)
- elif self.pooling_method == "sum":
- x = torch.sum(x, dim=1)
- elif self.pooling_method == "cls_token":
- x = x[:, 0]
- else:
- raise ValueError(f"Invalid pooling method: {self.pooling_method}")
-
- x = self.norm_f(x)
- preds = self.tabular_head(x)
-
- return preds
-
-[docs] def training_step(self, batch, batch_idx):
- """
- Processes a single batch during training, computes the loss, and logs training metrics.
-
- Parameters
- ----------
- batch : tuple
- A batch of data from the DataLoader, containing numerical features, categorical features, and labels.
- batch_idx : int
- The index of the batch within the epoch.
-
-
- Returns
- -------
- Tensor
- The computed loss for the batch.
- """
- cat_features, num_features, labels = batch
- preds = self(num_features=num_features, cat_features=cat_features)
-
- loss = self.loss_fct(preds.squeeze(), labels.float())
- self.log(
- "train_loss",
- loss,
- on_step=True,
- on_epoch=True,
- prog_bar=True,
- logger=True,
- )
- return loss
-
-[docs] def validation_step(self, batch, batch_idx):
- """
- Processes a single batch during validation, computes the loss, and logs validation metrics.
-
- Parameters
- ----------
- batch : tuple
- A batch of data from the DataLoader, containing numerical features, categorical features, and labels.
- batch_idx : int
- The index of the batch within the epoch.
- """
- cat_features, num_features, labels = batch
- preds = self(num_features=num_features, cat_features=cat_features)
-
- loss = self.loss_fct(preds.squeeze(), labels.float())
- self.log(
- "val_loss",
- loss,
- on_step=True,
- on_epoch=True,
- prog_bar=True,
- logger=True,
- )
-
- return loss
-
-[docs] def configure_optimizers(self):
- """
- Sets up the model's optimizer and learning rate scheduler based on the configurations provided.
-
- Returns
- -------
- dict
- A dictionary containing the optimizer and lr_scheduler configurations.
- """
- optimizer = torch.optim.Adam(
- self.parameters(), lr=self.lr, weight_decay=self.weight_decay
- )
- scheduler = {
- "scheduler": torch.optim.lr_scheduler.ReduceLROnPlateau(
- optimizer,
- mode="min",
- factor=self.lr_factor,
- patience=self.lr_patience,
- verbose=True,
- ),
- "monitor": "val_loss", # Name of the metric to monitor
- "interval": "epoch",
- "frequency": 1,
- }
-
- return {"optimizer": optimizer, "lr_scheduler": scheduler}
-
-
-
-
-
-import lightning as pl
-import torch
-import torch.nn as nn
-from ..utils.mamba_arch import Mamba
-from ..utils.mlp_utils import MLP
-from ..utils.normalization_layers import (
- RMSNorm,
- LayerNorm,
- LearnableLayerScaling,
- BatchNorm,
- InstanceNorm,
- GroupNorm,
-)
-from ..utils.configs import DefaultMambularConfig
-
-
-[docs]class BaseMambularRegressor(pl.LightningModule):
- """
- A PyTorch Lightning Module for regression tasks utilizing the Mamba architecture and various normalization techniques.
-
- Parameters
- ----------
- cat_feature_info : dict
- Dictionary containing information about categorical features.
- num_feature_info : dict
- Dictionary containing information about numerical features.
- config : DefaultMambularConfig, optional
- Configuration object containing default hyperparameters for the model (default is DefaultMambularConfig()).
- **kwargs : dict
- Additional keyword arguments.
-
- Attributes
- ----------
- lr : float
- Learning rate.
- lr_patience : int
- Patience for learning rate scheduler.
- weight_decay : float
- Weight decay for optimizer.
- lr_factor : float
- Factor by which the learning rate will be reduced.
- pooling_method : str
- Method to pool the features.
- cat_feature_info : dict
- Dictionary containing information about categorical features.
- num_feature_info : dict
- Dictionary containing information about numerical features.
- embedding_activation : callable
- Activation function for embeddings.
- mamba : Mamba
- Mamba architecture component.
- norm_f : nn.Module
- Normalization layer.
- num_embeddings : nn.ModuleList
- Module list for numerical feature embeddings.
- cat_embeddings : nn.ModuleList
- Module list for categorical feature embeddings.
- tabular_head : MLP
- Multi-layer perceptron head for tabular data.
- cls_token : nn.Parameter
- Class token parameter.
- loss_fct : nn.Module
- Loss function.
- embedding_norm : nn.Module, optional
- Layer normalization applied after embedding if specified.
- """
-
- def __init__(
- self,
- cat_feature_info,
- num_feature_info,
- config: DefaultMambularConfig = DefaultMambularConfig(),
- **kwargs,
- ):
- super().__init__()
-
- # Save all hyperparameters
- self.save_hyperparameters(ignore=["cat_feature_info", "num_feature_info"])
-
- # Assigning values from hyperparameters
- self.lr = self.hparams.get("lr", config.lr)
- self.lr_patience = self.hparams.get("lr_patience", config.lr_patience)
- self.weight_decay = self.hparams.get("weight_decay", config.weight_decay)
- self.lr_factor = self.hparams.get("lr_factor", config.lr_factor)
- self.pooling_method = self.hparams.get("pooling_method", config.pooling_method)
- self.cat_feature_info = cat_feature_info
- self.num_feature_info = num_feature_info
-
- self.embedding_activation = self.hparams.get(
- "num_embedding_activation", config.num_embedding_activation
- )
-
- # Additional layers and components initialization based on hyperparameters
- self.mamba = Mamba(
- d_model=self.hparams.get("d_model", config.d_model),
- n_layers=self.hparams.get("n_layers", config.n_layers),
- expand_factor=self.hparams.get("expand_factor", config.expand_factor),
- bias=self.hparams.get("bias", config.bias),
- d_conv=self.hparams.get("d_conv", config.d_conv),
- conv_bias=self.hparams.get("conv_bias", config.conv_bias),
- dropout=self.hparams.get("dropout", config.dropout),
- dt_rank=self.hparams.get("dt_rank", config.dt_rank),
- d_state=self.hparams.get("d_state", config.d_state),
- dt_scale=self.hparams.get("dt_scale", config.dt_scale),
- dt_init=self.hparams.get("dt_init", config.dt_init),
- dt_max=self.hparams.get("dt_max", config.dt_max),
- dt_min=self.hparams.get("dt_min", config.dt_min),
- dt_init_floor=self.hparams.get("dt_init_floor", config.dt_init_floor),
- norm=globals()[self.hparams.get("norm", config.norm)],
- activation=self.hparams.get("activation", config.activation),
- )
-
- # Set the normalization layer dynamically
- norm_layer = self.hparams.get("norm", config.norm)
- if norm_layer == "RMSNorm":
- self.norm_f = RMSNorm(self.hparams.get("d_model", config.d_model))
- elif norm_layer == "LayerNorm":
- self.norm_f = LayerNorm(self.hparams.get("d_model", config.d_model))
- elif norm_layer == "BatchNorm":
- self.norm_f = BatchNorm(self.hparams.get("d_model", config.d_model))
- elif norm_layer == "InstanceNorm":
- self.norm_f = InstanceNorm(self.hparams.get("d_model", config.d_model))
- elif norm_layer == "GroupNorm":
- self.norm_f = GroupNorm(1, self.hparams.get("d_model", config.d_model))
- elif norm_layer == "LearnableLayerScaling":
- self.norm_f = LearnableLayerScaling(
- self.hparams.get("d_model", config.d_model)
- )
- else:
- raise ValueError(f"Unsupported normalization layer: {norm_layer}")
-
- self.num_embeddings = nn.ModuleList(
- [
- nn.Sequential(
- nn.Linear(
- input_shape,
- self.hparams.get("d_model", config.d_model),
- bias=False,
- ),
- self.embedding_activation,
- )
- for feature_name, input_shape in num_feature_info.items()
- ]
- )
-
- self.cat_embeddings = nn.ModuleList(
- [
- nn.Embedding(
- num_categories + 1, self.hparams.get("d_model", config.d_model)
- )
- for feature_name, num_categories in cat_feature_info.items()
- ]
- )
-
- head_activation = self.hparams.get("head_activation", config.head_activation)
-
- self.tabular_head = MLP(
- self.hparams.get("d_model", config.d_model),
- hidden_units_list=self.hparams.get(
- "head_layer_sizes", config.head_layer_sizes
- ),
- dropout_rate=self.hparams.get("head_dropout", config.head_dropout),
- use_skip_layers=self.hparams.get(
- "head_skip_layers", config.head_skip_layers
- ),
- activation_fn=head_activation,
- use_batch_norm=self.hparams.get(
- "head_use_batch_norm", config.head_use_batch_norm
- ),
- )
-
- self.cls_token = nn.Parameter(
- torch.zeros(1, 1, self.hparams.get("d_model", config.d_model))
- )
-
- self.loss_fct = nn.MSELoss()
-
- if self.hparams.get("layer_norm_after_embedding"):
- self.embedding_norm = nn.LayerNorm(
- self.hparams.get("d_model", config.d_model)
- )
-
-[docs] def forward(self, num_features, cat_features):
- """
- Defines the forward pass of the regressor.
-
- Parameters
- ----------
- cat_features : Tensor
- Tensor containing the categorical features.
- num_features : Tensor
- Tensor containing the numerical features.
-
-
- Returns
- -------
- Tensor
- The output predictions of the model for regression tasks.
- """
-
- batch_size = (
- cat_features[0].size(0) if cat_features != [] else num_features[0].size(0)
- )
- cls_tokens = self.cls_token.expand(batch_size, -1, -1)
-
- # Process categorical features if present
- if len(self.cat_embeddings) > 0 and cat_features:
- cat_embeddings = [
- emb(cat_features[i]) for i, emb in enumerate(self.cat_embeddings)
- ]
- cat_embeddings = torch.stack(cat_embeddings, dim=1)
- cat_embeddings = torch.squeeze(cat_embeddings, dim=2)
- if self.hparams.get("layer_norm_after_embedding"):
- cat_embeddings = self.embedding_norm(cat_embeddings)
- else:
- cat_embeddings = None
-
- # Process numerical features if present
- if len(self.num_embeddings) > 0 and num_features:
- num_embeddings = [
- emb(num_features[i]) for i, emb in enumerate(self.num_embeddings)
- ]
- num_embeddings = torch.stack(num_embeddings, dim=1)
- if self.hparams.get("layer_norm_after_embedding"):
- num_embeddings = self.embedding_norm(num_embeddings)
- else:
- num_embeddings = None
-
- # Combine embeddings if both types are present, otherwise use whichever is available
-
- if cat_embeddings is not None and num_embeddings is not None:
- x = torch.cat([cls_tokens, cat_embeddings, num_embeddings], dim=1)
- elif cat_embeddings is not None:
- x = torch.cat([cls_tokens, cat_embeddings], dim=1)
- elif num_embeddings is not None:
- x = torch.cat([cls_tokens, num_embeddings], dim=1)
- else:
- raise ValueError("No features provided to the model.")
-
- x = self.mamba(x)
-
- # Apply pooling based on the specified method
- if self.pooling_method == "avg":
- x = torch.mean(x, dim=1)
- elif self.pooling_method == "max":
- x, _ = torch.max(x, dim=1)
- elif self.pooling_method == "sum":
- x = torch.sum(x, dim=1)
- elif self.pooling_method == "cls_token":
- x = x[:, 0]
- else:
- raise ValueError(f"Invalid pooling method: {self.pooling_method}")
-
- x = self.norm_f(x)
- preds = self.tabular_head(x)
-
- return preds
-
-[docs] def training_step(self, batch, batch_idx):
- """
- Defines the forward pass of the regressor.
-
- Parameters
- ----------
- cat_features : Tensor
- Tensor containing the categorical features.
- num_features : Tensor
- Tensor containing the numerical features.
-
-
- Returns
- -------
- Tensor
- The output predictions of the model for regression tasks.
- """
- cat_features, num_features, labels = batch
- preds = self(num_features=num_features, cat_features=cat_features)
-
- loss = self.loss_fct(preds.squeeze(), labels.float())
- self.log(
- "train_loss",
- loss,
- on_step=True,
- on_epoch=True,
- prog_bar=True,
- logger=True,
- )
- return loss
-
-[docs] def validation_step(self, batch, batch_idx):
- """
- Processes a single batch during validation, computes the loss, and logs validation metrics.
-
- Parameters
- ----------
- batch : tuple
- A batch of data from the DataLoader, containing numerical features, categorical features, and labels.
- batch_idx : int
- The index of the batch within the epoch.
- """
- cat_features, num_features, labels = batch
- preds = self(num_features=num_features, cat_features=cat_features)
-
- loss = self.loss_fct(preds.squeeze(), labels.float())
- self.log(
- "val_loss",
- loss,
- on_step=True,
- on_epoch=True,
- prog_bar=True,
- logger=True,
- )
-
- return loss
-
-[docs] def configure_optimizers(self):
- """
- Sets up the model's optimizer and learning rate scheduler based on the configurations provided.
-
- Returns
- -------
- dict
- A dictionary containing the optimizer and lr_scheduler configurations.
- """
- optimizer = torch.optim.Adam(
- self.parameters(), lr=self.lr, weight_decay=self.weight_decay
- )
- scheduler = {
- "scheduler": torch.optim.lr_scheduler.ReduceLROnPlateau(
- optimizer,
- mode="min",
- factor=self.lr_factor,
- patience=self.lr_patience,
- verbose=True,
- ),
- "monitor": "val_loss", # Name of the metric to monitor
- "interval": "epoch",
- "frequency": 1,
- }
-
- return {"optimizer": optimizer, "lr_scheduler": scheduler}
-
-
-
-
-
-import lightning as pl
-import numpy as np
-import warnings
-import pandas as pd
-import torch
-from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint
-from sklearn.base import BaseEstimator
-from sklearn.metrics import accuracy_score
-from sklearn.model_selection import train_test_split
-from torch.utils.data import DataLoader
-
-from ..base_models.classifier import BaseMambularClassifier
-from ..utils.configs import DefaultMambularConfig
-from ..utils.dataset import MambularDataModule, MambularDataset
-from ..utils.preprocessor import Preprocessor
-
-
-[docs]class MambularClassifier(BaseEstimator):
- """
- A classifier that mimics scikit-learn's API using PyTorch Lightning and a custom architecture.
-
- This classifier is designed to work with tabular data and provides a flexible interface for specifying model
- configurations and preprocessing steps. It integrates smoothly with scikit-learn's utilities, such as cross-validation
- and grid search.
-
- Parameters
- ----------
- # configuration parameters
- lr : float, optional
- Learning rate for the optimizer. Default is 1e-4.
- lr_patience : int, optional
- Number of epochs with no improvement on the validation loss to wait before reducing the learning rate. Default is 10.
- weight_decay : float, optional
- Weight decay (L2 penalty) coefficient. Default is 1e-6.
- lr_factor : float, optional
- Factor by which the learning rate will be reduced. Default is 0.1.
- d_model : int, optional
- Dimension of the model. Default is 64.
- n_layers : int, optional
- Number of layers. Default is 8.
- expand_factor : int, optional
- Expansion factor. Default is 2.
- bias : bool, optional
- Whether to use bias. Default is False.
- d_conv : int, optional
- Dimension of the convolution. Default is 16.
- conv_bias : bool, optional
- Whether to use bias in the convolution. Default is True.
- dropout : float, optional
- Dropout rate in the mamba blocks. Default is 0.05.
- dt_rank : str, optional
- Rank of the time dimension. Default is "auto".
- d_state : int, optional
- State dimension. Default is 16.
- dt_scale : float, optional
- Scale of the time dimension. Default is 1.0.
- dt_init : str, optional
- Initialization method for the time dimension. Default is "random".
- dt_max : float, optional
- Maximum value for the time dimension. Default is 0.1.
- dt_min : float, optional
- Minimum value for the time dimension. Default is 1e-3.
- dt_init_floor : float, optional
- Floor value for the time dimension initialization. Default is 1e-4.
- norm : str, optional
- Normalization method. Default is 'RMSNorm'.
- activation : callable, optional
- Activation function. Default is nn.SELU().
- num_embedding_activation : callable, optional
- Activation function for numerical embeddings. Default is nn.Identity().
- head_layer_sizes : list, optional
- Sizes of the layers in the head. Default is [64, 64, 32].
- head_dropout : float, optional
- Dropout rate for the head. Default is 0.5.
- head_skip_layers : bool, optional
- Whether to use skip layers in the head. Default is False.
- head_activation : callable, optional
- Activation function for the head. Default is nn.SELU().
- head_use_batch_norm : bool, optional
- Whether to use batch normalization in the head. Default is False.
-
- # Preprocessor Parameters
- n_bins : int, optional
- The number of bins to use for numerical feature binning. Default is 50.
- numerical_preprocessing : str, optional
- The preprocessing strategy for numerical features. Default is 'ple'.
- use_decision_tree_bins : bool, optional
- If True, uses decision tree regression/classification to determine optimal bin edges for numerical feature binning. Default is False.
- binning_strategy : str, optional
- Defines the strategy for binning numerical features. Default is 'uniform'.
- task : str, optional
- Indicates the type of machine learning task ('regression' or 'classification'). Default is 'regression'.
- cat_cutoff: float or int, optional
- Indicates the cutoff after which integer values are treated as categorical. If float, it's treated as a percentage. If int, it's the maximum number of unique values for a column to be considered categorical. Default is 3%
- treat_all_integers_as_numerical : bool, optional
- If True, all integer columns will be treated as numerical, regardless of their unique value count or proportion. Default is False
-
-
-
- Attributes
- ----------
- config : MambularConfig
- Configuration object that holds model-specific settings.
- preprocessor : Preprocessor
- Preprocessor object for handling feature preprocessing like normalization and encoding.
- model : BaseMambularClassifier or None
- The underlying PyTorch Lightning model, instantiated upon calling the `fit` method.
- """
-
- def __init__(self, **kwargs):
- # Known config arguments
- config_arg_names = [
- "lr",
- "lr_patience",
- "weight_decay",
- "lr_factor",
- "d_model",
- "n_layers",
- "expand_factor",
- "bias",
- "d_conv",
- "conv_bias",
- "dropout",
- "dt_rank",
- "d_state",
- "dt_scale",
- "dt_init",
- "dt_max",
- "dt_min",
- "dt_init_floor",
- "norm",
- "activation",
- "num_embedding_activation",
- "head_layer_sizes",
- "head_dropout",
- "head_skip_layers",
- "head_activation",
- "head_use_batch_norm",
- ]
-
- preprocessor_arg_names = [
- "n_bins",
- "numerical_preprocessing",
- "use_decision_tree_bins",
- "binning_strategy",
- "task",
- "cat_cutoff",
- "treat_all_integers_as_numerical",
- ]
-
- self.config_kwargs = {k: v for k, v in kwargs.items() if k in config_arg_names}
- self.config = DefaultMambularConfig(**self.config_kwargs)
-
- preprocessor_kwargs = {
- k: v for k, v in kwargs.items() if k in preprocessor_arg_names
- }
- # Raise a warning if task is set to 'classification'
- if preprocessor_kwargs.get("task") == "regression":
- warnings.warn(
- "The task in preprocessing binning is set to 'regression'. MambularClassifier is designed for classification tasks.",
- UserWarning,
- )
-
- if "task" not in list(preprocessor_kwargs.keys()):
- preprocessor_kwargs["task"] = "classification"
-
- self.preprocessor = Preprocessor(**preprocessor_kwargs)
- self.model = None
-
-[docs] def get_params(self, deep=True):
- """
- Get parameters for this estimator.
-
- Parameters
- ----------
- deep : bool, default=True
- If True, will return the parameters for this estimator and contained subobjects that are estimators.
-
-
- Returns
- -------
- params : dict
- Parameter names mapped to their values.
- """
-
- params = self.config_kwargs # Parameters used to initialize MambularConfig
-
- # If deep=True, include parameters from nested components like preprocessor
- if deep:
- # Assuming Preprocessor has a get_params method
- preprocessor_params = {
- "preprocessor__" + key: value
- for key, value in self.preprocessor.get_params().items()
- }
- params.update(preprocessor_params)
-
- return params
-
-[docs] def set_params(self, **parameters):
- """
- Set the parameters of this estimator.
-
- Parameters
- ----------
- **parameters : dict
- Estimator parameters.
-
-
- Returns
- -------
- self : object
- Estimator instance.
- """
- # Update config_kwargs with provided parameters
- valid_config_keys = self.config_kwargs.keys()
- config_updates = {k: v for k, v in parameters.items() if k in valid_config_keys}
- self.config_kwargs.update(config_updates)
-
- # Update the config object
- for key, value in config_updates.items():
- setattr(self.config, key, value)
-
- # Handle preprocessor parameters (prefixed with 'preprocessor__')
- preprocessor_params = {
- k.split("__")[1]: v
- for k, v in parameters.items()
- if k.startswith("preprocessor__")
- }
- if preprocessor_params:
- # Assuming Preprocessor has a set_params method
- self.preprocessor.set_params(**preprocessor_params)
-
- return self
-
-[docs] def split_data(self, X, y, val_size, random_state):
- """
- Split the dataset into training and validation sets.
-
- Parameters
- ----------
- X : array-like of shape (n_samples, n_features)
- The input samples.
- y : array-like of shape (n_samples,)
- The target values.
- val_size : float
- The proportion of the dataset to include in the validation split.
- random_state : int
- Controls the shuffling applied to the data before applying the split.
-
-
- Returns
- -------
- X_train, X_val, y_train, y_val : arrays
- The split datasets.
- """
- X_train, X_val, y_train, y_val = train_test_split(
- X, y, test_size=val_size, random_state=random_state
- )
-
- return X_train, X_val, y_train, y_val
-
-[docs] def preprocess_data(self, X_train, y_train, X_val, y_val, batch_size, shuffle):
- """
- Preprocess the training and validation data and create corresponding DataLoaders.
-
- Parameters
- ----------
- X_train : array-like of shape (n_samples, n_features)
- The training input samples.
- y_train : array-like of shape (n_samples,)
- The training target values.
- X_val : array-like of shape (n_samples, n_features)
- The validation input samples.
- y_val : array-like of shape (n_samples,)
- The validation target values.
- batch_size : int
- Size of mini-batches for the DataLoader.
- shuffle : bool
- Whether to shuffle the training data before splitting into batches.
-
-
- Returns
- -------
- data_module : MambularDataModule
- An instance of MambularDataModule containing training and validation DataLoaders.
- """
- self.preprocessor.fit(
- pd.concat([X_train, X_val], axis=0).reset_index(drop=True),
- np.concatenate((y_train, y_val), axis=0),
- )
- train_preprocessed_data = self.preprocessor.transform(X_train)
- val_preprocessed_data = self.preprocessor.transform(X_val)
-
- # Update feature info based on the actual processed data
- (
- self.cat_feature_info,
- self.num_feature_info,
- ) = self.preprocessor.get_feature_info()
-
- # Initialize lists for tensors
- train_cat_tensors = []
- train_num_tensors = []
- val_cat_tensors = []
- val_num_tensors = []
-
- # Populate tensors for categorical features, if present in processed data
- for key in self.cat_feature_info:
- cat_key = "cat_" + key # Assuming categorical keys are prefixed with 'cat_'
- if cat_key in train_preprocessed_data:
- train_cat_tensors.append(
- torch.tensor(train_preprocessed_data[cat_key], dtype=torch.long)
- )
- if cat_key in val_preprocessed_data:
- val_cat_tensors.append(
- torch.tensor(val_preprocessed_data[cat_key], dtype=torch.long)
- )
-
- binned_key = "num_" + key # for binned features
- if binned_key in train_preprocessed_data:
- train_cat_tensors.append(
- torch.tensor(train_preprocessed_data[binned_key], dtype=torch.long)
- )
-
- if binned_key in val_preprocessed_data:
- val_cat_tensors.append(
- torch.tensor(val_preprocessed_data[binned_key], dtype=torch.long)
- )
-
- # Populate tensors for numerical features, if present in processed data
- for key in self.num_feature_info:
- num_key = "num_" + key # Assuming numerical keys are prefixed with 'num_'
- if num_key in train_preprocessed_data:
- train_num_tensors.append(
- torch.tensor(train_preprocessed_data[num_key], dtype=torch.float32)
- )
- if num_key in val_preprocessed_data:
- val_num_tensors.append(
- torch.tensor(val_preprocessed_data[num_key], dtype=torch.float32)
- )
-
- train_labels = torch.tensor(y_train, dtype=torch.long)
- val_labels = torch.tensor(y_val, dtype=torch.long)
-
- # Create datasets
- train_dataset = MambularDataset(
- train_cat_tensors, train_num_tensors, train_labels, regression=False
- )
- val_dataset = MambularDataset(
- val_cat_tensors, val_num_tensors, val_labels, regression=False
- )
-
- # Create dataloaders
- train_dataloader = DataLoader(
- train_dataset, batch_size=batch_size, shuffle=shuffle
- )
- val_dataloader = DataLoader(val_dataset, batch_size=batch_size)
-
- return MambularDataModule(train_dataloader, val_dataloader)
-
-[docs] def preprocess_test_data(self, X):
- """
- Preprocesses the test data and creates tensors for categorical and numerical features.
-
- Parameters
- ----------
- X : DataFrame or array-like, shape (n_samples, n_features)
- Test feature set.
-
-
- Returns
- -------
- cat_tensors : list of Tensors
- List of tensors for each categorical feature.
- num_tensors : list of Tensors
- List of tensors for each numerical feature.
- """
- processed_data = self.preprocessor.transform(X)
-
- # Initialize lists for tensors
- cat_tensors = []
- num_tensors = []
-
- # Populate tensors for categorical features
- for key in self.cat_feature_info:
- cat_key = "cat_" + key # Assuming categorical keys are prefixed with 'cat_'
- if cat_key in processed_data:
- cat_tensors.append(
- torch.tensor(processed_data[cat_key], dtype=torch.long)
- )
-
- binned_key = "num_" + key # for binned features
- if binned_key in processed_data:
- cat_tensors.append(
- torch.tensor(processed_data[binned_key], dtype=torch.long)
- )
-
- # Populate tensors for numerical features
- for key in self.num_feature_info:
- num_key = "num_" + key # Assuming numerical keys are prefixed with 'num_'
- if num_key in processed_data:
- num_tensors.append(
- torch.tensor(processed_data[num_key], dtype=torch.float32)
- )
-
- return cat_tensors, num_tensors
-
-[docs] def fit(
- self,
- X,
- y,
- val_size: float = 0.2,
- X_val=None,
- y_val=None,
- max_epochs: int = 100,
- random_state: int = 101,
- batch_size: int = 128,
- shuffle: bool = True,
- patience: int = 15,
- monitor: str = "val_loss",
- mode: str = "min",
- lr: float = 1e-4,
- lr_patience: int = 10,
- factor: float = 0.1,
- weight_decay: float = 1e-06,
- **trainer_kwargs
- ):
- """
- Fit the model to the given training data, optionally using a separate validation set.
-
- Parameters
- ----------
- X : array-like or pd.DataFrame of shape (n_samples, n_features)
- The training input samples.
- y : array-like of shape (n_samples,) or (n_samples, n_outputs)
- The target values (class labels in classification, real numbers in regression).
- val_size : float, default=0.2
- The proportion of the dataset to include in the validation split if `X_val` is None. Ignored if `X_val` is provided.
- X_val : array-like or pd.DataFrame of shape (n_samples, n_features), optional
- The validation input samples. If provided, `X` and `y` are not split and this data is used for validation.
- y_val : array-like of shape (n_samples,) or (n_samples, n_outputs), optional
- The validation target values. Required if `X_val` is provided.
- max_epochs : int, default=100
- Maximum number of epochs for training.
- random_state : int, default=101
- Seed used by the random number generator for shuffling the data if `X_val` is not provided.
- batch_size : int, default=64
- Number of samples per gradient update.
- shuffle : bool, default=True
- Whether to shuffle the training data before each epoch if `X_val` is not provided.
- patience : int, default=10
- Number of epochs with no improvement after which training will be stopped if using early stopping.
- monitor : str, default="val_loss"
- Quantity to be monitored for early stopping.
- mode : str, default="min"
- One of {"min", "max"}. In "min" mode, training will stop when the quantity monitored has stopped decreasing; in "max" mode, it will stop when the quantity monitored has stopped increasing.
- lr : float, default=1e-3
- Learning rate for the optimizer.
- lr_patience : int, default=10
- Number of epochs with no improvement after which the learning rate will be reduced.
- factor : float, default=0.75
- Factor by which the learning rate will be reduced. new_lr = lr * factor.
- weight_decay : float, default=0.025
- Weight decay (L2 penalty) parameter.
- **trainer_kwargs : dict
- Additional keyword arguments to be passed to the PyTorch Lightning Trainer constructor.
-
-
- Returns
- -------
- self : object
- The fitted estimator.
- """
-
- if not isinstance(X, pd.DataFrame):
- X = pd.DataFrame(X)
- if X_val:
- if not isinstance(X_val, pd.DataFrame):
- X_val = pd.DataFrame(X_val)
-
- num_classes = len(np.unique(y))
-
- if not X_val:
- X_train, X_val, y_train, y_val = self.split_data(
- X, y, val_size, random_state
- )
-
- self.data_module = self.preprocess_data(
- X_train, y_train, X_val, y_val, batch_size, shuffle
- )
-
- self.model = BaseMambularClassifier(
- num_classes=num_classes,
- config=self.config,
- cat_feature_info=self.cat_feature_info,
- num_feature_info=self.num_feature_info,
- lr=lr,
- lr_patience=lr_patience,
- lr_factor=factor,
- weight_decay=weight_decay,
- )
-
- early_stop_callback = EarlyStopping(
- monitor=monitor, min_delta=0.00, patience=patience, verbose=False, mode=mode
- )
-
- checkpoint_callback = ModelCheckpoint(
- monitor="val_loss", # Adjust according to your validation metric
- mode="min",
- save_top_k=1,
- dirpath="model_checkpoints", # Specify the directory to save checkpoints
- filename="best_model",
- )
-
- # Initialize the trainer and train the model
- trainer = pl.Trainer(
- max_epochs=max_epochs,
- callbacks=[early_stop_callback, checkpoint_callback],
- **trainer_kwargs
- )
- trainer.fit(self.model, self.data_module)
-
- best_model_path = checkpoint_callback.best_model_path
- if best_model_path:
- checkpoint = torch.load(best_model_path)
- self.model.load_state_dict(checkpoint["state_dict"])
-
- return self
-
-[docs] def predict(self, X):
- """
- Predict the class labels for the given input samples.
-
- Parameters
- ----------
- X : array-like or pd.DataFrame of shape (n_samples, n_features)
- The input samples to predict.
-
-
- Returns
- -------
- predictions : ndarray of shape (n_samples,)
- Predicted class labels for each input sample.
-
-
- Notes
- -----
- The method preprocesses the input data using the same preprocessor used during training,
- sets the model to evaluation mode, and then performs inference to predict the class labels.
- The predictions are converted from a PyTorch tensor to a NumPy array before being returned.
-
- """
-
- # Preprocess the data
- if not isinstance(X, pd.DataFrame):
- X = pd.DataFrame(X)
- device = next(self.model.parameters()).device
- cat_tensors, num_tensors = self.preprocess_test_data(X)
- if isinstance(cat_tensors, list):
- cat_tensors = [tensor.to(device) for tensor in cat_tensors]
- else:
- cat_tensors = cat_tensors.to(device)
-
- if isinstance(num_tensors, list):
- num_tensors = [tensor.to(device) for tensor in num_tensors]
- else:
- num_tensors = num_tensors.to(device)
-
- # Set the model to evaluation mode
- self.model.eval()
-
- # Perform inference
- with torch.no_grad():
- logits = self.model(num_features=num_tensors, cat_features=cat_tensors)
-
- # Check the shape of the logits to determine binary or multi-class classification
- if logits.shape[1] == 1:
- # Binary classification
- probabilities = torch.sigmoid(logits)
- predictions = (probabilities > 0.5).long().squeeze()
- else:
- # Multi-class classification
- probabilities = torch.softmax(logits, dim=1)
- predictions = torch.argmax(probabilities, dim=1)
-
- # Convert predictions to NumPy array and return
- return predictions.cpu().numpy()
-
-[docs] def predict_proba(self, X):
- """
- Predict class probabilities for the given input samples.
-
- Parameters
- ----------
- X : array-like or pd.DataFrame of shape (n_samples, n_features)
- The input samples for which to predict class probabilities.
-
-
- Notes
- -----
- The method preprocesses the input data using the same preprocessor used during training,
- sets the model to evaluation mode, and then performs inference to predict the class probabilities.
- Softmax is applied to the logits to obtain probabilities, which are then converted from a PyTorch tensor
- to a NumPy array before being returned.
-
-
- Examples
- --------
- >>> from sklearn.metrics import accuracy_score, precision_score, f1_score, roc_auc_score
- >>> # Define the metrics you want to evaluate
- >>> metrics = {
- ... 'Accuracy': (accuracy_score, False),
- ... 'Precision': (precision_score, False),
- ... 'F1 Score': (f1_score, False),
- ... 'AUC Score': (roc_auc_score, True)
- ... }
- >>> # Assuming 'X_test' and 'y_test' are your test dataset and labels
- >>> # Evaluate using the specified metrics
- >>> results = classifier.evaluate(X_test, y_test, metrics=metrics)
-
-
- Returns
- -------
- probabilities : ndarray of shape (n_samples, n_classes)
- Predicted class probabilities for each input sample.
-
- """
- # Preprocess the data
- if not isinstance(X, pd.DataFrame):
- X = pd.DataFrame(X)
- device = next(self.model.parameters()).device
- cat_tensors, num_tensors = self.preprocess_test_data(X)
- if isinstance(cat_tensors, list):
- cat_tensors = [tensor.to(device) for tensor in cat_tensors]
- else:
- cat_tensors = cat_tensors.to(device)
-
- if isinstance(num_tensors, list):
- num_tensors = [tensor.to(device) for tensor in num_tensors]
- else:
- num_tensors = num_tensors.to(device)
-
- # Set the model to evaluation mode
- self.model.eval()
-
- # Perform inference
- with torch.no_grad():
- logits = self.model(num_features=num_tensors, cat_features=cat_tensors)
- if logits.shape[1] > 1:
- probabilities = torch.softmax(logits, dim=1)
- else:
- probabilities = torch.sigmoid(logits)
-
- # Convert probabilities to NumPy array and return
- return probabilities.cpu().numpy()
-
-[docs] def evaluate(self, X, y_true, metrics=None):
- """
- Evaluate the model on the given data using specified metrics.
-
- Parameters
- ----------
- X : array-like or pd.DataFrame of shape (n_samples, n_features)
- The input samples to predict.
- y_true : array-like of shape (n_samples,)
- The true class labels against which to evaluate the predictions.
- metrics : dict
- A dictionary where keys are metric names and values are tuples containing the metric function
- and a boolean indicating whether the metric requires probability scores (True) or class labels (False).
-
-
- Returns
- -------
- scores : dict
- A dictionary with metric names as keys and their corresponding scores as values.
-
-
- Notes
- -----
- This method uses either the `predict` or `predict_proba` method depending on the metric requirements.
- """
- # Ensure input is in the correct format
- if metrics is None:
- metrics = {"Accuracy": (accuracy_score, False)}
-
- if not isinstance(X, pd.DataFrame):
- X = pd.DataFrame(X)
-
- # Initialize dictionary to store results
- scores = {}
-
- # Generate class probabilities if any metric requires them
- if any(use_proba for _, use_proba in metrics.values()):
- probabilities = self.predict_proba(X)
-
- # Generate class labels if any metric requires them
- if any(not use_proba for _, use_proba in metrics.values()):
- predictions = self.predict(X)
-
- # Compute each metric
- for metric_name, (metric_func, use_proba) in metrics.items():
- if use_proba:
- scores[metric_name] = metric_func(y_true, probabilities)
- else:
- scores[metric_name] = metric_func(y_true, predictions)
-
- return scores
-
-
-
-
-
-import lightning as pl
-import numpy as np
-import warnings
-import pandas as pd
-import properscoring as ps
-import torch
-from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint
-from sklearn.base import BaseEstimator
-from sklearn.metrics import accuracy_score, mean_squared_error
-from sklearn.model_selection import train_test_split
-from torch.utils.data import DataLoader
-
-from ..base_models.distributional import BaseMambularLSS
-from ..utils.configs import DefaultMambularConfig
-from ..utils.dataset import MambularDataModule, MambularDataset
-from ..utils.distributional_metrics import (
- beta_brier_score,
- dirichlet_error,
- gamma_deviance,
- inverse_gamma_loss,
- negative_binomial_deviance,
- poisson_deviance,
- student_t_loss,
-)
-from ..utils.preprocessor import Preprocessor
-
-
-[docs]class MambularLSS(BaseEstimator):
- """
- MambularLSS is a machine learning estimator that is designed for structured data,
- incorporating both preprocessing and a deep learning model. The estimator
- integrates configurable components for data preprocessing and the neural network model,
- facilitating end-to-end training and prediction workflows.
-
- The initialization of this class separates configuration arguments for the model and
- the preprocessor, allowing for flexible adjustment of parameters.
-
- Parameters
- ----------
- # configuration parameters
- lr : float, optional
- Learning rate for the optimizer. Default is 1e-4.
- lr_patience : int, optional
- Number of epochs with no improvement on the validation loss to wait before reducing the learning rate. Default is 10.
- weight_decay : float, optional
- Weight decay (L2 penalty) coefficient. Default is 1e-6.
- lr_factor : float, optional
- Factor by which the learning rate will be reduced. Default is 0.1.
- d_model : int, optional
- Dimension of the model. Default is 64.
- n_layers : int, optional
- Number of layers. Default is 8.
- expand_factor : int, optional
- Expansion factor. Default is 2.
- bias : bool, optional
- Whether to use bias. Default is False.
- d_conv : int, optional
- Dimension of the convolution. Default is 16.
- conv_bias : bool, optional
- Whether to use bias in the convolution. Default is True.
- dropout : float, optional
- Dropout rate in the mamba blocks. Default is 0.05.
- dt_rank : str, optional
- Rank of the time dimension. Default is "auto".
- d_state : int, optional
- State dimension. Default is 16.
- dt_scale : float, optional
- Scale of the time dimension. Default is 1.0.
- dt_init : str, optional
- Initialization method for the time dimension. Default is "random".
- dt_max : float, optional
- Maximum value for the time dimension. Default is 0.1.
- dt_min : float, optional
- Minimum value for the time dimension. Default is 1e-3.
- dt_init_floor : float, optional
- Floor value for the time dimension initialization. Default is 1e-4.
- norm : str, optional
- Normalization method. Default is 'RMSNorm'.
- activation : callable, optional
- Activation function. Default is nn.SELU().
- num_embedding_activation : callable, optional
- Activation function for numerical embeddings. Default is nn.Identity().
- head_layer_sizes : list, optional
- Sizes of the layers in the head. Default is [64, 64, 32].
- head_dropout : float, optional
- Dropout rate for the head. Default is 0.5.
- head_skip_layers : bool, optional
- Whether to use skip layers in the head. Default is False.
- head_activation : callable, optional
- Activation function for the head. Default is nn.SELU().
- head_use_batch_norm : bool, optional
- Whether to use batch normalization in the head. Default is False.
-
- # Preprocessor Parameters
- n_bins : int, optional
- The number of bins to use for numerical feature binning. Default is 50.
- numerical_preprocessing : str, optional
- The preprocessing strategy for numerical features. Default is 'ple'.
- use_decision_tree_bins : bool, optional
- If True, uses decision tree regression/classification to determine optimal bin edges for numerical feature binning. Default is False.
- binning_strategy : str, optional
- Defines the strategy for binning numerical features. Default is 'uniform'.
- task : str, optional
- Indicates the type of machine learning task ('regression' or 'classification'). Default is 'regression'.
- cat_cutoff: float or int, optional
- Indicates the cutoff after which integer values are treated as categorical. If float, it's treated as a percentage. If int, it's the maximum number of unique values for a column to be considered categorical. Default is 3%
- treat_all_integers_as_numerical : bool, optional
- If True, all integer columns will be treated as numerical, regardless of their unique value count or proportion. Default is False
-
-
-
- Attributes
- ----------
- config : MambularConfig
- Configuration object that holds model-specific settings.
- preprocessor : Preprocessor
- Preprocessor object for handling feature preprocessing like normalization and encoding.
- model : BaseMambularClassifier or None
- The underlying PyTorch Lightning model, instantiated upon calling the `fit` method.
- """
-
- def __init__(self, **kwargs):
- # Known config arguments
- config_arg_names = [
- "lr",
- "lr_patience",
- "weight_decay",
- "lr_factor",
- "d_model",
- "n_layers",
- "expand_factor",
- "bias",
- "d_conv",
- "conv_bias",
- "dropout",
- "dt_rank",
- "d_state",
- "dt_scale",
- "dt_init",
- "dt_max",
- "dt_min",
- "dt_init_floor",
- "norm",
- "activation",
- "num_embedding_activation",
- "head_layer_sizes",
- "head_dropout",
- "head_skip_layers",
- "head_activation",
- "head_use_batch_norm",
- ]
-
- preprocessor_arg_names = [
- "n_bins",
- "numerical_preprocessing",
- "use_decision_tree_bins",
- "binning_strategy",
- "task",
- "cat_cutoff",
- "treat_all_integers_as_numerical",
- ]
-
- self.config_kwargs = {k: v for k, v in kwargs.items() if k in config_arg_names}
- self.config = DefaultMambularConfig(**self.config_kwargs)
-
- preprocessor_kwargs = {
- k: v for k, v in kwargs.items() if k in preprocessor_arg_names
- }
- # Raise a warning if task is set to 'classification'
- if preprocessor_kwargs.get("task") == "regression":
- warnings.warn(
- "The task in preprocessing binning is set to 'regression'. Make sure that this is correct for your distributional family ",
- UserWarning,
- )
-
- self.preprocessor = Preprocessor(**preprocessor_kwargs)
- self.model = None
-
-[docs] def get_params(self, deep=True):
- """
- Get parameters for this estimator, optionally including parameters from nested components
- like the preprocessor.
-
- Parameters
- ----------
- deep : bool, default=True
- If True, return parameters of nested components.
-
-
- Returns
- -------
- dict
- A dictionary mapping parameter names to their values. For nested components,
- parameter names are prefixed accordingly (e.g., 'preprocessor__<param_name>').
- """
-
- params = self.config_kwargs # Parameters used to initialize MambularConfig
-
- # If deep=True, include parameters from nested components like preprocessor
- if deep:
- # Assuming Preprocessor has a get_params method
- preprocessor_params = {
- "preprocessor__" + key: value
- for key, value in self.preprocessor.get_params().items()
- }
- params.update(preprocessor_params)
-
- return params
-
-[docs] def set_params(self, **parameters):
- """
- Set the parameters of this estimator, allowing for modifications to both the configuration
- and preprocessor parameters. Parameters not recognized as configuration arguments are
- assumed to be preprocessor arguments.
-
- Parameters
- ----------
- **parameters: Arbitrary keyword arguments where keys are parameter names and values
- are the new parameter values.
-
-
- Returns
- -------
- self: This instance with updated parameters.
- """
- # Update config_kwargs with provided parameters
- valid_config_keys = self.config_kwargs.keys()
- config_updates = {k: v for k, v in parameters.items() if k in valid_config_keys}
- self.config_kwargs.update(config_updates)
-
- # Update the config object
- for key, value in config_updates.items():
- setattr(self.config, key, value)
-
- # Handle preprocessor parameters (prefixed with 'preprocessor__')
- preprocessor_params = {
- k.split("__")[1]: v
- for k, v in parameters.items()
- if k.startswith("preprocessor__")
- }
- if preprocessor_params:
- # Assuming Preprocessor has a set_params method
- self.preprocessor.set_params(**preprocessor_params)
-
- return self
-
-[docs] def split_data(self, X, y, val_size, random_state):
- """
- Split the dataset into training and validation sets.
-
- Parameters
- ----------
- X : array-like
- Features of the dataset.
- y : array-like
- Target values.
- val_size : float
- The proportion of the dataset to include in the validation split.
- random_state : int, optional
- The seed used by the random number generator for reproducibility.
-
-
- Returns
- -------
- tuple
- A tuple containing split datasets (X_train, X_val, y_train, y_val).
- """
- X_train, X_val, y_train, y_val = train_test_split(
- X, y, test_size=val_size, random_state=random_state
- )
-
- return X_train, X_val, y_train, y_val
-
-[docs] def preprocess_data(self, X_train, y_train, X_val, y_val, batch_size, shuffle):
- """
- Preprocess the training and validation data, fit the preprocessor on the training data,
- and transform both training and validation data. This method also initializes tensors
- for categorical and numerical features and labels, and prepares DataLoader objects for
- both datasets.
-
- Parameters
- ----------
- X_train : array-like
- Training features.
- y_train : array-like
- Training target values.
- X_val : array-like
- Validation features.
- y_val : array-like
- Validation target values.
- batch_size : int
- Batch size for DataLoader objects.
- shuffle : bool
- Whether to shuffle the training data in the DataLoader.
-
-
- Returns
- -------
- MambularDataModule
- An object containing DataLoaders for training and validation datasets.
- """
- self.preprocessor.fit(
- pd.concat([X_train, X_val], axis=0).reset_index(drop=True),
- np.concatenate((y_train, y_val), axis=0),
- )
- train_preprocessed_data = self.preprocessor.transform(X_train)
- val_preprocessed_data = self.preprocessor.transform(X_val)
-
- # Update feature info based on the actual processed data
- (
- self.cat_feature_info,
- self.num_feature_info,
- ) = self.preprocessor.get_feature_info()
-
- # Initialize lists for tensors
- train_cat_tensors = []
- train_num_tensors = []
- val_cat_tensors = []
- val_num_tensors = []
-
- # Populate tensors for categorical features, if present in processed data
- for key in self.cat_feature_info:
- cat_key = "cat_" + key # Assuming categorical keys are prefixed with 'cat_'
- if cat_key in train_preprocessed_data:
- train_cat_tensors.append(
- torch.tensor(train_preprocessed_data[cat_key], dtype=torch.long)
- )
- if cat_key in val_preprocessed_data:
- val_cat_tensors.append(
- torch.tensor(val_preprocessed_data[cat_key], dtype=torch.long)
- )
-
- binned_key = "num_" + key # for binned features
- if binned_key in train_preprocessed_data:
- train_cat_tensors.append(
- torch.tensor(train_preprocessed_data[binned_key], dtype=torch.long)
- )
-
- if binned_key in val_preprocessed_data:
- val_cat_tensors.append(
- torch.tensor(val_preprocessed_data[binned_key], dtype=torch.long)
- )
-
- # Populate tensors for numerical features, if present in processed data
- for key in self.num_feature_info:
- num_key = "num_" + key # Assuming numerical keys are prefixed with 'num_'
- if num_key in train_preprocessed_data:
- train_num_tensors.append(
- torch.tensor(train_preprocessed_data[num_key], dtype=torch.float32)
- )
- if num_key in val_preprocessed_data:
- val_num_tensors.append(
- torch.tensor(val_preprocessed_data[num_key], dtype=torch.float32)
- )
-
- train_labels = torch.tensor(y_train, dtype=torch.float32)
- val_labels = torch.tensor(y_val, dtype=torch.float32)
-
- # Create datasets
- train_dataset = MambularDataset(
- train_cat_tensors, train_num_tensors, train_labels
- )
- val_dataset = MambularDataset(val_cat_tensors, val_num_tensors, val_labels)
-
- # Create dataloaders
- train_dataloader = DataLoader(
- train_dataset, batch_size=batch_size, shuffle=shuffle
- )
- val_dataloader = DataLoader(val_dataset, batch_size=batch_size)
-
- return MambularDataModule(train_dataloader, val_dataloader)
-
-[docs] def preprocess_test_data(self, X):
- """
- Preprocess test data using the fitted preprocessor. This method prepares tensors for
- categorical and numerical features based on the preprocessed test data.
-
- Parameters
- ----------
- X : array-like
- Test features to preprocess.
-
-
- Returns
- -------
- tuple
- A tuple containing lists of tensors for categorical and numerical features.
- """
- processed_data = self.preprocessor.transform(X)
-
- # Initialize lists for tensors
- cat_tensors = []
- num_tensors = []
-
- # Populate tensors for categorical features
- for key in self.cat_feature_info:
- cat_key = "cat_" + key # Assuming categorical keys are prefixed with 'cat_'
- if cat_key in processed_data:
- cat_tensors.append(
- torch.tensor(processed_data[cat_key], dtype=torch.long)
- )
-
- binned_key = "num_" + key # for binned features
- if binned_key in processed_data:
- cat_tensors.append(
- torch.tensor(processed_data[binned_key], dtype=torch.long)
- )
-
- # Populate tensors for numerical features
- for key in self.num_feature_info:
- num_key = "num_" + key # Assuming numerical keys are prefixed with 'num_'
- if num_key in processed_data:
- num_tensors.append(
- torch.tensor(processed_data[num_key], dtype=torch.float32)
- )
-
- return cat_tensors, num_tensors
-
-[docs] def fit(
- self,
- X,
- y,
- family,
- val_size: float = 0.2,
- X_val=None,
- y_val=None,
- max_epochs: int = 100,
- random_state: int = 101,
- batch_size: int = 128,
- shuffle: bool = True,
- patience: int = 15,
- monitor: str = "val_loss",
- mode: str = "min",
- lr: float = 1e-4,
- lr_patience: int = 10,
- factor: float = 0.1,
- weight_decay: float = 1e-06,
- **trainer_kwargs
- ):
- """
- Fits the model to the provided data, using the specified loss distribution family for the prediction task.
-
- Parameters
- ----------
- X : DataFrame or array-like, shape (n_samples, n_features)
- Training features.
- y : array-like, shape (n_samples,) or (n_samples, n_targets)
- Target values for training.
- family : str
- The name of the distribution family to use for the loss function. Examples include 'normal' for regression tasks.
- val_size : float, default=0.2
- Proportion of the dataset to include in the validation split if `X_val` is None.
- X_val : DataFrame or array-like, shape (n_samples, n_features), optional
- Validation features. If provided, `X` and `y` are not split.
- y_val : array-like, shape (n_samples,) or (n_samples, n_targets), optional
- Validation target values. Required if `X_val` is provided.
- max_epochs : int, default=100
- Maximum number of epochs for training.
- random_state : int, default=101
- Seed used by the random number generator for shuffling the data.
- batch_size : int, default=64
- Number of samples per gradient update.
- shuffle : bool, default=True
- Whether to shuffle the training data before each epoch.
- patience : int, default=10
- Number of epochs with no improvement on the validation metric to wait before early stopping.
- monitor : str, default="val_loss"
- The metric to monitor for early stopping.
- mode : str, default="min"
- In 'min' mode, training will stop when the quantity monitored has stopped decreasing;
- in 'max' mode, it will stop when the quantity monitored has stopped increasing.
- lr : float, default=1e-3
- Learning rate for the optimizer.
- lr_patience : int, default=10
- Number of epochs with no improvement on the validation metric to wait before reducing the learning rate.
- factor : float, default=0.75
- Factor by which the learning rate will be reduced.
- weight_decay : float, default=0.025
- Weight decay (L2 penalty) parameter.
- **trainer_kwargs : dict
- Additional keyword arguments for PyTorch Lightning's Trainer class.
-
-
- Returns
- -------
- self : object
- The fitted estimator.
- """
- if not isinstance(X, pd.DataFrame):
- X = pd.DataFrame(X)
- if X_val:
- if not isinstance(X_val, pd.DataFrame):
- X_val = pd.DataFrame(X_val)
-
- if not X_val:
- X_train, X_val, y_train, y_val = self.split_data(
- X, y, val_size, random_state
- )
-
- data_module = self.preprocess_data(
- X_train, y_train, X_val, y_val, batch_size, shuffle
- )
-
- self.model = BaseMambularLSS(
- family=family,
- config=self.config,
- cat_feature_info=self.cat_feature_info,
- num_feature_info=self.num_feature_info,
- lr=lr,
- lr_patience=lr_patience,
- lr_factor=factor,
- weight_decay=weight_decay,
- )
-
- early_stop_callback = EarlyStopping(
- monitor=monitor, min_delta=0.00, patience=patience, verbose=False, mode=mode
- )
-
- checkpoint_callback = ModelCheckpoint(
- monitor="val_loss",
- mode="min",
- save_top_k=1,
- dirpath="model_checkpoints",
- filename="best_model",
- )
-
- # Initialize the trainer and train the model
- trainer = pl.Trainer(
- max_epochs=max_epochs,
- callbacks=[early_stop_callback, checkpoint_callback],
- **trainer_kwargs,
- )
- trainer.fit(self.model, data_module)
-
- best_model_path = checkpoint_callback.best_model_path
- if best_model_path:
- checkpoint = torch.load(best_model_path)
- self.model.load_state_dict(checkpoint["state_dict"])
-
- return self
-
-[docs] def predict(self, X, raw=False):
- """
- Predicts target values for the given input samples using the fitted model.
-
- Parameters
- ----------
- X : DataFrame or array-like, shape (n_samples, n_features)
- The input samples for which to predict target values.
-
-
- Returns
- -------
- predictions : ndarray, shape (n_samples,) or (n_samples, n_distributional_parameters)
- The predicted target values.
- """
- # Preprocess the data
- if not isinstance(X, pd.DataFrame):
- X = pd.DataFrame(X)
- device = next(self.model.parameters()).device
- cat_tensors, num_tensors = self.preprocess_test_data(X)
- if isinstance(cat_tensors, list):
- cat_tensors = [tensor.to(device) for tensor in cat_tensors]
- else:
- cat_tensors = cat_tensors.to(device)
-
- if isinstance(num_tensors, list):
- num_tensors = [tensor.to(device) for tensor in num_tensors]
- else:
- num_tensors = num_tensors.to(device)
-
- # Set the model to evaluation mode
- self.model.eval()
-
- # Perform inference
- with torch.no_grad():
- predictions = self.model(num_features=num_tensors, cat_features=cat_tensors)
-
- if not raw:
- return self.model.family(predictions).cpu().numpy()
-
- # Convert predictions to NumPy array and return
- else:
- return predictions.cpu().numpy()
-
-[docs] def evaluate(self, X, y_true, metrics=None, distribution_family=None):
- """
- Evaluate the model on the given data using specified metrics tailored to the distribution type.
-
- Parameters
- ----------
- X : DataFrame or array-like, shape (n_samples, n_features)
- Input samples.
- y_true : DataFrame or array-like, shape (n_samples,) or (n_samples, n_outputs)
- True target values.
- metrics : dict, optional
- A dictionary where keys are metric names and values are the metric functions.
- If None, default metrics based on the detected or specified distribution_family are used.
- distribution_family : str, optional
- Specifies the distribution family the model is predicting for. If None, it will attempt to infer based
- on the model's settings.
-
-
- Returns
- -------
- scores : dict
- A dictionary with metric names as keys and their corresponding scores as values.
- """
- # Infer distribution family from model settings if not provided
- if distribution_family is None:
- distribution_family = getattr(self.model, "distribution_family", "normal")
-
- # Setup default metrics if none are provided
- if metrics is None:
- metrics = self.get_default_metrics(distribution_family)
-
- # Make predictions
- predictions = self.predict(X, raw=False)
-
- # Initialize dictionary to store results
- scores = {}
-
- # Compute each metric
- for metric_name, metric_func in metrics.items():
- scores[metric_name] = metric_func(y_true, predictions)
-
- return scores
-
-[docs] def get_default_metrics(self, distribution_family):
- """
- Provides default metrics based on the distribution family.
-
- Parameters
- ----------
- distribution_family : str
- The distribution family for which to provide default metrics.
-
-
- Returns
- -------
- metrics : dict
- A dictionary of default metric functions.
- """
- default_metrics = {
- "normal": {
- "MSE": lambda y, pred: mean_squared_error(y, pred[:, 0]),
- "CRPS": lambda y, pred: np.mean(
- [
- ps.crps_gaussian(y[i], mu=pred[i, 0], sig=np.sqrt(pred[i, 1]))
- for i in range(len(y))
- ]
- ),
- },
- "poisson": {"Poisson Deviance": poisson_deviance},
- "gamma": {"Gamma Deviance": gamma_deviance},
- "beta": {"Brier Score": beta_brier_score},
- "dirichlet": {"Dirichlet Error": dirichlet_error},
- "studentt": {"Student-T Loss": student_t_loss},
- "negativebinom": {"Negative Binomial Deviance": negative_binomial_deviance},
- "inversegamma": {"Inverse Gamma Loss": inverse_gamma_loss},
- "categorical": {"Accuracy": accuracy_score},
- }
- return default_metrics.get(distribution_family, {})
-
-
-
-
-
-import lightning as pl
-import numpy as np
-import pandas as pd
-import warnings
-
-import torch
-from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint
-from sklearn.base import BaseEstimator
-from sklearn.decomposition import PCA
-from sklearn.metrics import accuracy_score
-from sklearn.model_selection import train_test_split
-from torch.utils.data import DataLoader
-
-from ..base_models.embedding_classifier import BaseEmbeddingMambularClassifier
-from ..utils.configs import DefaultMambularConfig
-from ..utils.dataset import EmbeddingMambularDataset, MambularDataModule
-from ..utils.preprocessor import Preprocessor
-
-
-[docs]class EmbeddingMambularClassifier(BaseEstimator):
- """
- Provides an scikit-learn-like interface for the ProteinMambularClassifier, making it compatible with
- scikit-learn's utilities and workflow. This class encapsulates the PyTorch Lightning model, preprocessing,
- and data loading, offering methods for fitting, predicting, and probability estimation in a manner akin
- to scikit-learn's API.
-
- Parameters
- ----------
- # configuration parameters
- lr : float, optional
- Learning rate for the optimizer. Default is 1e-4.
- lr_patience : int, optional
- Number of epochs with no improvement on the validation loss to wait before reducing the learning rate. Default is 10.
- weight_decay : float, optional
- Weight decay (L2 penalty) coefficient. Default is 1e-6.
- lr_factor : float, optional
- Factor by which the learning rate will be reduced. Default is 0.1.
- d_model : int, optional
- Dimension of the model. Default is 64.
- n_layers : int, optional
- Number of layers. Default is 8.
- expand_factor : int, optional
- Expansion factor. Default is 2.
- bias : bool, optional
- Whether to use bias. Default is False.
- d_conv : int, optional
- Dimension of the convolution. Default is 16.
- conv_bias : bool, optional
- Whether to use bias in the convolution. Default is True.
- dropout : float, optional
- Dropout rate in the mamba blocks. Default is 0.05.
- dt_rank : str, optional
- Rank of the time dimension. Default is "auto".
- d_state : int, optional
- State dimension. Default is 16.
- dt_scale : float, optional
- Scale of the time dimension. Default is 1.0.
- dt_init : str, optional
- Initialization method for the time dimension. Default is "random".
- dt_max : float, optional
- Maximum value for the time dimension. Default is 0.1.
- dt_min : float, optional
- Minimum value for the time dimension. Default is 1e-3.
- dt_init_floor : float, optional
- Floor value for the time dimension initialization. Default is 1e-4.
- norm : str, optional
- Normalization method. Default is 'RMSNorm'.
- activation : callable, optional
- Activation function. Default is nn.SELU().
- num_embedding_activation : callable, optional
- Activation function for numerical embeddings. Default is nn.Identity().
- head_layer_sizes : list, optional
- Sizes of the layers in the head. Default is [64, 64, 32].
- head_dropout : float, optional
- Dropout rate for the head. Default is 0.5.
- head_skip_layers : bool, optional
- Whether to use skip layers in the head. Default is False.
- head_activation : callable, optional
- Activation function for the head. Default is nn.SELU().
- head_use_batch_norm : bool, optional
- Whether to use batch normalization in the head. Default is False.
-
- # Preprocessor Parameters
- n_bins : int, optional
- The number of bins to use for numerical feature binning. Default is 50.
- numerical_preprocessing : str, optional
- The preprocessing strategy for numerical features. Default is 'ple'.
- use_decision_tree_bins : bool, optional
- If True, uses decision tree regression/classification to determine optimal bin edges for numerical feature binning. Default is False.
- binning_strategy : str, optional
- Defines the strategy for binning numerical features. Default is 'uniform'.
- task : str, optional
- Indicates the type of machine learning task ('regression' or 'classification'). Default is 'regression'.
- cat_cutoff: float or int, optional
- Indicates the cutoff after which integer values are treated as categorical. If float, it's treated as a percentage. If int, it's the maximum number of unique values for a column to be considered categorical. Default is 3%
- treat_all_integers_as_numerical : bool, optional
- If True, all integer columns will be treated as numerical, regardless of their unique value count or proportion. Default is False
-
-
- Attributes
- ----------
- config : MambularConfig
- Configuration object containing model-specific parameters.
- preprocessor : Preprocessor
- Preprocessor object for data preprocessing steps.
- model : BaseEmbeddingMambularRegressor
- The neural network model, initialized after the `fit` method is called.
- """
-
- def __init__(self, **kwargs):
- # Known config arguments
- config_arg_names = [
- "lr",
- "lr_patience",
- "weight_decay",
- "lr_factor",
- "d_model",
- "n_layers",
- "expand_factor",
- "bias",
- "d_conv",
- "conv_bias",
- "dropout",
- "dt_rank",
- "d_state",
- "dt_scale",
- "dt_init",
- "dt_max",
- "dt_min",
- "dt_init_floor",
- "norm",
- "activation",
- "num_embedding_activation",
- "head_layer_sizes",
- "head_dropout",
- "head_skip_layers",
- "head_activation",
- "head_use_batch_norm",
- ]
-
- preprocessor_arg_names = [
- "n_bins",
- "numerical_preprocessing",
- "use_decision_tree_bins",
- "binning_strategy",
- "task",
- "cat_cutoff",
- "treat_all_integers_as_numerical",
- ]
-
- self.config_kwargs = {k: v for k, v in kwargs.items() if k in config_arg_names}
- self.config = DefaultMambularConfig(**self.config_kwargs)
-
- preprocessor_kwargs = {
- k: v for k, v in kwargs.items() if k in preprocessor_arg_names
- }
- if "numerical_preprocessing" not in list(preprocessor_kwargs.keys()):
- preprocessor_kwargs["numerical_preprocessing"] = "standardization"
-
- # Raise a warning if task is set to 'classification'
- if preprocessor_kwargs.get("task") == "regression":
- warnings.warn(
- "The task is set to 'regression'. This model is designed for classification tasks.",
- UserWarning,
- )
-
- if "task" not in list(preprocessor_kwargs.keys()):
- preprocessor_kwargs["task"] = "classification"
-
- self.preprocessor = Preprocessor(**preprocessor_kwargs)
- self.model = None
-
-[docs] def get_params(self, deep=True):
- """
- Get parameters for this estimator.
-
- Parameters
- ----------
- deep : bool, default=True
- If True, will return the parameters for this estimator and contained subobjects that are estimators.
-
-
- Returns
- -------
- params : dict
- Parameter names mapped to their values.
- """
- params = self.config_kwargs # Parameters used to initialize MambularConfig
-
- # If deep=True, include parameters from nested components like preprocessor
- if deep:
- # Assuming Preprocessor has a get_params method
- preprocessor_params = {
- "preprocessor__" + key: value
- for key, value in self.preprocessor.get_params().items()
- }
- params.update(preprocessor_params)
-
- return params
-
-[docs] def set_params(self, **parameters):
- """
- Set the parameters of this estimator.
-
- Parameters
- ----------
- **parameters : dict
- Estimator parameters.
-
-
- Returns
- -------
- self : object
- Estimator instance.
- """
- # Update config_kwargs with provided parameters
- valid_config_keys = self.config_kwargs.keys()
- config_updates = {k: v for k, v in parameters.items() if k in valid_config_keys}
- self.config_kwargs.update(config_updates)
-
- # Update the config object
- for key, value in config_updates.items():
- setattr(self.config, key, value)
-
- # Handle preprocessor parameters (prefixed with 'preprocessor__')
- preprocessor_params = {
- k.split("__")[1]: v
- for k, v in parameters.items()
- if k.startswith("preprocessor__")
- }
- if preprocessor_params:
- # Assuming Preprocessor has a set_params method
- self.preprocessor.set_params(**preprocessor_params)
-
- return self
-
-[docs] def split_data(self, X, y, val_size, random_state):
- """
- Split the dataset into training and validation sets.
-
- Parameters
- ----------
- X : array-like of shape (n_samples, n_features)
- The input samples.
- y : array-like of shape (n_samples,)
- The target values.
- val_size : float
- The proportion of the dataset to include in the validation split.
- random_state : int
- Controls the shuffling applied to the data before applying the split.
-
-
- Returns
- -------
- X_train, X_val, y_train, y_val : arrays
- The split datasets.
- """
- X_train, X_val, y_train, y_val = train_test_split(
- X, y, test_size=val_size, random_state=random_state
- )
-
- return X_train, X_val, y_train, y_val
-
-[docs] def preprocess_data(self, X_train, y_train, X_val, y_val, batch_size, shuffle):
- """
- Preprocess the training and validation data and create corresponding DataLoaders.
-
- Parameters
- ----------
- X_train : array-like of shape (n_samples, n_features)
- The training input samples.
- y_train : array-like of shape (n_samples,)
- The training target values.
- X_val : array-like of shape (n_samples, n_features)
- The validation input samples.
- y_val : array-like of shape (n_samples,)
- The validation target values.
- batch_size : int
- Size of mini-batches for the DataLoader.
- shuffle : bool
- Whether to shuffle the training data before splitting into batches.
-
-
- Returns
- -------
- data_module : MambularDataModule
- An instance of MambularDataModule containing training and validation DataLoaders.
- """
- self.preprocessor.fit(
- pd.concat([X_train, X_val], axis=0).reset_index(drop=True),
- np.concatenate((y_train, y_val), axis=0),
- )
- train_preprocessed_data = self.preprocessor.transform(X_train)
- val_preprocessed_data = self.preprocessor.transform(X_val)
-
- # Update feature info based on the actual processed data
- (
- self.cat_feature_info,
- self.num_feature_info,
- ) = self.preprocessor.get_feature_info()
-
- # Initialize lists for tensors
- train_cat_tensors = []
- train_num_tensors = []
- val_cat_tensors = []
- val_num_tensors = []
-
- # Populate tensors for categorical features, if present in processed data
- for key in self.cat_feature_info:
- cat_key = "cat_" + key # Assuming categorical keys are prefixed with 'cat_'
- if cat_key in train_preprocessed_data:
- train_cat_tensors.append(
- torch.tensor(train_preprocessed_data[cat_key], dtype=torch.long)
- )
- if cat_key in val_preprocessed_data:
- val_cat_tensors.append(
- torch.tensor(val_preprocessed_data[cat_key], dtype=torch.long)
- )
-
- binned_key = "num_" + key # for binned features
- if binned_key in train_preprocessed_data:
- train_cat_tensors.append(
- torch.tensor(train_preprocessed_data[binned_key], dtype=torch.long)
- )
-
- if binned_key in val_preprocessed_data:
- val_cat_tensors.append(
- torch.tensor(val_preprocessed_data[binned_key], dtype=torch.long)
- )
-
- # Populate tensors for numerical features, if present in processed data
- for key in self.num_feature_info:
- num_key = "num_" + str(
- key
- ) # Assuming numerical keys are prefixed with 'num_'
- if num_key in train_preprocessed_data:
- train_num_tensors.append(
- torch.tensor(train_preprocessed_data[num_key], dtype=torch.float32)
- )
- if num_key in val_preprocessed_data:
- val_num_tensors.append(
- torch.tensor(val_preprocessed_data[num_key], dtype=torch.float32)
- )
-
- train_labels = torch.tensor(y_train, dtype=torch.long)
- val_labels = torch.tensor(y_val, dtype=torch.long)
-
- # Create datasets
- train_dataset = EmbeddingMambularDataset(
- train_cat_tensors, train_num_tensors, train_labels, regression=False
- )
- val_dataset = EmbeddingMambularDataset(
- val_cat_tensors, val_num_tensors, val_labels, regression=False
- )
-
- # Create dataloaders
- train_dataloader = DataLoader(
- train_dataset, batch_size=batch_size, shuffle=shuffle
- )
- val_dataloader = DataLoader(val_dataset, batch_size=batch_size)
-
- return MambularDataModule(train_dataloader, val_dataloader)
-
-[docs] def preprocess_test_data(self, X):
- """
- Preprocesses the test data and creates tensors for categorical and numerical features.
-
- Parameters
- ----------
- X : DataFrame or array-like, shape (n_samples, n_features)
- Test feature set.
-
-
- Returns
- -------
- cat_tensors : list of Tensors
- List of tensors for each categorical feature.
- num_tensors : list of Tensors
- List of tensors for each numerical feature.
- """
- processed_data = self.preprocessor.transform(X)
-
- # Initialize lists for tensors
- cat_tensors = []
- num_tensors = []
-
- # Populate tensors for categorical features
- for key in self.cat_feature_info:
- cat_key = "cat_" + str(
- key
- ) # Assuming categorical keys are prefixed with 'cat_'
- if cat_key in processed_data:
- cat_tensors.append(
- torch.tensor(processed_data[cat_key], dtype=torch.long)
- )
-
- binned_key = "num_" + str(key) # for binned features
- if binned_key in processed_data:
- cat_tensors.append(
- torch.tensor(processed_data[binned_key], dtype=torch.long)
- )
-
- # Populate tensors for numerical features
- for key in self.num_feature_info:
- num_key = "num_" + str(
- key
- ) # Assuming numerical keys are prefixed with 'num_'
- if num_key in processed_data:
- num_tensors.append(
- torch.tensor(processed_data[num_key], dtype=torch.float32)
- )
-
- return cat_tensors, num_tensors
-
-[docs] def fit(
- self,
- X,
- y,
- val_size=0.2,
- X_val=None,
- y_val=None,
- max_epochs=100,
- random_state=101,
- batch_size=64,
- shuffle=True,
- patience=10,
- monitor="val_loss",
- mode="min",
- lr=1e-3,
- lr_patience=10,
- factor=0.75,
- weight_decay=0.025,
- raw_embeddings=False,
- seq_size=20,
- pca=False,
- reduced_dims=16,
- **trainer_kwargs
- ):
- """
- Fits the model to the given dataset.
-
- Parameters
- ----------
- X : pandas DataFrame or array-like
- Feature matrix for training.
- y : array-like
- Target vector.
- val_size : float, optional
- Fraction of the data to use for validation if X_val is None.
- X_val : pandas DataFrame or array-like, optional
- Feature matrix for validation.
- y_val : array-like, optional
- Target vector for validation.
- max_epochs : int, default=100
- Maximum number of epochs for training.
- random_state : int, optional
- Seed for random number generators.
- batch_size : int, default=32
- Size of batches for training and validation.
- shuffle : bool, default=True
- Whether to shuffle training data before each epoch.
- patience : int, default=10
- Patience for early stopping based on val_loss.
- monitor : str, default='val_loss'
- Metric to monitor for early stopping.
- mode : str, default='min'
- Mode for early stopping ('min' or 'max').
- lr : float, default=0.001
- Learning rate for the optimizer.
- lr_patience : int, default=5
- Patience for learning rate reduction.
- factor : float, default=0.1
- Factor for learning rate reduction.
- weight_decay : float, default=0.0
- Weight decay for the optimizer.
- raw_embeddings : bool, default=False
- Whether to use raw features or embeddings.
- seq_size : int, optional
- Sequence size for embeddings, relevant if raw_embeddings is False.
- **trainer_kwargs : dict
- Additional arguments for the PyTorch Lightning Trainer.
-
-
- Returns
- -------
- self : object
- The fitted estimator.
- """
- if not isinstance(X, pd.DataFrame):
- X = pd.DataFrame(X)
- if X_val:
- if not isinstance(X_val, pd.DataFrame):
- X_val = pd.DataFrame(X_val)
-
- # Apply PCA if indicated
- if pca:
- pca_transformer = PCA(n_components=reduced_dims)
- X = pca_transformer.fit_transform(
- X
- ) # Fit and transform the PCA on the complete dataset
- if X_val is not None:
- X_val = pca_transformer.transform(
- X_val
- ) # Transform validation data with the same PCA model
-
- raw_embeddings = True
-
- if not X_val:
- X_train, X_val, y_train, y_val = self.split_data(
- X, y, val_size, random_state
- )
- else:
- X_train = X
- y_train = y
-
- data_module = self.preprocess_data(
- X_train, y_train, X_val, y_val, batch_size, shuffle
- )
-
- if raw_embeddings:
- self.config.d_model = X.shape[1]
-
- num_classes = len(np.unique(y))
-
- self.model = BaseEmbeddingMambularClassifier(
- num_classes=num_classes,
- config=self.config,
- cat_feature_info=self.cat_feature_info,
- num_feature_info=self.num_feature_info,
- lr=lr,
- lr_patience=lr_patience,
- lr_factor=factor,
- weight_decay=weight_decay,
- raw_embeddings=raw_embeddings,
- seq_size=seq_size,
- )
-
- early_stop_callback = EarlyStopping(
- monitor=monitor, min_delta=0.00, patience=patience, verbose=False, mode=mode
- )
-
- checkpoint_callback = ModelCheckpoint(
- monitor="val_loss", # Adjust according to your validation metric
- mode="min",
- save_top_k=1,
- dirpath="model_checkpoints", # Specify the directory to save checkpoints
- filename="best_model",
- )
-
- # Initialize the trainer and train the model
- trainer = pl.Trainer(
- max_epochs=max_epochs,
- callbacks=[early_stop_callback, checkpoint_callback],
- **trainer_kwargs
- )
- trainer.fit(self.model, data_module)
-
- best_model_path = checkpoint_callback.best_model_path
- if best_model_path:
- checkpoint = torch.load(best_model_path)
- self.model.load_state_dict(checkpoint["state_dict"])
-
- return self
-
-[docs] def predict(self, X):
- """
- Predict the class labels for the given input samples.
-
- Parameters
- ----------
- X : array-like or pd.DataFrame of shape (n_samples, n_features)
- The input samples to predict.
-
-
- Returns
- -------
- predictions : ndarray of shape (n_samples,)
- Predicted class labels for each input sample.
-
-
- Notes
- -----
- The method preprocesses the input data using the same preprocessor used during training,
- sets the model to evaluation mode, and then performs inference to predict the class labels.
- The predictions are converted from a PyTorch tensor to a NumPy array before being returned.
- """
- # Preprocess the data
- if not isinstance(X, pd.DataFrame):
- X = pd.DataFrame(X)
-
- if hasattr(self, "pca_transformer"):
- X = pd.DataFrame(self.pca_transformer.transform(X))
-
- cat_tensors, num_tensors = self.preprocess_test_data(X)
- device = next(self.model.parameters()).device
- cat_tensors, num_tensors = self.preprocess_test_data(X)
- if isinstance(cat_tensors, list):
- cat_tensors = [tensor.to(device) for tensor in cat_tensors]
- else:
- cat_tensors = cat_tensors.to(device)
-
- if isinstance(num_tensors, list):
- num_tensors = [tensor.to(device) for tensor in num_tensors]
- else:
- num_tensors = num_tensors.to(device)
-
- # Set the model to evaluation mode
- self.model.eval()
-
- # Perform inference
- with torch.no_grad():
- logits = self.model(num_features=num_tensors, cat_features=cat_tensors)
-
- # Check the shape of the logits to determine binary or multi-class classification
- if logits.shape[1] == 1:
- # Binary classification
- probabilities = torch.sigmoid(logits)
- predictions = (probabilities > 0.5).long().squeeze()
- else:
- # Multi-class classification
- probabilities = torch.softmax(logits, dim=1)
- predictions = torch.argmax(probabilities, dim=1)
-
- # Convert predictions to NumPy array and return
- return predictions.cpu().numpy()
-
-[docs] def predict_proba(self, X):
- """
- Predict class probabilities for the given input samples.
-
- Parameters
- ----------
- X : array-like or pd.DataFrame of shape (n_samples, n_features)
- The input samples for which to predict class probabilities.
-
-
- Returns
- -------
- probabilities : ndarray of shape (n_samples, n_classes)
- Predicted class probabilities for each input sample.
-
-
- Notes
- -----
- The method preprocesses the input data using the same preprocessor used during training,
- sets the model to evaluation mode, and then performs inference to predict the class probabilities.
- Softmax is applied to the logits to obtain probabilities, which are then converted from a PyTorch tensor
- to a NumPy array before being returned.
- """
- # Preprocess the data
- if not isinstance(X, pd.DataFrame):
- X = pd.DataFrame(X)
- cat_tensors, num_tensors = self.preprocess_test_data(X)
- device = next(self.model.parameters()).device
- cat_tensors, num_tensors = self.preprocess_test_data(X)
- if isinstance(cat_tensors, list):
- cat_tensors = [tensor.to(device) for tensor in cat_tensors]
- else:
- cat_tensors = cat_tensors.to(device)
-
- if isinstance(num_tensors, list):
- num_tensors = [tensor.to(device) for tensor in num_tensors]
- else:
- num_tensors = num_tensors.to(device)
-
- # Set the model to evaluation mode
- self.model.eval()
-
- # Perform inference
- with torch.no_grad():
- logits = self.model(num_features=num_tensors, cat_features=cat_tensors)
- if logits.shape[1] > 1:
- probabilities = torch.softmax(logits, dim=1)
- else:
- probabilities = torch.sigmoid(logits)
-
- # Convert probabilities to NumPy array and return
- return probabilities.cpu().numpy()
-
-[docs] def evaluate(self, X, y_true, metrics=None):
- """
- Evaluate the model on the given data using specified metrics.
-
- Parameters
- ----------
- X : array-like or pd.DataFrame of shape (n_samples, n_features)
- The input samples to predict.
- y_true : array-like of shape (n_samples,)
- The true class labels against which to evaluate the predictions.
- metrics : dict
- A dictionary where keys are metric names and values are tuples containing the metric function
- and a boolean indicating whether the metric requires probability scores (True) or class labels (False).
-
-
- Returns
- -------
- scores : dict
- A dictionary with metric names as keys and their corresponding scores as values.
-
-
- Notes
- -----
- This method uses either the `predict` or `predict_proba` method depending on the metric requirements.
- """
- # Ensure input is in the correct format
- if metrics is None:
- metrics = {"Accuracy": (accuracy_score, False)}
-
- if not isinstance(X, pd.DataFrame):
- X = pd.DataFrame(X)
-
- # Initialize dictionary to store results
- scores = {}
-
- # Generate class probabilities if any metric requires them
- if any(use_proba for _, use_proba in metrics.values()):
- probabilities = self.predict_proba(X)
-
- # Generate class labels if any metric requires them
- if any(not use_proba for _, use_proba in metrics.values()):
- predictions = self.predict(X)
-
- # Compute each metric
- for metric_name, (metric_func, use_proba) in metrics.items():
- if use_proba:
- scores[metric_name] = metric_func(y_true, probabilities)
- else:
- scores[metric_name] = metric_func(y_true, predictions)
-
- return scores
-
-
-
-
-
-import warnings
-
-import lightning as pl
-import numpy as np
-import pandas as pd
-import torch
-from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint
-from sklearn.base import BaseEstimator
-from sklearn.decomposition import PCA
-from sklearn.metrics import mean_squared_error
-from sklearn.model_selection import train_test_split
-from torch.utils.data import DataLoader
-
-from ..base_models.embedding_regressor import BaseEmbeddingMambularRegressor
-from ..utils.configs import DefaultMambularConfig
-from ..utils.dataset import EmbeddingMambularDataset, MambularDataModule
-from ..utils.preprocessor import Preprocessor
-
-
-[docs]class EmbeddingMambularRegressor(BaseEstimator):
- """
- An sklearn-like interface for the ProteinMambularRegressor, making it compatible with sklearn's utilities
- and workflows. This class wraps the PyTorch Lightning model and preprocessor, providing methods for fitting,
- predicting, and setting/getting parameters in a way that mimics sklearn's API.
-
- Parameters
- ----------
- # configuration parameters
- lr : float, optional
- Learning rate for the optimizer. Default is 1e-4.
- lr_patience : int, optional
- Number of epochs with no improvement on the validation loss to wait before reducing the learning rate. Default is 10.
- weight_decay : float, optional
- Weight decay (L2 penalty) coefficient. Default is 1e-6.
- lr_factor : float, optional
- Factor by which the learning rate will be reduced. Default is 0.1.
- d_model : int, optional
- Dimension of the model. Default is 64.
- n_layers : int, optional
- Number of layers. Default is 8.
- expand_factor : int, optional
- Expansion factor. Default is 2.
- bias : bool, optional
- Whether to use bias. Default is False.
- d_conv : int, optional
- Dimension of the convolution. Default is 16.
- conv_bias : bool, optional
- Whether to use bias in the convolution. Default is True.
- dropout : float, optional
- Dropout rate in the mamba blocks. Default is 0.05.
- dt_rank : str, optional
- Rank of the time dimension. Default is "auto".
- d_state : int, optional
- State dimension. Default is 16.
- dt_scale : float, optional
- Scale of the time dimension. Default is 1.0.
- dt_init : str, optional
- Initialization method for the time dimension. Default is "random".
- dt_max : float, optional
- Maximum value for the time dimension. Default is 0.1.
- dt_min : float, optional
- Minimum value for the time dimension. Default is 1e-3.
- dt_init_floor : float, optional
- Floor value for the time dimension initialization. Default is 1e-4.
- norm : str, optional
- Normalization method. Default is 'RMSNorm'.
- activation : callable, optional
- Activation function. Default is nn.SELU().
- num_embedding_activation : callable, optional
- Activation function for numerical embeddings. Default is nn.Identity().
- head_layer_sizes : list, optional
- Sizes of the layers in the head. Default is [64, 64, 32].
- head_dropout : float, optional
- Dropout rate for the head. Default is 0.5.
- head_skip_layers : bool, optional
- Whether to use skip layers in the head. Default is False.
- head_activation : callable, optional
- Activation function for the head. Default is nn.SELU().
- head_use_batch_norm : bool, optional
- Whether to use batch normalization in the head. Default is False.
-
- # Preprocessor Parameters
- n_bins : int, optional
- The number of bins to use for numerical feature binning. Default is 50.
- numerical_preprocessing : str, optional
- The preprocessing strategy for numerical features. Default is 'ple'.
- use_decision_tree_bins : bool, optional
- If True, uses decision tree regression/classification to determine optimal bin edges for numerical feature binning. Default is False.
- binning_strategy : str, optional
- Defines the strategy for binning numerical features. Default is 'uniform'.
- task : str, optional
- Indicates the type of machine learning task ('regression' or 'classification'). Default is 'regression'.
- cat_cutoff: float or int, optional
- Indicates the cutoff after which integer values are treated as categorical. If float, it's treated as a percentage. If int, it's the maximum number of unique values for a column to be considered categorical. Default is 3%
- treat_all_integers_as_numerical : bool, optional
- If True, all integer columns will be treated as numerical, regardless of their unique value count or proportion. Default is False
-
-
- Attributes
- ----------
- config : MambularConfig
- Configuration object containing model-specific parameters.
- preprocessor : Preprocessor
- Preprocessor object for data preprocessing steps.
- model : BaseEmbeddingMambularRegressor
- The neural network model, initialized after the `fit` method is called.
- """
-
- def __init__(self, **kwargs):
- # Known config arguments
- config_arg_names = [
- "lr",
- "lr_patience",
- "weight_decay",
- "lr_factor",
- "d_model",
- "n_layers",
- "expand_factor",
- "bias",
- "d_conv",
- "conv_bias",
- "dropout",
- "dt_rank",
- "d_state",
- "dt_scale",
- "dt_init",
- "dt_max",
- "dt_min",
- "dt_init_floor",
- "norm",
- "activation",
- "num_embedding_activation",
- "head_layer_sizes",
- "head_dropout",
- "head_skip_layers",
- "head_activation",
- "head_use_batch_norm",
- ]
-
- preprocessor_arg_names = [
- "n_bins",
- "numerical_preprocessing",
- "use_decision_tree_bins",
- "binning_strategy",
- "task",
- "cat_cutoff",
- "treat_all_integers_as_numerical",
- ]
-
- self.config_kwargs = {k: v for k,
- v in kwargs.items() if k in config_arg_names}
- self.config = DefaultMambularConfig(**self.config_kwargs)
-
- preprocessor_kwargs = {
- k: v for k, v in kwargs.items() if k in preprocessor_arg_names
- }
- if "numerical_preprocessing" not in list(preprocessor_kwargs.keys()):
- preprocessor_kwargs["numerical_preprocessing"] = "standardization"
-
- self.preprocessor = Preprocessor(**preprocessor_kwargs)
- self.model = None
-
- # Raise a warning if task is set to 'classification'
- if preprocessor_kwargs.get("task") == "classification":
- warnings.warn(
- "The task is set to 'classification'. MambularRegressor is designed for regression tasks.",
- UserWarning,
- )
-
-[docs] def get_params(self, deep=True):
- """
- Get parameters for this estimator. Overrides the BaseEstimator method.
-
- Parameters
- ----------
- deep : bool, default=True
- If True, returns the parameters for this estimator and contained sub-objects that are estimators.
-
- Returns
- -------
- params : dict
- Parameter names mapped to their values.
- """
- params = self.config_kwargs # Parameters used to initialize MambularConfig
-
- # If deep=True, include parameters from nested components like preprocessor
- if deep:
- # Assuming Preprocessor has a get_params method
- preprocessor_params = {
- "preprocessor__" + key: value
- for key, value in self.preprocessor.get_params().items()
- }
- params.update(preprocessor_params)
-
- return params
-
-[docs] def set_params(self, **parameters):
- """
- Set the parameters of this estimator. Overrides the BaseEstimator method.
-
- Parameters
- ----------
- **parameters : dict
- Estimator parameters to be set.
-
-
- Returns
- -------
- self : object
- The instance with updated parameters.
- """
- # Update config_kwargs with provided parameters
- valid_config_keys = self.config_kwargs.keys()
- config_updates = {k: v for k,
- v in parameters.items() if k in valid_config_keys}
- self.config_kwargs.update(config_updates)
-
- # Update the config object
- for key, value in config_updates.items():
- setattr(self.config, key, value)
-
- # Handle preprocessor parameters (prefixed with 'preprocessor__')
- preprocessor_params = {
- k.split("__")[1]: v
- for k, v in parameters.items()
- if k.startswith("preprocessor__")
- }
- if preprocessor_params:
- # Assuming Preprocessor has a set_params method
- self.preprocessor.set_params(**preprocessor_params)
-
- return self
-
-[docs] def split_data(self, X, y, val_size, random_state):
- """
- Splits the dataset into training and validation sets.
-
- Parameters
- ----------
- X : array-like or DataFrame, shape (n_samples, n_features)
- Input features.
- y : array-like, shape (n_samples,) or (n_samples, n_targets)
- Target values.
- val_size : float
- The proportion of the dataset to include in the validation split.
- random_state : int
- Controls the shuffling applied to the data before applying the split.
-
-
- Returns
- -------
- X_train, X_val, y_train, y_val : arrays
- The split datasets.
- """
- X_train, X_val, y_train, y_val = train_test_split(
- X, y, test_size=val_size, random_state=random_state
- )
-
- return X_train, X_val, y_train, y_val
-
-[docs] def preprocess_data(self, X_train, y_train, X_val, y_val, batch_size, shuffle):
- """
- Preprocesses the training and validation data, and creates DataLoaders for them.
-
- Parameters
- ----------
- X_train : DataFrame or array-like, shape (n_samples_train, n_features)
- Training feature set.
- y_train : array-like, shape (n_samples_train,)
- Training target values.
- X_val : DataFrame or array-like, shape (n_samples_val, n_features)
- Validation feature set.
- y_val : array-like, shape (n_samples_val,)
- Validation target values.
- batch_size : int
- Size of batches for the DataLoader.
- shuffle : bool
- Whether to shuffle the training data in the DataLoader.
-
-
- Returns
- -------
- data_module : MambularDataModule
- An instance of MambularDataModule containing the training and validation DataLoaders.
- """
- self.preprocessor.fit(
- pd.concat([X_train, X_val], axis=0).reset_index(drop=True),
- np.concatenate((y_train, y_val), axis=0),
- )
- train_preprocessed_data = self.preprocessor.transform(X_train)
- val_preprocessed_data = self.preprocessor.transform(X_val)
-
- # Update feature info based on the actual processed data
- (
- self.cat_feature_info,
- self.num_feature_info,
- ) = self.preprocessor.get_feature_info()
-
- # Initialize lists for tensors
- train_cat_tensors = []
- train_num_tensors = []
- val_cat_tensors = []
- val_num_tensors = []
-
- # Populate tensors for categorical features, if present in processed data
- for key in self.cat_feature_info:
- cat_key = "cat_" + key # Assuming categorical keys are prefixed with 'cat_'
- if cat_key in train_preprocessed_data:
- train_cat_tensors.append(
- torch.tensor(
- train_preprocessed_data[cat_key], dtype=torch.long)
- )
- if cat_key in val_preprocessed_data:
- val_cat_tensors.append(
- torch.tensor(
- val_preprocessed_data[cat_key], dtype=torch.long)
- )
-
- binned_key = "num_" + key # for binned features
- if binned_key in train_preprocessed_data:
- train_cat_tensors.append(
- torch.tensor(
- train_preprocessed_data[binned_key], dtype=torch.long)
- )
-
- if binned_key in val_preprocessed_data:
- val_cat_tensors.append(
- torch.tensor(
- val_preprocessed_data[binned_key], dtype=torch.long)
- )
-
- # Populate tensors for numerical features, if present in processed data
- for key in self.num_feature_info:
- num_key = "num_" + str(
- key
- ) # Assuming numerical keys are prefixed with 'num_'
- if num_key in train_preprocessed_data:
- train_num_tensors.append(
- torch.tensor(
- train_preprocessed_data[num_key], dtype=torch.float32)
- )
- if num_key in val_preprocessed_data:
- val_num_tensors.append(
- torch.tensor(
- val_preprocessed_data[num_key], dtype=torch.float32)
- )
-
- train_labels = torch.tensor(y_train, dtype=torch.float32)
- val_labels = torch.tensor(y_val, dtype=torch.float32)
-
- # Create datasets
- train_dataset = EmbeddingMambularDataset(
- train_cat_tensors, train_num_tensors, train_labels
- )
- val_dataset = EmbeddingMambularDataset(
- val_cat_tensors, val_num_tensors, val_labels
- )
-
- # Create dataloaders
- train_dataloader = DataLoader(
- train_dataset, batch_size=batch_size, shuffle=shuffle
- )
- val_dataloader = DataLoader(val_dataset, batch_size=batch_size)
-
- return MambularDataModule(train_dataloader, val_dataloader)
-
-[docs] def preprocess_test_data(self, X):
- """
- Preprocesses the test data and creates tensors for categorical and numerical features.
-
- Parameters
- ----------
- X : DataFrame or array-like, shape (n_samples, n_features)
- Test feature set.
-
-
- Returns
- -------
- cat_tensors : list of Tensors
- List of tensors for each categorical feature.
- num_tensors : list of Tensors
- List of tensors for each numerical feature.
- """
- processed_data = self.preprocessor.transform(X)
-
- # Initialize lists for tensors
- cat_tensors = []
- num_tensors = []
-
- # Populate tensors for categorical features
- for key in self.cat_feature_info:
- cat_key = "cat_" + str(
- key
- ) # Assuming categorical keys are prefixed with 'cat_'
- if cat_key in processed_data:
- cat_tensors.append(
- torch.tensor(processed_data[cat_key], dtype=torch.long)
- )
-
- binned_key = "num_" + str(key) # for binned features
- if binned_key in processed_data:
- cat_tensors.append(
- torch.tensor(processed_data[binned_key], dtype=torch.long)
- )
-
- # Populate tensors for numerical features
- for key in self.num_feature_info:
- num_key = "num_" + str(
- key
- ) # Assuming numerical keys are prefixed with 'num_'
- if num_key in processed_data:
- num_tensors.append(
- torch.tensor(processed_data[num_key], dtype=torch.float32)
- )
-
- return cat_tensors, num_tensors
-
-[docs] def fit(
- self,
- X,
- y,
- val_size=0.2,
- X_val=None,
- y_val=None,
- max_epochs=100,
- random_state=101,
- batch_size=64,
- shuffle=True,
- patience=10,
- monitor="val_loss",
- mode="min",
- lr=1e-3,
- lr_patience=10,
- factor=0.75,
- weight_decay=0.025,
- raw_embeddings=False,
- seq_size=20,
- pca=False,
- **trainer_kwargs
- ):
- """
- Fits the ProteinMambularRegressor model to the training data.
-
- Parameters
- ----------
- X : array-like or DataFrame
- The training input samples.
- y : array-like
- The target values (class labels for classification, real numbers for regression).
- val_size : float, optional
- The proportion of the dataset to include in the validation split if `X_val` is not provided.
- X_val : array-like or DataFrame, optional
- The validation input samples.
- y_val : array-like, optional
- The validation target values.
- max_epochs : int, optional
- The maximum number of epochs for training.
- random_state : int, optional
- The seed used by the random number generator.
- batch_size : int, optional
- Size of the batches for training.
- shuffle : bool, optional
- Whether to shuffle the training data.
- patience : int, optional
- Patience for early stopping.
- monitor : str, optional
- Quantity to be monitored for early stopping.
- mode : str, optional
- One of {'auto', 'min', 'max'}. In 'min' mode, training will stop when the quantity monitored has stopped decreasing;
- in 'max' mode, it will stop when the quantity monitored has stopped increasing.
- lr : float, optional
- Learning rate for the optimizer.
- lr_patience : int, optional
- Number of epochs with no improvement after which the learning rate will be reduced.
- factor : float, optional
- Factor by which the learning rate will be reduced.
- weight_decay : float, optional
- Weight decay coefficient for regularization in the optimizer.
- raw_embeddings : bool, optional
- Whether to use raw numerical features directly or to process them into embeddings.
- seq_size : int, optional
- The sequence size for processing numerical features when not using raw embeddings.
- **trainer_kwargs : dict
- Additional keyword arguments for the PyTorch Lightning Trainer.
-
- Returns
- -------
- self : object
- Returns an instance of self.
- """
-
- if not isinstance(X, pd.DataFrame):
- X = pd.DataFrame(X)
- if X_val:
- if not isinstance(X_val, pd.DataFrame):
- X_val = pd.DataFrame(X_val)
-
- # Apply PCA if indicated
- if pca:
- self.pca_transformer = PCA(n_components=seq_size)
- X = pd.DataFrame(
- self.pca_transformer.fit_transform(X)
- ) # Fit and transform the PCA on the complete dataset
- if X_val is not None:
- X_val = pd.DataFrame(
- self.pca_transformer.transform(X_val)
- ) # Transform validation data with the same PCA model
-
- raw_embeddings = True
-
- if not X_val:
- X_train, X_val, y_train, y_val = self.split_data(
- X, y, val_size, random_state
- )
- else:
- X_train = X
- y_train = y
-
- data_module = self.preprocess_data(
- X_train, y_train, X_val, y_val, batch_size, shuffle
- )
-
- if raw_embeddings:
- self.config.d_model = X.shape[1]
-
- self.model = BaseEmbeddingMambularRegressor(
- config=self.config,
- cat_feature_info=self.cat_feature_info,
- num_feature_info=self.num_feature_info,
- lr=lr,
- lr_patience=lr_patience,
- lr_factor=factor,
- weight_decay=weight_decay,
- raw_embeddings=raw_embeddings,
- seq_size=seq_size,
- )
-
- early_stop_callback = EarlyStopping(
- monitor=monitor, min_delta=0.00, patience=patience, verbose=False, mode=mode
- )
-
- checkpoint_callback = ModelCheckpoint(
- monitor="val_loss", # Adjust according to your validation metric
- mode="min",
- save_top_k=1,
- dirpath="model_checkpoints", # Specify the directory to save checkpoints
- filename="best_model",
- )
-
- # Initialize the trainer and train the model
- trainer = pl.Trainer(
- max_epochs=max_epochs,
- callbacks=[early_stop_callback, checkpoint_callback],
- **trainer_kwargs
- )
- trainer.fit(self.model, data_module)
-
- best_model_path = checkpoint_callback.best_model_path
- if best_model_path:
- checkpoint = torch.load(best_model_path)
- self.model.load_state_dict(checkpoint["state_dict"])
-
- return self
-
-[docs] def predict(self, X):
- """
- Predicts target values for the given input samples.
-
- Parameters
- ----------
- X : DataFrame or array-like, shape (n_samples, n_features)
- The input samples for which to predict target values.
-
-
- Returns
- -------
- predictions : ndarray, shape (n_samples,) or (n_samples, n_outputs)
- The predicted target values.
- """
- # Preprocess the data
- if not isinstance(X, pd.DataFrame):
- X = pd.DataFrame(X)
-
- if hasattr(self, "pca_transformer"):
- X = pd.DataFrame(self.pca_transformer.transform(X))
-
- device = next(self.model.parameters()).device
- cat_tensors, num_tensors = self.preprocess_test_data(X)
- if isinstance(cat_tensors, list):
- cat_tensors = [tensor.to(device) for tensor in cat_tensors]
- else:
- cat_tensors = cat_tensors.to(device)
-
- if isinstance(num_tensors, list):
- num_tensors = [tensor.to(device) for tensor in num_tensors]
- else:
- num_tensors = num_tensors.to(device)
-
- # Set the model to evaluation mode
- self.model.eval()
-
- # Perform inference
- with torch.no_grad():
- predictions = self.model(
- num_features=num_tensors, cat_features=cat_tensors)
-
- # Convert predictions to NumPy array and return
- return predictions.cpu().numpy()
-
-[docs] def evaluate(self, X, y_true, metrics=None):
- """
- Evaluate the model on the given data using specified metrics.
-
- Parameters
- ----------
- X : array-like or pd.DataFrame of shape (n_samples, n_features)
- The input samples to predict.
- y_true : array-like of shape (n_samples,) or (n_samples, n_outputs)
- The true target values against which to evaluate the predictions.
- metrics : dict
- A dictionary where keys are metric names and values are the metric functions.
-
-
- Notes
- -----
- This method uses the `predict` method to generate predictions and computes each metric.
-
-
- Examples
- --------
- >>> from sklearn.metrics import mean_squared_error, r2_score
- >>> regressor = EmbeddingMambularRegressor()
- >>> regressor.fit(X_train, y_train)
- >>> metrics = {
- ... 'Mean Squared Error': mean_squared_error,
- ... 'R2 Score': r2_score
- ... }
- >>> # Assuming 'X_test' and 'y_test' are your test dataset and labels
- >>> # Evaluate using the specified metrics
- >>> results = regressor.evaluate(X_test, y_test, metrics=metrics)
-
-
- Returns
- -------
- scores : dict
- A dictionary with metric names as keys and their corresponding scores as values.
-
- """
- if metrics is None:
- metrics = {"Mean Squared Error": mean_squared_error}
-
- # Ensure input is in the correct format
- if not isinstance(X, pd.DataFrame):
- X = pd.DataFrame(X)
-
- # Generate predictions using the trained model
- predictions = self.predict(X)
-
- # Initialize dictionary to store results
- scores = {}
-
- # Compute each metric
- for metric_name, metric_func in metrics.items():
- scores[metric_name] = metric_func(y_true, predictions)
-
- return scores
-
-
-
-
-
-import lightning as pl
-import pandas as pd
-import torch
-from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint
-from sklearn.base import BaseEstimator
-from sklearn.metrics import mean_squared_error
-from sklearn.model_selection import train_test_split
-from torch.utils.data import DataLoader
-import warnings
-import numpy as np
-from ..base_models.regressor import BaseMambularRegressor
-from ..utils.dataset import MambularDataModule, MambularDataset
-from ..utils.preprocessor import Preprocessor
-from ..utils.configs import DefaultMambularConfig
-
-
-[docs]class MambularRegressor(BaseEstimator):
- """
- A regressor implemented using PyTorch Lightning that follows the scikit-learn API conventions.
- This class is designed to work with tabular data, offering a straightforward way to specify
- model configurations and preprocessing steps. It integrates seamlessly with scikit-learn's tools
- such as cross-validation and grid search.
-
- Parameters
- ----------
- # configuration parameters
- lr : float, optional
- Learning rate for the optimizer. Default is 1e-4.
- lr_patience : int, optional
- Number of epochs with no improvement on the validation loss to wait before reducing the learning rate. Default is 10.
- weight_decay : float, optional
- Weight decay (L2 penalty) coefficient. Default is 1e-6.
- lr_factor : float, optional
- Factor by which the learning rate will be reduced. Default is 0.1.
- d_model : int, optional
- Dimension of the model. Default is 64.
- n_layers : int, optional
- Number of layers. Default is 8.
- expand_factor : int, optional
- Expansion factor. Default is 2.
- bias : bool, optional
- Whether to use bias. Default is False.
- d_conv : int, optional
- Dimension of the convolution. Default is 16.
- conv_bias : bool, optional
- Whether to use bias in the convolution. Default is True.
- dropout : float, optional
- Dropout rate in the mamba blocks. Default is 0.05.
- dt_rank : str, optional
- Rank of the time dimension. Default is "auto".
- d_state : int, optional
- State dimension. Default is 16.
- dt_scale : float, optional
- Scale of the time dimension. Default is 1.0.
- dt_init : str, optional
- Initialization method for the time dimension. Default is "random".
- dt_max : float, optional
- Maximum value for the time dimension. Default is 0.1.
- dt_min : float, optional
- Minimum value for the time dimension. Default is 1e-3.
- dt_init_floor : float, optional
- Floor value for the time dimension initialization. Default is 1e-4.
- norm : str, optional
- Normalization method. Default is 'RMSNorm'.
- activation : callable, optional
- Activation function. Default is nn.SELU().
- num_embedding_activation : callable, optional
- Activation function for numerical embeddings. Default is nn.Identity().
- head_layer_sizes : list, optional
- Sizes of the layers in the head. Default is [64, 64, 32].
- head_dropout : float, optional
- Dropout rate for the head. Default is 0.5.
- head_skip_layers : bool, optional
- Whether to use skip layers in the head. Default is False.
- head_activation : callable, optional
- Activation function for the head. Default is nn.SELU().
- head_use_batch_norm : bool, optional
- Whether to use batch normalization in the head. Default is False.
-
- # Preprocessor Parameters
- n_bins : int, optional
- The number of bins to use for numerical feature binning. Default is 50.
- numerical_preprocessing : str, optional
- The preprocessing strategy for numerical features. Default is 'ple'.
- use_decision_tree_bins : bool, optional
- If True, uses decision tree regression/classification to determine optimal bin edges for numerical feature binning. Default is False.
- binning_strategy : str, optional
- Defines the strategy for binning numerical features. Default is 'uniform'.
- task : str, optional
- Indicates the type of machine learning task ('regression' or 'classification'). Default is 'regression'.
- cat_cutoff: float or int, optional
- Indicates the cutoff after which integer values are treated as categorical. If float, it's treated as a percentage. If int, it's the maximum number of unique values for a column to be considered categorical. Default is 3%
- treat_all_integers_as_numerical : bool, optional
- If True, all integer columns will be treated as numerical, regardless of their unique value count or proportion. Default is False
-
-
-
- Attributes
- ----------
- config : DefaultConfig
- An object storing the configuration settings for the model.
- preprocessor : Preprocessor
- An object responsible for preprocessing the input data, such as encoding categorical variables and scaling numerical features.
- model : BaseMambularRegressor or None
- The underlying regression model, which is a PyTorch Lightning module. It is instantiated when the `fit` method is called.
- """
-
- def __init__(self, **kwargs):
- # Known config arguments
- config_arg_names = [
- "lr",
- "lr_patience",
- "weight_decay",
- "lr_factor",
- "d_model",
- "n_layers",
- "expand_factor",
- "bias",
- "d_conv",
- "conv_bias",
- "dropout",
- "dt_rank",
- "d_state",
- "dt_scale",
- "dt_init",
- "dt_max",
- "dt_min",
- "dt_init_floor",
- "norm",
- "activation",
- "num_embedding_activation",
- "head_layer_sizes",
- "head_dropout",
- "head_skip_layers",
- "head_activation",
- "head_use_batch_norm",
- ]
-
- preprocessor_arg_names = [
- "n_bins",
- "numerical_preprocessing",
- "use_decision_tree_bins",
- "binning_strategy",
- "task",
- "cat_cutoff",
- "treat_all_integers_as_numerical",
- ]
-
- self.config_kwargs = {k: v for k, v in kwargs.items() if k in config_arg_names}
- self.config = DefaultMambularConfig(**self.config_kwargs)
-
- preprocessor_kwargs = {
- k: v for k, v in kwargs.items() if k in preprocessor_arg_names
- }
-
- self.preprocessor = Preprocessor(**preprocessor_kwargs)
- self.model = None
-
- # Raise a warning if task is set to 'classification'
- if preprocessor_kwargs.get("task") == "classification":
- warnings.warn(
- "The task is set to 'classification'. MambularRegressor is designed for regression tasks.",
- UserWarning,
- )
-
-[docs] def get_params(self, deep=True):
- """
- Get parameters for this estimator. Overrides the BaseEstimator method.
-
- Parameters
- ----------
- deep : bool, default=True
- If True, returns the parameters for this estimator and contained sub-objects that are estimators.
-
- Returns
- -------
- params : dict
- Parameter names mapped to their values.
- """
- params = self.config_kwargs # Parameters used to initialize DefaultConfig
-
- # If deep=True, include parameters from nested components like preprocessor
- if deep:
- # Assuming Preprocessor has a get_params method
- preprocessor_params = {
- "preprocessor__" + key: value
- for key, value in self.preprocessor.get_params().items()
- }
- params.update(preprocessor_params)
-
- return params
-
-[docs] def set_params(self, **parameters):
- """
- Set the parameters of this estimator. Overrides the BaseEstimator method.
-
- Parameters
- ----------
- **parameters : dict
- Estimator parameters to be set.
-
- Returns
- -------
- self : object
- The instance with updated parameters.
- """
- # Update config_kwargs with provided parameters
- valid_config_keys = self.config_kwargs.keys()
- config_updates = {k: v for k, v in parameters.items() if k in valid_config_keys}
- self.config_kwargs.update(config_updates)
-
- # Update the config object
- for key, value in config_updates.items():
- setattr(self.config, key, value)
-
- # Handle preprocessor parameters (prefixed with 'preprocessor__')
- preprocessor_params = {
- k.split("__")[1]: v
- for k, v in parameters.items()
- if k.startswith("preprocessor__")
- }
- if preprocessor_params:
- # Assuming Preprocessor has a set_params method
- self.preprocessor.set_params(**preprocessor_params)
-
- return self
-
-[docs] def split_data(self, X, y, val_size, random_state):
- """
- Splits the dataset into training and validation sets.
-
- Parameters
- ----------
- X : array-like or DataFrame, shape (n_samples, n_features)
- Input features.
- y : array-like, shape (n_samples,) or (n_samples, n_targets)
- Target values.
- val_size : float
- The proportion of the dataset to include in the validation split.
- random_state : int
- Controls the shuffling applied to the data before applying the split.
-
-
- Returns
- -------
- X_train, X_val, y_train, y_val : arrays
- The split datasets.
- """
- X_train, X_val, y_train, y_val = train_test_split(
- X, y, test_size=val_size, random_state=random_state
- )
-
- return X_train, X_val, y_train, y_val
-
-[docs] def preprocess_data(self, X_train, y_train, X_val, y_val, batch_size, shuffle):
- """
- Preprocesses the training and validation data, and creates DataLoaders for them.
-
- Parameters
- ----------
- X_train : DataFrame or array-like, shape (n_samples_train, n_features)
- Training feature set.
- y_train : array-like, shape (n_samples_train,)
- Training target values.
- X_val : DataFrame or array-like, shape (n_samples_val, n_features)
- Validation feature set.
- y_val : array-like, shape (n_samples_val,)
- Validation target values.
- batch_size : int
- Size of batches for the DataLoader.
- shuffle : bool
- Whether to shuffle the training data in the DataLoader.
-
-
- Returns
- -------
- data_module : MambularDataModule
- An instance of MambularDataModule containing the training and validation DataLoaders.
- """
- self.preprocessor.fit(
- pd.concat([X_train, X_val], axis=0).reset_index(drop=True),
- np.concatenate((y_train, y_val), axis=0),
- )
- train_preprocessed_data = self.preprocessor.transform(X_train)
- val_preprocessed_data = self.preprocessor.transform(X_val)
-
- # Update feature info based on the actual processed data
- (
- self.cat_feature_info,
- self.num_feature_info,
- ) = self.preprocessor.get_feature_info()
-
- # Initialize lists for tensors
- train_cat_tensors = []
- train_num_tensors = []
- val_cat_tensors = []
- val_num_tensors = []
-
- # Populate tensors for categorical features, if present in processed data
- for key in self.cat_feature_info:
- cat_key = "cat_" + key # Assuming categorical keys are prefixed with 'cat_'
- if cat_key in train_preprocessed_data:
- train_cat_tensors.append(
- torch.tensor(train_preprocessed_data[cat_key], dtype=torch.long)
- )
- if cat_key in val_preprocessed_data:
- val_cat_tensors.append(
- torch.tensor(val_preprocessed_data[cat_key], dtype=torch.long)
- )
-
- binned_key = "num_" + key # for binned features
- if binned_key in train_preprocessed_data:
- train_cat_tensors.append(
- torch.tensor(train_preprocessed_data[binned_key], dtype=torch.long)
- )
-
- if binned_key in val_preprocessed_data:
- val_cat_tensors.append(
- torch.tensor(val_preprocessed_data[binned_key], dtype=torch.long)
- )
-
- # Populate tensors for numerical features, if present in processed data
- for key in self.num_feature_info:
- num_key = "num_" + key # Assuming numerical keys are prefixed with 'num_'
- if num_key in train_preprocessed_data:
- train_num_tensors.append(
- torch.tensor(train_preprocessed_data[num_key], dtype=torch.float32)
- )
- if num_key in val_preprocessed_data:
- val_num_tensors.append(
- torch.tensor(val_preprocessed_data[num_key], dtype=torch.float32)
- )
-
- train_labels = torch.tensor(y_train, dtype=torch.float32)
- val_labels = torch.tensor(y_val, dtype=torch.float32)
-
- # Create datasets
- train_dataset = MambularDataset(
- train_cat_tensors, train_num_tensors, train_labels
- )
- val_dataset = MambularDataset(val_cat_tensors, val_num_tensors, val_labels)
-
- # Create dataloaders
- train_dataloader = DataLoader(
- train_dataset, batch_size=batch_size, shuffle=shuffle
- )
- val_dataloader = DataLoader(val_dataset, batch_size=batch_size)
-
- return MambularDataModule(train_dataloader, val_dataloader)
-
-[docs] def preprocess_test_data(self, X):
- """
- Preprocesses the test data and creates tensors for categorical and numerical features.
-
- Parameters
- ----------
- X : DataFrame or array-like, shape (n_samples, n_features)
- Test feature set.
-
-
- Returns
- -------
- cat_tensors : list of Tensors
- List of tensors for each categorical feature.
- num_tensors : list of Tensors
- List of tensors for each numerical feature.
- """
- processed_data = self.preprocessor.transform(X)
-
- # Initialize lists for tensors
- cat_tensors = []
- num_tensors = []
-
- # Populate tensors for categorical features
- for key in self.cat_feature_info:
- cat_key = "cat_" + key # Assuming categorical keys are prefixed with 'cat_'
- if cat_key in processed_data:
- cat_tensors.append(
- torch.tensor(processed_data[cat_key], dtype=torch.long)
- )
-
- binned_key = "num_" + key # for binned features
- if binned_key in processed_data:
- cat_tensors.append(
- torch.tensor(processed_data[binned_key], dtype=torch.long)
- )
-
- # Populate tensors for numerical features
- for key in self.num_feature_info:
- num_key = "num_" + key # Assuming numerical keys are prefixed with 'num_'
- if num_key in processed_data:
- num_tensors.append(
- torch.tensor(processed_data[num_key], dtype=torch.float32)
- )
-
- return cat_tensors, num_tensors
-
-[docs] def fit(
- self,
- X,
- y,
- val_size: float = 0.2,
- X_val=None,
- y_val=None,
- max_epochs: int = 100,
- random_state: int = 101,
- batch_size: int = 128,
- shuffle: bool = True,
- patience: int = 15,
- monitor: str = "val_loss",
- mode: str = "min",
- lr: float = 1e-4,
- lr_patience: int = 10,
- factor: float = 0.1,
- weight_decay: float = 1e-06,
- **trainer_kwargs
- ):
- """
- Trains the regression model using the provided training data. Optionally, a separate validation set can be used.
-
- Parameters
- ----------
- X : DataFrame or array-like, shape (n_samples, n_features)
- The training input samples.
- y : array-like, shape (n_samples,) or (n_samples, n_targets)
- The target values (real numbers).
- val_size : float, default=0.2
- The proportion of the dataset to include in the validation split if `X_val` is None. Ignored if `X_val` is provided.
- X_val : DataFrame or array-like, shape (n_samples, n_features), optional
- The validation input samples. If provided, `X` and `y` are not split and this data is used for validation.
- y_val : array-like, shape (n_samples,) or (n_samples, n_targets), optional
- The validation target values. Required if `X_val` is provided.
- max_epochs : int, default=100
- Maximum number of epochs for training.
- random_state : int, default=101
- Controls the shuffling applied to the data before applying the split.
- batch_size : int, default=64
- Number of samples per gradient update.
- shuffle : bool, default=True
- Whether to shuffle the training data before each epoch.
- patience : int, default=10
- Number of epochs with no improvement on the validation loss to wait before early stopping.
- monitor : str, default="val_loss"
- The metric to monitor for early stopping.
- mode : str, default="min"
- Whether the monitored metric should be minimized (`min`) or maximized (`max`).
- lr : float, default=1e-3
- Learning rate for the optimizer.
- lr_patience : int, default=10
- Number of epochs with no improvement on the validation loss to wait before reducing the learning rate.
- factor : float, default=0.1
- Factor by which the learning rate will be reduced.
- weight_decay : float, default=0.025
- Weight decay (L2 penalty) coefficient.
- **trainer_kwargs : Additional keyword arguments for PyTorch Lightning's Trainer class.
-
-
- Returns
- -------
- self : object
- The fitted regressor.
- """
- if not isinstance(X, pd.DataFrame):
- X = pd.DataFrame(X)
- if X_val:
- if not isinstance(X_val, pd.DataFrame):
- X_val = pd.DataFrame(X_val)
-
- if not X_val:
- X_train, X_val, y_train, y_val = self.split_data(
- X, y, val_size, random_state
- )
-
- self.data_module = self.preprocess_data(
- X_train, y_train, X_val, y_val, batch_size, shuffle
- )
-
- self.model = BaseMambularRegressor(
- config=self.config,
- cat_feature_info=self.cat_feature_info,
- num_feature_info=self.num_feature_info,
- lr=lr,
- lr_patience=lr_patience,
- lr_factor=factor,
- weight_decay=weight_decay,
- )
-
- early_stop_callback = EarlyStopping(
- monitor=monitor, min_delta=0.00, patience=patience, verbose=False, mode=mode
- )
-
- checkpoint_callback = ModelCheckpoint(
- monitor="val_loss", # Adjust according to your validation metric
- mode="min",
- save_top_k=1,
- dirpath="model_checkpoints", # Specify the directory to save checkpoints
- filename="best_model",
- )
-
- # Initialize the trainer and train the model
- trainer = pl.Trainer(
- max_epochs=max_epochs,
- callbacks=[early_stop_callback, checkpoint_callback],
- **trainer_kwargs
- )
- trainer.fit(self.model, self.data_module)
-
- best_model_path = checkpoint_callback.best_model_path
- if best_model_path:
- checkpoint = torch.load(best_model_path)
- self.model.load_state_dict(checkpoint["state_dict"])
-
- return self
-
-[docs] def predict(self, X):
- """
- Predicts target values for the given input samples.
-
- Parameters
- ----------
- X : DataFrame or array-like, shape (n_samples, n_features)
- The input samples for which to predict target values.
-
-
- Returns
- -------
- predictions : ndarray, shape (n_samples,) or (n_samples, n_outputs)
- The predicted target values.
- """
- # Preprocess the data
- if not isinstance(X, pd.DataFrame):
- X = pd.DataFrame(X)
- device = next(self.model.parameters()).device
- cat_tensors, num_tensors = self.preprocess_test_data(X)
- if isinstance(cat_tensors, list):
- cat_tensors = [tensor.to(device) for tensor in cat_tensors]
- else:
- cat_tensors = cat_tensors.to(device)
-
- if isinstance(num_tensors, list):
- num_tensors = [tensor.to(device) for tensor in num_tensors]
- else:
- num_tensors = num_tensors.to(device)
-
- # Set the model to evaluation mode
- self.model.eval()
-
- # Perform inference
- with torch.no_grad():
- predictions = self.model(num_features=num_tensors, cat_features=cat_tensors)
-
- # Convert predictions to NumPy array and return
- return predictions.cpu().numpy()
-
-[docs] def evaluate(self, X, y_true, metrics=None):
- """
- Evaluate the model on the given data using specified metrics.
-
- Parameters
- ----------
- X : array-like or pd.DataFrame of shape (n_samples, n_features)
- The input samples to predict.
- y_true : array-like of shape (n_samples,) or (n_samples, n_outputs)
- The true target values against which to evaluate the predictions.
- metrics : dict
- A dictionary where keys are metric names and values are the metric functions.
-
-
- Notes
- -----
- This method uses the `predict` method to generate predictions and computes each metric.
-
-
- Examples
- --------
- >>> from sklearn.metrics import mean_squared_error, r2_score
- >>> from sklearn.model_selection import train_test_split
- >>> from mambular.models import MambularRegressor
- >>> metrics = {
- ... 'Mean Squared Error': mean_squared_error,
- ... 'R2 Score': r2_score
- ... }
- >>> # Assuming 'X_test' and 'y_test' are your test dataset and labels
- >>> # Evaluate using the specified metrics
- >>> results = regressor.evaluate(X_test, y_test, metrics=metrics)
-
-
- Returns
- -------
- scores : dict
- A dictionary with metric names as keys and their corresponding scores as values.
- """
- if metrics is None:
- metrics = {"Mean Squared Error": mean_squared_error}
-
- # Ensure input is in the correct format
- if not isinstance(X, pd.DataFrame):
- X = pd.DataFrame(X)
-
- # Generate predictions using the trained model
- predictions = self.predict(X)
-
- # Initialize dictionary to store results
- scores = {}
-
- # Compute each metric
- for metric_name, metric_func in metrics.items():
- scores[metric_name] = metric_func(y_true, predictions)
-
- return scores
-
-
-
-
-
-import numpy as np
-import pandas as pd
-from sklearn.compose import ColumnTransformer
-from sklearn.exceptions import NotFittedError
-from sklearn.impute import SimpleImputer
-from sklearn.pipeline import Pipeline
-from sklearn.preprocessing import KBinsDiscretizer, MinMaxScaler, StandardScaler
-from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
-
-from .ple_encoding import PLE
-from .prepro_utils import ContinuousOrdinalEncoder, CustomBinner, OneHotFromOrdinal
-
-
-[docs]class Preprocessor:
- """
- A comprehensive preprocessor for structured data, capable of handling both numerical and categorical features.
- It supports various preprocessing strategies for numerical data, including binning, one-hot encoding,
- standardization, and normalization. Categorical features can be transformed using continuous ordinal encoding.
- Additionally, it allows for the use of decision tree-derived bin edges for numerical feature binning.
-
- The class is designed to work seamlessly with pandas DataFrames, facilitating easy integration into
- machine learning pipelines.
-
- Parameters
- ----------
- n_bins (int): The number of bins to use for numerical feature binning. This parameter is relevant
- only if `numerical_preprocessing` is set to 'binning' or 'one_hot'.
- numerical_preprocessing (str): The preprocessing strategy for numerical features. Valid options are
- 'binning', 'one_hot', 'standardization', and 'normalization'.
- use_decision_tree_bins (bool): If True, uses decision tree regression/classification to determine
- optimal bin edges for numerical feature binning. This parameter is
- relevant only if `numerical_preprocessing` is set to 'binning' or 'one_hot'.
- binning_strategy (str): Defines the strategy for binning numerical features. Options include 'uniform',
- 'quantile', or other sklearn-compatible strategies.
- task (str): Indicates the type of machine learning task ('regression' or 'classification'). This can
- influence certain preprocessing behaviors, especially when using decision tree-based binning.
- binning_strategy (str): Defines the strategy for binning numerical features. Options include 'uniform',
- 'quantile', or other sklearn-compatible strategies.
- task (str): Indicates the type of machine learning task ('regression' or 'classification'). This can
- influence certain preprocessing behaviors, especially when using decision tree-based binning.
- cat_cutoff (float or int): Indicates the cutoff after which integer values are treated as categorical.
- If float, it's treated as a percentage. If int, it's the maximum number of
- unique values for a column to be considered categorical.
- treat_all_integers_as_numerical (bool): If True, all integer columns will be treated as numerical, regardless
- of their unique value count or proportion.
-
-
- Attributes
- ----------
- column_transformer (ColumnTransformer): An instance of sklearn's ColumnTransformer that holds the
- configured preprocessing pipelines for different feature types.
- fitted (bool): Indicates whether the preprocessor has been fitted to the data.
-
- """
-
- def __init__(
- self,
- n_bins=50,
- numerical_preprocessing="ple",
- use_decision_tree_bins=False,
- binning_strategy="uniform",
- task="regression",
- cat_cutoff=0.03,
- treat_all_integers_as_numerical=False,
- ):
- self.n_bins = n_bins
- self.numerical_preprocessing = numerical_preprocessing.lower()
- if self.numerical_preprocessing not in [
- "ple",
- "binning",
- "one_hot",
- "standardization",
- "normalization",
- ]:
- raise ValueError(
- "Invalid numerical_preprocessing value. Supported values are 'ple', 'binning', 'one_hot', 'standardization', and 'normalization'."
- )
- self.numerical_preprocessing = numerical_preprocessing.lower()
- if self.numerical_preprocessing not in [
- "ple",
- "binning",
- "one_hot",
- "standardization",
- "normalization",
- ]:
- raise ValueError(
- "Invalid numerical_preprocessing value. Supported values are 'ple', 'binning', 'one_hot', 'standardization', and 'normalization'."
- )
- self.use_decision_tree_bins = use_decision_tree_bins
- self.column_transformer = None
- self.fitted = False
- self.binning_strategy = binning_strategy
- self.task = task
- self.cat_cutoff = cat_cutoff
- self.treat_all_integers_as_numerical = treat_all_integers_as_numerical
-
- def set_params(self, **params):
- for key, value in params.items():
- setattr(self, key, value)
- return self
-
- def _detect_column_types(self, X):
- """
- Identifies and separates the features in the dataset into numerical and categorical types based on the data type
- and the proportion of unique values.
-
- Parameters
- ----------
- X (DataFrame or dict): The input dataset, where the features are columns in a DataFrame or keys in a dict.
-
-
- Returns
- -------
- tuple: A tuple containing two lists, the first with the names of numerical features and the second with the names of categorical features.
- """
- categorical_features = []
- numerical_features = []
-
- if isinstance(X, dict):
- X = pd.DataFrame(X)
-
- for col in X.columns:
- num_unique_values = X[col].nunique()
- total_samples = len(X[col])
-
- if self.treat_all_integers_as_numerical and X[col].dtype.kind == "i":
- numerical_features.append(col)
- else:
- if isinstance(self.cat_cutoff, float):
- cutoff_condition = (
- num_unique_values / total_samples
- ) < self.cat_cutoff
- elif isinstance(self.cat_cutoff, int):
- cutoff_condition = num_unique_values < self.cat_cutoff
- else:
- raise ValueError(
- "cat_cutoff should be either a float or an integer."
- )
-
- if X[col].dtype.kind not in "iufc" or (
- X[col].dtype.kind == "i" and cutoff_condition
- ):
- categorical_features.append(col)
- else:
- numerical_features.append(col)
-
- return numerical_features, categorical_features
-
-[docs] def fit(self, X, y=None):
- """
- Fits the preprocessor to the data by identifying feature types and configuring the appropriate transformations for each feature.
- It sets up a column transformer with a pipeline of transformations for numerical and categorical features based on the specified preprocessing strategy.
-
- Parameters
- ----------
- X (DataFrame or dict): The input dataset to fit the preprocessor on.
- y (array-like, optional): The target variable. Required if `use_decision_tree_bins` is True for determining optimal bin edges using decision trees.
-
-
- Returns
- -------
- self: The fitted Preprocessor instance.
- """
- if isinstance(X, dict):
- X = pd.DataFrame(X)
-
- numerical_features, categorical_features = self._detect_column_types(X)
- transformers = []
-
- if numerical_features:
- for feature in numerical_features:
- numeric_transformer_steps = [
- ("imputer", SimpleImputer(strategy="mean"))
- ]
-
- if self.numerical_preprocessing in ["binning", "one_hot"]:
- bins = (
- self._get_decision_tree_bins(X[[feature]], y, [feature])
- if self.use_decision_tree_bins
- else self.n_bins
- )
- if isinstance(bins, int):
- numeric_transformer_steps.extend(
- [
- (
- "discretizer",
- KBinsDiscretizer(
- n_bins=bins
- if isinstance(bins, int)
- else len(bins) - 1,
- encode="ordinal",
- strategy=self.binning_strategy,
- subsample=200_000 if len(X) > 200_000 else None,
- ),
- ),
- ]
- )
- else:
- numeric_transformer_steps.extend(
- [
- (
- "discretizer",
- CustomBinner(bins=bins),
- ),
- ]
- )
-
- if self.numerical_preprocessing == "one_hot":
- numeric_transformer_steps.extend(
- [
- ("onehot_from_ordinal", OneHotFromOrdinal()),
- ]
- )
-
- elif self.numerical_preprocessing == "standardization":
- numeric_transformer_steps.append(("scaler", StandardScaler()))
-
- elif self.numerical_preprocessing == "normalization":
- numeric_transformer_steps.append(("normalizer", MinMaxScaler()))
-
- elif self.numerical_preprocessing == "ple":
- numeric_transformer_steps.append(("normalizer", MinMaxScaler()))
- numeric_transformer_steps.append(
- ("ple", PLE(n_bins=self.n_bins, task=self.task))
- )
-
- elif self.numerical_preprocessing == "ple":
- numeric_transformer_steps.append(("normalizer", MinMaxScaler()))
- numeric_transformer_steps.append(
- ("ple", PLE(n_bins=self.n_bins, task=self.task))
- )
-
- numeric_transformer = Pipeline(numeric_transformer_steps)
-
- transformers.append((f"num_{feature}", numeric_transformer, [feature]))
-
- if categorical_features:
- for feature in categorical_features:
- # Create a pipeline for each categorical feature
- categorical_transformer = Pipeline(
- [
- ("imputer", SimpleImputer(strategy="most_frequent")),
- (
- "continuous_ordinal",
- ContinuousOrdinalEncoder(),
- ),
- ]
- )
- # Append the transformer for the current categorical feature
- transformers.append(
- (f"cat_{feature}", categorical_transformer, [feature])
- )
-
- self.column_transformer = ColumnTransformer(
- transformers=transformers, remainder="passthrough"
- )
- self.column_transformer.fit(X, y)
-
- self.fitted = True
-
- def _get_decision_tree_bins(self, X, y, numerical_features):
- """
- Uses decision tree models to determine optimal bin edges for numerical feature binning. This method is used when `use_decision_tree_bins` is True.
-
- Parameters
- ----------
- X (DataFrame): The input dataset containing only the numerical features for which the bin edges are to be determined.
- y (array-like): The target variable for training the decision tree models.
- numerical_features (list of str): The names of the numerical features for which the bin edges are to be determined.
-
-
- Returns
- -------
- list: A list of arrays, where each array contains the bin edges determined by the decision tree for a numerical feature.
- """
- bins = []
- for feature in numerical_features:
- tree_model = (
- DecisionTreeClassifier(max_depth=3)
- if y.dtype.kind in "bi"
- else DecisionTreeRegressor(max_depth=3)
- )
- tree_model.fit(X[[feature]], y)
- thresholds = tree_model.tree_.threshold[tree_model.tree_.feature != -2]
- bin_edges = np.sort(np.unique(thresholds))
-
- bins.append(
- np.concatenate(([X[feature].min()], bin_edges, [X[feature].max()]))
- )
- return bins
-
-[docs] def transform(self, X):
- """
- Transforms the input data using the preconfigured column transformer and converts the output into a dictionary
- format with keys corresponding to transformed feature names and values as arrays of transformed data.
-
- This method converts the sparse or dense matrix returned by the column transformer into a more accessible
- dictionary format, where each key-value pair represents a feature and its transformed data.
- Transforms the input data using the preconfigured column transformer and converts the output into a dictionary
- format with keys corresponding to transformed feature names and values as arrays of transformed data.
-
- This method converts the sparse or dense matrix returned by the column transformer into a more accessible
- dictionary format, where each key-value pair represents a feature and its transformed data.
-
- Parameters
- ----------
- X (DataFrame): The input data to be transformed.
- X (DataFrame): The input data to be transformed.
-
-
- Returns
- -------
- dict: A dictionary where keys are the names of the features (as per the transformations defined in the
- column transformer) and the values are numpy arrays of the transformed data.
- """
- if not self.fitted:
- raise NotFittedError(
- "The preprocessor must be fitted before transforming new data. Use .fit or .fit_transform"
- )
- transformed_X = self.column_transformer.transform(X)
-
- # Now let's convert this into a dictionary of arrays, one per column
- transformed_dict = self._split_transformed_output(X, transformed_X)
- return transformed_dict
-
- def _split_transformed_output(self, X, transformed_X):
- """
- Splits the transformed data array into a dictionary where keys correspond to the original column names or
- feature groups and values are the transformed data for those columns.
-
- This helper method is utilized within `transform` to segregate the transformed data based on the
- specification in the column transformer, assigning each transformed section to its corresponding feature name.
-
- Parameters
- ----------
- X (DataFrame): The original input data, used for determining shapes and transformations.
- transformed_X (numpy array): The transformed data as a numpy array, outputted by the column transformer.
-
-
- Returns
- -------
- dict: A dictionary mapping each transformation's name to its respective numpy array of transformed data.
- The type of each array (int or float) is determined based on the type of transformation applied.
- """
- start = 0
- transformed_dict = {}
- for (
- name,
- transformer,
- columns,
- ) in self.column_transformer.transformers_:
- if transformer != "drop":
- end = start + transformer.transform(X[[columns[0]]]).shape[1]
- dtype = int if "cat" in name else float
- transformed_dict[name] = transformed_X[:, start:end].astype(dtype)
- start = end
-
- return transformed_dict
-
-[docs] def fit_transform(self, X, y=None):
- """
- Fits the preprocessor to the data and then transforms the data using the fitted preprocessing pipelines. This is a convenience method that combines `fit` and `transform`.
-
- Parameters
- ----------
- X (DataFrame or dict): The input dataset to fit the preprocessor on and then transform.
- y (array-like, optional): The target variable. Required if `use_decision_tree_bins` is True.
-
-
- Returns
- -------
- dict: A dictionary with the transformed data, where keys are the base feature names and values are the transformed features as arrays.
- """
- self.fit(X, y)
- self.fitted = True
- return self.transform(X)
-
-[docs] def get_feature_info(self):
- """
- Retrieves information about how features are encoded within the model's preprocessor.
- This method identifies the type of encoding applied to each feature, categorizing them into binned or ordinal
- encodings and other types of encodings (e.g., one-hot encoding after discretization).
-
- This method should only be called after the preprocessor has been fitted, as it relies on the structure and
- configuration of the `column_transformer` attribute.
-
-
- Raises
- ------
- RuntimeError: If the `column_transformer` is not yet fitted, indicating that the preprocessor must be
- fitted before invoking this method.
-
-
- Returns
- -------
- tuple of (dict, dict):
- - The first dictionary maps feature names to their respective number of bins or categories if they are
- processed using discretization or ordinal encoding.
- - The second dictionary includes feature names with other encoding details, such as the dimension of
- features after encoding transformations (e.g., one-hot encoding dimensions).
-
- """
- binned_or_ordinal_info = {}
- other_encoding_info = {}
-
- if not self.column_transformer:
- raise RuntimeError("The preprocessor has not been fitted yet.")
-
- for (
- name,
- transformer_pipeline,
- columns,
- ) in self.column_transformer.transformers_:
- steps = [step[0] for step in transformer_pipeline.steps]
-
- for feature_name in columns:
- # Handle features processed with discretization
- if "discretizer" in steps:
- step = transformer_pipeline.named_steps["discretizer"]
- n_bins = step.n_bins_[0] if hasattr(step, "n_bins_") else None
-
- # Check if discretization is followed by one-hot encoding
- if "onehot_from_ordinal" in steps:
- # Classify as other encoding due to the expanded feature dimensions from one-hot encoding
- other_encoding_info[
- feature_name
- ] = n_bins # Number of bins before one-hot encoding
- print(
- f"Numerical Feature (Discretized & One-Hot Encoded): {feature_name}, Number of bins before one-hot encoding: {n_bins}"
- )
- else:
- # Only discretization without subsequent one-hot encoding
- binned_or_ordinal_info[feature_name] = n_bins
- print(
- f"Numerical Feature (Binned): {feature_name}, Number of bins: {n_bins}"
- )
-
- # Handle features processed with continuous ordinal encoding
- elif "continuous_ordinal" in steps:
- step = transformer_pipeline.named_steps["continuous_ordinal"]
- n_categories = len(step.mapping_[columns.index(feature_name)])
- binned_or_ordinal_info[feature_name] = n_categories
- print(
- f"Categorical Feature (Ordinal Encoded): {feature_name}, Number of unique categories: {n_categories}"
- )
-
- # Handle other numerical feature encodings
- else:
- last_step = transformer_pipeline.steps[-1][1]
- if hasattr(last_step, "transform"):
- transformed_feature = last_step.transform(
- np.zeros((1, len(columns)))
- )
- other_encoding_info[feature_name] = transformed_feature.shape[1]
- print(
- f"Feature: {feature_name} (Other Encoding), Encoded feature dimension: {transformed_feature.shape[1]}"
- f"Feature: {feature_name} (Other Encoding), Encoded feature dimension: {transformed_feature.shape[1]}"
- )
-
- print("-" * 50)
-
- return binned_or_ordinal_info, other_encoding_info
-Short
- */ - .o-tooltip--left { - position: relative; - } - - .o-tooltip--left:after { - opacity: 0; - visibility: hidden; - position: absolute; - content: attr(data-tooltip); - padding: .2em; - font-size: .8em; - left: -.2em; - background: grey; - color: white; - white-space: nowrap; - z-index: 2; - border-radius: 2px; - transform: translateX(-102%) translateY(0); - transition: opacity 0.2s cubic-bezier(0.64, 0.09, 0.08, 1), transform 0.2s cubic-bezier(0.64, 0.09, 0.08, 1); -} - -.o-tooltip--left:hover:after { - display: block; - opacity: 1; - visibility: visible; - transform: translateX(-100%) translateY(0); - transition: opacity 0.2s cubic-bezier(0.64, 0.09, 0.08, 1), transform 0.2s cubic-bezier(0.64, 0.09, 0.08, 1); - transition-delay: .5s; -} - -/* By default the copy button shouldn't show up when printing a page */ -@media print { - button.copybtn { - display: none; - } -} diff --git a/docs/_build/html/_static/copybutton.js b/docs/_build/html/_static/copybutton.js deleted file mode 100644 index 2ea7ff3..0000000 --- a/docs/_build/html/_static/copybutton.js +++ /dev/null @@ -1,248 +0,0 @@ -// Localization support -const messages = { - 'en': { - 'copy': 'Copy', - 'copy_to_clipboard': 'Copy to clipboard', - 'copy_success': 'Copied!', - 'copy_failure': 'Failed to copy', - }, - 'es' : { - 'copy': 'Copiar', - 'copy_to_clipboard': 'Copiar al portapapeles', - 'copy_success': '¡Copiado!', - 'copy_failure': 'Error al copiar', - }, - 'de' : { - 'copy': 'Kopieren', - 'copy_to_clipboard': 'In die Zwischenablage kopieren', - 'copy_success': 'Kopiert!', - 'copy_failure': 'Fehler beim Kopieren', - }, - 'fr' : { - 'copy': 'Copier', - 'copy_to_clipboard': 'Copier dans le presse-papier', - 'copy_success': 'Copié !', - 'copy_failure': 'Échec de la copie', - }, - 'ru': { - 'copy': 'Скопировать', - 'copy_to_clipboard': 'Скопировать в буфер', - 'copy_success': 'Скопировано!', - 'copy_failure': 'Не удалось скопировать', - }, - 'zh-CN': { - 'copy': '复制', - 'copy_to_clipboard': '复制到剪贴板', - 'copy_success': '复制成功!', - 'copy_failure': '复制失败', - }, - 'it' : { - 'copy': 'Copiare', - 'copy_to_clipboard': 'Copiato negli appunti', - 'copy_success': 'Copiato!', - 'copy_failure': 'Errore durante la copia', - } -} - -let locale = 'en' -if( document.documentElement.lang !== undefined - && messages[document.documentElement.lang] !== undefined ) { - locale = document.documentElement.lang -} - -let doc_url_root = DOCUMENTATION_OPTIONS.URL_ROOT; -if (doc_url_root == '#') { - doc_url_root = ''; -} - -/** - * SVG files for our copy buttons - */ -let iconCheck = `` - -// If the user specified their own SVG use that, otherwise use the default -let iconCopy = ``; -if (!iconCopy) { - iconCopy = `` -} - -/** - * Set up copy/paste for code blocks - */ - -const runWhenDOMLoaded = cb => { - if (document.readyState != 'loading') { - cb() - } else if (document.addEventListener) { - document.addEventListener('DOMContentLoaded', cb) - } else { - document.attachEvent('onreadystatechange', function() { - if (document.readyState == 'complete') cb() - }) - } -} - -const codeCellId = index => `codecell${index}` - -// Clears selected text since ClipboardJS will select the text when copying -const clearSelection = () => { - if (window.getSelection) { - window.getSelection().removeAllRanges() - } else if (document.selection) { - document.selection.empty() - } -} - -// Changes tooltip text for a moment, then changes it back -// We want the timeout of our `success` class to be a bit shorter than the -// tooltip and icon change, so that we can hide the icon before changing back. -var timeoutIcon = 2000; -var timeoutSuccessClass = 1500; - -const temporarilyChangeTooltip = (el, oldText, newText) => { - el.setAttribute('data-tooltip', newText) - el.classList.add('success') - // Remove success a little bit sooner than we change the tooltip - // So that we can use CSS to hide the copybutton first - setTimeout(() => el.classList.remove('success'), timeoutSuccessClass) - setTimeout(() => el.setAttribute('data-tooltip', oldText), timeoutIcon) -} - -// Changes the copy button icon for two seconds, then changes it back -const temporarilyChangeIcon = (el) => { - el.innerHTML = iconCheck; - setTimeout(() => {el.innerHTML = iconCopy}, timeoutIcon) -} - -const addCopyButtonToCodeCells = () => { - // If ClipboardJS hasn't loaded, wait a bit and try again. This - // happens because we load ClipboardJS asynchronously. - if (window.ClipboardJS === undefined) { - setTimeout(addCopyButtonToCodeCells, 250) - return - } - - // Add copybuttons to all of our code cells - const COPYBUTTON_SELECTOR = 'div.highlight pre'; - const codeCells = document.querySelectorAll(COPYBUTTON_SELECTOR) - codeCells.forEach((codeCell, index) => { - const id = codeCellId(index) - codeCell.setAttribute('id', id) - - const clipboardButton = id => - `` - codeCell.insertAdjacentHTML('afterend', clipboardButton(id)) - }) - -function escapeRegExp(string) { - return string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); // $& means the whole matched string -} - -/** - * Removes excluded text from a Node. - * - * @param {Node} target Node to filter. - * @param {string} exclude CSS selector of nodes to exclude. - * @returns {DOMString} Text from `target` with text removed. - */ -function filterText(target, exclude) { - const clone = target.cloneNode(true); // clone as to not modify the live DOM - if (exclude) { - // remove excluded nodes - clone.querySelectorAll(exclude).forEach(node => node.remove()); - } - return clone.innerText; -} - -// Callback when a copy button is clicked. Will be passed the node that was clicked -// should then grab the text and replace pieces of text that shouldn't be used in output -function formatCopyText(textContent, copybuttonPromptText, isRegexp = false, onlyCopyPromptLines = true, removePrompts = true, copyEmptyLines = true, lineContinuationChar = "", hereDocDelim = "") { - var regexp; - var match; - - // Do we check for line continuation characters and "HERE-documents"? - var useLineCont = !!lineContinuationChar - var useHereDoc = !!hereDocDelim - - // create regexp to capture prompt and remaining line - if (isRegexp) { - regexp = new RegExp('^(' + copybuttonPromptText + ')(.*)') - } else { - regexp = new RegExp('^(' + escapeRegExp(copybuttonPromptText) + ')(.*)') - } - - const outputLines = []; - var promptFound = false; - var gotLineCont = false; - var gotHereDoc = false; - const lineGotPrompt = []; - for (const line of textContent.split('\n')) { - match = line.match(regexp) - if (match || gotLineCont || gotHereDoc) { - promptFound = regexp.test(line) - lineGotPrompt.push(promptFound) - if (removePrompts && promptFound) { - outputLines.push(match[2]) - } else { - outputLines.push(line) - } - gotLineCont = line.endsWith(lineContinuationChar) & useLineCont - if (line.includes(hereDocDelim) & useHereDoc) - gotHereDoc = !gotHereDoc - } else if (!onlyCopyPromptLines) { - outputLines.push(line) - } else if (copyEmptyLines && line.trim() === '') { - outputLines.push(line) - } - } - - // If no lines with the prompt were found then just use original lines - if (lineGotPrompt.some(v => v === true)) { - textContent = outputLines.join('\n'); - } - - // Remove a trailing newline to avoid auto-running when pasting - if (textContent.endsWith("\n")) { - textContent = textContent.slice(0, -1) - } - return textContent -} - - -var copyTargetText = (trigger) => { - var target = document.querySelector(trigger.attributes['data-clipboard-target'].value); - - // get filtered text - let exclude = '.linenos'; - - let text = filterText(target, exclude); - return formatCopyText(text, '', false, true, true, true, '', '') -} - - // Initialize with a callback so we can modify the text before copy - const clipboard = new ClipboardJS('.copybtn', {text: copyTargetText}) - - // Update UI with error/success messages - clipboard.on('success', event => { - clearSelection() - temporarilyChangeTooltip(event.trigger, messages[locale]['copy'], messages[locale]['copy_success']) - temporarilyChangeIcon(event.trigger) - }) - - clipboard.on('error', event => { - temporarilyChangeTooltip(event.trigger, messages[locale]['copy'], messages[locale]['copy_failure']) - }) -} - -runWhenDOMLoaded(addCopyButtonToCodeCells) \ No newline at end of file diff --git a/docs/_build/html/_static/copybutton_funcs.js b/docs/_build/html/_static/copybutton_funcs.js deleted file mode 100644 index dbe1aaa..0000000 --- a/docs/_build/html/_static/copybutton_funcs.js +++ /dev/null @@ -1,73 +0,0 @@ -function escapeRegExp(string) { - return string.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); // $& means the whole matched string -} - -/** - * Removes excluded text from a Node. - * - * @param {Node} target Node to filter. - * @param {string} exclude CSS selector of nodes to exclude. - * @returns {DOMString} Text from `target` with text removed. - */ -export function filterText(target, exclude) { - const clone = target.cloneNode(true); // clone as to not modify the live DOM - if (exclude) { - // remove excluded nodes - clone.querySelectorAll(exclude).forEach(node => node.remove()); - } - return clone.innerText; -} - -// Callback when a copy button is clicked. Will be passed the node that was clicked -// should then grab the text and replace pieces of text that shouldn't be used in output -export function formatCopyText(textContent, copybuttonPromptText, isRegexp = false, onlyCopyPromptLines = true, removePrompts = true, copyEmptyLines = true, lineContinuationChar = "", hereDocDelim = "") { - var regexp; - var match; - - // Do we check for line continuation characters and "HERE-documents"? - var useLineCont = !!lineContinuationChar - var useHereDoc = !!hereDocDelim - - // create regexp to capture prompt and remaining line - if (isRegexp) { - regexp = new RegExp('^(' + copybuttonPromptText + ')(.*)') - } else { - regexp = new RegExp('^(' + escapeRegExp(copybuttonPromptText) + ')(.*)') - } - - const outputLines = []; - var promptFound = false; - var gotLineCont = false; - var gotHereDoc = false; - const lineGotPrompt = []; - for (const line of textContent.split('\n')) { - match = line.match(regexp) - if (match || gotLineCont || gotHereDoc) { - promptFound = regexp.test(line) - lineGotPrompt.push(promptFound) - if (removePrompts && promptFound) { - outputLines.push(match[2]) - } else { - outputLines.push(line) - } - gotLineCont = line.endsWith(lineContinuationChar) & useLineCont - if (line.includes(hereDocDelim) & useHereDoc) - gotHereDoc = !gotHereDoc - } else if (!onlyCopyPromptLines) { - outputLines.push(line) - } else if (copyEmptyLines && line.trim() === '') { - outputLines.push(line) - } - } - - // If no lines with the prompt were found then just use original lines - if (lineGotPrompt.some(v => v === true)) { - textContent = outputLines.join('\n'); - } - - // Remove a trailing newline to avoid auto-running when pasting - if (textContent.endsWith("\n")) { - textContent = textContent.slice(0, -1) - } - return textContent -} diff --git a/docs/_build/html/_static/doctools.js b/docs/_build/html/_static/doctools.js deleted file mode 100644 index d06a71d..0000000 --- a/docs/_build/html/_static/doctools.js +++ /dev/null @@ -1,156 +0,0 @@ -/* - * doctools.js - * ~~~~~~~~~~~ - * - * Base JavaScript utilities for all Sphinx HTML documentation. - * - * :copyright: Copyright 2007-2023 by the Sphinx team, see AUTHORS. - * :license: BSD, see LICENSE for details. - * - */ -"use strict"; - -const BLACKLISTED_KEY_CONTROL_ELEMENTS = new Set([ - "TEXTAREA", - "INPUT", - "SELECT", - "BUTTON", -]); - -const _ready = (callback) => { - if (document.readyState !== "loading") { - callback(); - } else { - document.addEventListener("DOMContentLoaded", callback); - } -}; - -/** - * Small JavaScript module for the documentation. - */ -const Documentation = { - init: () => { - Documentation.initDomainIndexTable(); - Documentation.initOnKeyListeners(); - }, - - /** - * i18n support - */ - TRANSLATIONS: {}, - PLURAL_EXPR: (n) => (n === 1 ? 0 : 1), - LOCALE: "unknown", - - // gettext and ngettext don't access this so that the functions - // can safely bound to a different name (_ = Documentation.gettext) - gettext: (string) => { - const translated = Documentation.TRANSLATIONS[string]; - switch (typeof translated) { - case "undefined": - return string; // no translation - case "string": - return translated; // translation exists - default: - return translated[0]; // (singular, plural) translation tuple exists - } - }, - - ngettext: (singular, plural, n) => { - const translated = Documentation.TRANSLATIONS[singular]; - if (typeof translated !== "undefined") - return translated[Documentation.PLURAL_EXPR(n)]; - return n === 1 ? singular : plural; - }, - - addTranslations: (catalog) => { - Object.assign(Documentation.TRANSLATIONS, catalog.messages); - Documentation.PLURAL_EXPR = new Function( - "n", - `return (${catalog.plural_expr})` - ); - Documentation.LOCALE = catalog.locale; - }, - - /** - * helper function to focus on search bar - */ - focusSearchBar: () => { - document.querySelectorAll("input[name=q]")[0]?.focus(); - }, - - /** - * Initialise the domain index toggle buttons - */ - initDomainIndexTable: () => { - const toggler = (el) => { - const idNumber = el.id.substr(7); - const toggledRows = document.querySelectorAll(`tr.cg-${idNumber}`); - if (el.src.substr(-9) === "minus.png") { - el.src = `${el.src.substr(0, el.src.length - 9)}plus.png`; - toggledRows.forEach((el) => (el.style.display = "none")); - } else { - el.src = `${el.src.substr(0, el.src.length - 8)}minus.png`; - toggledRows.forEach((el) => (el.style.display = "")); - } - }; - - const togglerElements = document.querySelectorAll("img.toggler"); - togglerElements.forEach((el) => - el.addEventListener("click", (event) => toggler(event.currentTarget)) - ); - togglerElements.forEach((el) => (el.style.display = "")); - if (DOCUMENTATION_OPTIONS.COLLAPSE_INDEX) togglerElements.forEach(toggler); - }, - - initOnKeyListeners: () => { - // only install a listener if it is really needed - if ( - !DOCUMENTATION_OPTIONS.NAVIGATION_WITH_KEYS && - !DOCUMENTATION_OPTIONS.ENABLE_SEARCH_SHORTCUTS - ) - return; - - document.addEventListener("keydown", (event) => { - // bail for input elements - if (BLACKLISTED_KEY_CONTROL_ELEMENTS.has(document.activeElement.tagName)) return; - // bail with special keys - if (event.altKey || event.ctrlKey || event.metaKey) return; - - if (!event.shiftKey) { - switch (event.key) { - case "ArrowLeft": - if (!DOCUMENTATION_OPTIONS.NAVIGATION_WITH_KEYS) break; - - const prevLink = document.querySelector('link[rel="prev"]'); - if (prevLink && prevLink.href) { - window.location.href = prevLink.href; - event.preventDefault(); - } - break; - case "ArrowRight": - if (!DOCUMENTATION_OPTIONS.NAVIGATION_WITH_KEYS) break; - - const nextLink = document.querySelector('link[rel="next"]'); - if (nextLink && nextLink.href) { - window.location.href = nextLink.href; - event.preventDefault(); - } - break; - } - } - - // some keyboard layouts may need Shift to get / - switch (event.key) { - case "/": - if (!DOCUMENTATION_OPTIONS.ENABLE_SEARCH_SHORTCUTS) break; - Documentation.focusSearchBar(); - event.preventDefault(); - } - }); - }, -}; - -// quick alias for translations -const _ = Documentation.gettext; - -_ready(Documentation.init); diff --git a/docs/_build/html/_static/documentation_options.js b/docs/_build/html/_static/documentation_options.js deleted file mode 100644 index 133c7f5..0000000 --- a/docs/_build/html/_static/documentation_options.js +++ /dev/null @@ -1,14 +0,0 @@ -var DOCUMENTATION_OPTIONS = { - URL_ROOT: document.getElementById("documentation_options").getAttribute('data-url_root'), - VERSION: '0.1.4', - LANGUAGE: 'en', - COLLAPSE_INDEX: false, - BUILDER: 'html', - FILE_SUFFIX: '.html', - LINK_SUFFIX: '.html', - HAS_SOURCE: true, - SOURCELINK_SUFFIX: '', - NAVIGATION_WITH_KEYS: false, - SHOW_SEARCH_SUMMARY: true, - ENABLE_SEARCH_SHORTCUTS: true, -}; \ No newline at end of file diff --git a/docs/_build/html/_static/file.png b/docs/_build/html/_static/file.png deleted file mode 100644 index a858a41..0000000 Binary files a/docs/_build/html/_static/file.png and /dev/null differ diff --git a/docs/_build/html/_static/images/logo_binder.svg b/docs/_build/html/_static/images/logo_binder.svg deleted file mode 100644 index 45fecf7..0000000 --- a/docs/_build/html/_static/images/logo_binder.svg +++ /dev/null @@ -1,19 +0,0 @@ - - - diff --git a/docs/_build/html/_static/images/logo_colab.png b/docs/_build/html/_static/images/logo_colab.png deleted file mode 100644 index b7560ec..0000000 Binary files a/docs/_build/html/_static/images/logo_colab.png and /dev/null differ diff --git a/docs/_build/html/_static/images/logo_deepnote.svg b/docs/_build/html/_static/images/logo_deepnote.svg deleted file mode 100644 index fa77ebf..0000000 --- a/docs/_build/html/_static/images/logo_deepnote.svg +++ /dev/null @@ -1 +0,0 @@ - diff --git a/docs/_build/html/_static/images/logo_jupyterhub.svg b/docs/_build/html/_static/images/logo_jupyterhub.svg deleted file mode 100644 index 60cfe9f..0000000 --- a/docs/_build/html/_static/images/logo_jupyterhub.svg +++ /dev/null @@ -1 +0,0 @@ - diff --git a/docs/_build/html/_static/jquery.js b/docs/_build/html/_static/jquery.js deleted file mode 100644 index c4c6022..0000000 --- a/docs/_build/html/_static/jquery.js +++ /dev/null @@ -1,2 +0,0 @@ -/*! jQuery v3.6.0 | (c) OpenJS Foundation and other contributors | jquery.org/license */ -!function(e,t){"use strict";"object"==typeof module&&"object"==typeof module.exports?module.exports=e.document?t(e,!0):function(e){if(!e.document)throw new Error("jQuery requires a window with a document");return t(e)}:t(e)}("undefined"!=typeof window?window:this,function(C,e){"use strict";var t=[],r=Object.getPrototypeOf,s=t.slice,g=t.flat?function(e){return t.flat.call(e)}:function(e){return t.concat.apply([],e)},u=t.push,i=t.indexOf,n={},o=n.toString,v=n.hasOwnProperty,a=v.toString,l=a.call(Object),y={},m=function(e){return"function"==typeof e&&"number"!=typeof e.nodeType&&"function"!=typeof e.item},x=function(e){return null!=e&&e===e.window},E=C.document,c={type:!0,src:!0,nonce:!0,noModule:!0};function b(e,t,n){var r,i,o=(n=n||E).createElement("script");if(o.text=e,t)for(r in c)(i=t[r]||t.getAttribute&&t.getAttribute(r))&&o.setAttribute(r,i);n.head.appendChild(o).parentNode.removeChild(o)}function w(e){return null==e?e+"":"object"==typeof e||"function"==typeof e?n[o.call(e)]||"object":typeof e}var f="3.6.0",S=function(e,t){return new S.fn.init(e,t)};function p(e){var t=!!e&&"length"in e&&e.length,n=w(e);return!m(e)&&!x(e)&&("array"===n||0===t||"number"==typeof t&&0