autogluon.tabular 1.3.2b20250712__py3-none-any.whl → 1.3.2b20250714__py3-none-any.whl

autogluon/tabular/models/mitra/_internal/data/dataset_split.py

@@ -0,0 +1,53 @@
+import numpy as np
+from sklearn.model_selection import StratifiedKFold, train_test_split
+
+from ..._internal.config.enums import Task
+
+def make_dataset_split(x: np.ndarray, y: np.ndarray, task: Task, seed: int) -> tuple[np.ndarray, ...]:
+    # Splits the dataset into train and validation sets with ratio 80/20
+
+    if task == Task.REGRESSION:
+        return make_standard_dataset_split(x, y, seed=seed)
+
+    size_of_smallest_class = np.min(np.bincount(y))
+
+    if size_of_smallest_class >= 5:
+        # stratification needs have at least 5 samples in each class if split is 80/20
+        return make_stratified_dataset_split(x, y, seed=seed)
+    else:
+        return make_standard_dataset_split(x, y, seed=seed)
+    
+
+def make_stratified_dataset_split(x, y, n_splits=5, seed=0):
+
+    # Stratify doesn't shuffle the data, so we shuffle it first
+    permutation = np.random.permutation(len(y))
+    x, y = x[permutation], y[permutation]
+
+    min_samples_per_class = np.min(np.bincount(y))
+    
+    # Adjust n_splits based on both total samples and minimum samples per class
+    n_samples = len(y)
+    max_possible_splits = min(n_samples - 1, min_samples_per_class)
+    n_splits = min(n_splits, max_possible_splits)
+    
+    # Ensure we have at least 2 splits if possible
+    if n_samples >= 2 and min_samples_per_class >= 2:
+        n_splits = max(2, n_splits)
+    else:
+        # If we can't do stratified splitting, fall back to standard split
+        return make_standard_dataset_split(x, y, seed)
+
+    skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=seed)
+    indices = next(skf.split(x, y))
+    x_t_train, x_t_valid = x[indices[0]], x[indices[1]] # 80%, 20%
+    y_t_train, y_t_valid = y[indices[0]], y[indices[1]]
+
+    return x_t_train, x_t_valid, y_t_train, y_t_valid
+
+
+def make_standard_dataset_split(x, y, seed):
+        
+    return train_test_split(
+        x, y, test_size=0.2, random_state=seed,
+    )

autogluon/tabular/models/mitra/_internal/data/preprocessor.py

@@ -0,0 +1,420 @@
+from typing import Optional, Self
+
+import random
+import numpy as np
+from loguru import logger
+from sklearn.feature_selection import SelectKBest
+from sklearn.preprocessing import QuantileTransformer, StandardScaler, OrdinalEncoder
+from sklearn.compose import ColumnTransformer
+from sklearn.decomposition import TruncatedSVD
+from sklearn.pipeline import Pipeline, FeatureUnion
+from sklearn.base import BaseEstimator, TransformerMixin
+
+from ..._internal.config.enums import Task
+
+class NoneTransformer(BaseEstimator, TransformerMixin):
+    def fit(self, X, y=None):
+        return self
+    def transform(self, X):
+        return X
+    
+class Preprocessor():
+    """
+    This class is used to preprocess the data before it is pushed through the model.
+    The preprocessor assures that the data has the right shape and is normalized,
+    This way the model always gets the same input distribution, 
+    no matter whether the input data is synthetic or real.
+
+    """
+
+    def __init__(
+            self, 
+            dim_embedding: Optional[int],   # Size of the feature embedding. For some models this is None, which means the embedding does not depend on the number of features
+            n_classes: int,   # Actual number of classes in the dataset, assumed to be numbered 0, ..., n_classes - 1 
+            dim_output: int,  # Maximum number of classes the model has been trained on -> size of the output
+            use_quantile_transformer: bool,
+            use_feature_count_scaling: bool,
+            use_random_transforms: bool,
+            shuffle_classes: bool,
+            shuffle_features: bool,
+            random_mirror_regression: bool,
+            random_mirror_x: bool,
+            task: Task
+        ):
+
+        self.dim_embedding = dim_embedding
+        self.n_classes = n_classes
+        self.dim_output = dim_output
+        self.use_quantile_transformer = use_quantile_transformer
+        self.use_feature_count_scaling = use_feature_count_scaling
+        self.use_random_transforms = use_random_transforms
+        self.shuffle_classes = shuffle_classes
+        self.shuffle_features = shuffle_features
+        self.random_mirror_regression = random_mirror_regression
+        self.random_mirror_x = random_mirror_x
+        self.task = task
+    
+    def fit(self, X: np.ndarray, y: np.ndarray) -> Self:
+        """
+        X: np.ndarray [n_samples, n_features]
+        y: np.ndarray [n_samples]
+        """
+
+        if self.task == Task.CLASSIFICATION:
+            # We assume that y properly presents classes [0, 1, 2, ...] before passing to the preprocessor
+            # If the test set has a class that is not in the training set, we will throw an error
+
+            assert np.all(y < self.n_classes), "y contains class values that are not in the range of n_classes"
+
+        self.compute_pre_nan_mean(X)
+        X = self.impute_nan_features_with_mean(X)
+
+        self.determine_which_features_are_singular(X)
+        X = self.cutoff_singular_features(X, self.singular_features)
+
+        self.determine_which_features_to_select(X, y)
+        X = self.select_features(X)
+
+        if self.use_quantile_transformer:
+            # If use quantile transform is off, it means that the preprocessing will happen on the GPU.
+            X = self.fit_transform_quantile_transformer(X)
+        
+            self.mean, self.std = self.calc_mean_std(X)
+            X = self.normalize_by_mean_std(X, self.mean, self.std)
+        
+        if self.use_random_transforms:
+            X = self.transform_tabpfn(X)
+
+        if self.task == Task.CLASSIFICATION and self.shuffle_classes:
+            self.determine_shuffle_class_order()
+        
+        if self.shuffle_features:
+            self.determine_feature_order(X)
+
+        if self.task == Task.REGRESSION:
+            self.determine_mix_max_scale(y)
+
+        if self.task == Task.REGRESSION and self.random_mirror_regression:
+            self.determine_regression_mirror()
+
+        if self.random_mirror_x:
+            self.determine_mirror(X)
+
+        X[np.isnan(X)] = 0
+        X[np.isinf(X)] = 0
+
+        return self
+    
+
+    def transform_X(self, X: np.ndarray):
+
+        X = self.impute_nan_features_with_mean(X)
+        X = self.cutoff_singular_features(X, self.singular_features)
+        X = self.select_features(X)
+
+        if self.use_quantile_transformer:
+            # If use quantile transform is off, it means that the preprocessing will happen on the GPU.
+
+            X = self.quantile_transformer.transform(X)
+        
+            X = self.normalize_by_mean_std(X, self.mean, self.std)
+
+            if self.use_feature_count_scaling:
+                X = self.normalize_by_feature_count(X)
+        
+        if self.use_random_transforms:
+            X = self.random_transforms.transform(X)
+
+        if self.shuffle_features:
+            X = self.randomize_feature_order(X)
+
+        if self.random_mirror_x:
+            X = self.apply_random_mirror_x(X)
+
+        X = X.astype(np.float32)
+
+        X[np.isnan(X)] = 0
+        X[np.isinf(X)] = 0
+
+        return X
+
+
+    def transform_tabpfn(self, X: np.ndarray):
+       
+        n_samples = X.shape[0]
+        n_features = X.shape[1]
+        
+        use_config1 = random.random() < 0.5 
+        random_state = random.randint(0, 1000000)
+
+        if use_config1:
+            self.random_transforms = Pipeline([
+                ('quantile', QuantileTransformer(
+                    output_distribution="normal",
+                    n_quantiles=max(n_samples // 10, 2),
+                    random_state=random_state
+                )),
+                ('svd', FeatureUnion([
+                    ('passthrough', NoneTransformer()),
+                    ('svd', Pipeline([
+                        ('standard', StandardScaler(with_mean=False)),
+                        ('svd', TruncatedSVD(
+                            algorithm="arpack",
+                            n_components=max(1, min(n_samples // 10 + 1, n_features // 2)),
+                            random_state=random_state
+                        ))
+                    ]))
+                ]))
+            ])
+        else:
+            self.random_transforms = ColumnTransformer([
+                ('ordinal', OrdinalEncoder(
+                    handle_unknown="use_encoded_value",
+                    unknown_value=np.nan
+                ), [])  
+            ], remainder='passthrough')
+        
+        return self.random_transforms.fit_transform(X)
+    
+    
+    def transform_y(self, y: np.ndarray):
+
+        if self.task == Task.CLASSIFICATION:
+            # We assume that y properly presents classes [0, 1, 2, ...] before passing to the preprocessor
+            # If the test set has a class that is not in the training set, we will throw an error
+            assert np.all(y < self.n_classes), "y contains class values that are not in the range of n_classes"
+
+        if self.task == Task.CLASSIFICATION and self.shuffle_classes:
+            y = self.randomize_class_order(y)
+
+        if self.task == Task.REGRESSION:
+            y = self.normalize_y(y)
+
+        if self.task == Task.REGRESSION and self.random_mirror_regression:
+            y = self.apply_random_mirror_regression(y)
+
+        match self.task:
+            case Task.CLASSIFICATION:
+                y = y.astype(np.int64)
+            case Task.REGRESSION:
+                y = y.astype(np.float32)
+
+        return y
+    
+
+    def inverse_transform_y(self, y: np.ndarray):
+        # Function used during the prediction to transform the model output back to the original space
+        # For classification, y is assumed to be logits of shape [n_samples, n_classes]
+
+        match self.task:
+            case Task.CLASSIFICATION:
+                y = self.extract_correct_classes(y)
+
+                if self.shuffle_classes:
+                    y = self.undo_randomize_class_order(y)
+
+            case Task.REGRESSION:
+
+                if  self.random_mirror_regression:
+                    y = self.apply_random_mirror_regression(y)
+
+                y = self.undo_normalize_y(y)
+
+        return y
+
+    
+
+    def fit_transform_quantile_transformer(self, X: np.ndarray) -> np.ndarray:
+
+        n_obs, n_features = X.shape
+        n_quantiles = min(n_obs, 1000)
+        self.quantile_transformer = QuantileTransformer(n_quantiles=n_quantiles, output_distribution='normal')
+        X = self.quantile_transformer.fit_transform(X)
+
+        return X
+
+        
+
+    def determine_which_features_are_singular(self, x: np.ndarray) -> None:
+
+        self.singular_features = np.array([ len(np.unique(x_col)) for x_col in x.T ]) == 1
+        
+
+
+    def determine_which_features_to_select(self, x: np.ndarray, y: np.ndarray) -> None:
+
+        if self.dim_embedding is None:
+            # All features are selected
+            return
+
+        if x.shape[1] > self.dim_embedding:
+            logger.info(f"Number of features is capped at {self.dim_embedding}, but the dataset has {x.shape[1]} features. A subset of {self.dim_embedding} are selected using SelectKBest")
+
+            self.select_k_best = SelectKBest(k=self.dim_embedding)
+            self.select_k_best.fit(x, y)
+
+
+    def compute_pre_nan_mean(self, x: np.ndarray) -> None:
+        """
+        Computes the mean of the data before the NaNs are imputed
+        """
+        self.pre_nan_mean = np.nanmean(x, axis=0)
+
+
+    def impute_nan_features_with_mean(self, x: np.ndarray) -> np.ndarray:
+
+        inds = np.where(np.isnan(x))
+        x[inds] = np.take(self.pre_nan_mean, inds[1])
+        return x
+
+    
+    def select_features(self, x: np.ndarray) -> np.ndarray:
+
+        if self.dim_embedding is None:
+            # All features are selected
+            return x
+
+        if x.shape[1] > self.dim_embedding:
+            x = self.select_k_best.transform(x)
+
+        return x
+    
+
+    def cutoff_singular_features(self, x: np.ndarray, singular_features: np.ndarray) -> np.ndarray:
+
+        if singular_features.any():
+            x = x[:, ~singular_features]
+
+        return x
+
+
+    def calc_mean_std(self, x: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+        """
+        Calculates the mean and std of the training data
+        """
+        mean = x.mean(axis=0)
+        std = x.std(axis=0) + 1e-6
+        return mean, std
+    
+
+    def normalize_by_mean_std(self, x: np.ndarray, mean: np.ndarray, std: np.ndarray) -> np.ndarray:
+        """
+        Normalizes the data by the mean and std
+        """
+
+        x = (x - mean) / std
+        return x
+
+
+    def normalize_by_feature_count(self, x: np.ndarray) -> np.ndarray:
+        """
+        An interesting way of normalization by the tabPFN paper
+        """
+
+        assert self.dim_embedding is not None, "dim_embedding must be set to use this feature count scaling"
+
+        x = x * self.dim_embedding / x.shape[1]
+
+        return x
+
+
+
+    def extend_feature_dim_to_dim_embedding(self, x: np.ndarray, dim_embedding) -> np.ndarray:
+        """
+        Increases the number of features to the number of features the model has been trained on
+        """
+
+        assert self.dim_embedding is not None, "dim_embedding must be set to extend the feature dimension"
+
+        added_zeros = np.zeros((x.shape[0], dim_embedding - x.shape[1]), dtype=np.float32)
+        x = np.concatenate([x, added_zeros], axis=1)
+        return x
+    
+
+    def determine_mix_max_scale(self, y: np.ndarray) -> None:
+        self.y_min = y.min()
+        self.y_max = y.max()
+        assert self.y_min != self.y_max, "y_min and y_max are the same, cannot normalize, regression makes no sense"
+
+    
+    def normalize_y(self, y: np.ndarray) -> np.ndarray:
+        y = (y - self.y_min) / (self.y_max - self.y_min)
+        return y
+    
+
+    def undo_normalize_y(self, y: np.ndarray) -> np.ndarray:
+        y = y * (self.y_max - self.y_min) + self.y_min
+        return y
+    
+
+    def determine_regression_mirror(self) -> None:
+        self.regression_mirror = np.random.choice([True, False], size=(1,)).item()
+
+
+    def apply_random_mirror_regression(self, y: np.ndarray) -> np.ndarray:
+        if self.regression_mirror:
+            y = 1 - y
+        return y
+    
+
+    def determine_mirror(self, x: np.ndarray) -> None:
+
+        n_features = x.shape[1]
+        self.mirror = np.random.choice([1, -1], size=(1, n_features))
+
+
+    def apply_random_mirror_x(self, x: np.ndarray) -> np.ndarray:
+
+        x = x * self.mirror
+        return x
+
+
+    def determine_shuffle_class_order(self) -> None:
+
+        if self.shuffle_classes:
+            self.new_shuffle_classes = np.random.permutation(self.n_classes)
+        else:
+            self.new_shuffle_classes = np.arange(self.n_classes)
+
+    
+    def randomize_class_order(self, y: np.ndarray) -> np.ndarray:
+
+        mapping = { i: self.new_shuffle_classes[i] for i in range(self.n_classes) }
+        y = np.array([mapping[i.item()] for i in y], dtype=np.int64)
+
+        return y    
+    
+    
+    def undo_randomize_class_order(self, y_logits: np.ndarray) -> np.ndarray:
+        """
+        We assume y_logits has shape [n_samples, n_classes]
+        """
+
+        # mapping = {self.new_shuffle_classes[i]: i for i in range(self.n_classes)}
+        mapping = {i: self.new_shuffle_classes[i] for i in range(self.n_classes)}
+        y = np.concatenate([y_logits[:, mapping[i]:mapping[i]+1] for i in range(self.n_classes)], axis=1)
+       
+        return y
+    
+
+    def extract_correct_classes(self, y_logits: np.ndarray) -> np.ndarray:
+        # Even though our network might be able to support 10 classes,
+        # If the problem only has three classes, we should give three classes as output.
+        # We assume y_logits has shape [n_samples, n_classes]
+        y_logits = y_logits[:, :self.n_classes]
+        return y_logits
+
+
+
+    def determine_feature_order(self, x: np.ndarray) -> None:
+
+        n_features = x.shape[1]
+        self.new_feature_order = np.random.permutation(n_features)
+
+
+
+    def randomize_feature_order(self, x: np.ndarray) -> np.ndarray:
+
+        x = x[:, self.new_feature_order]
+
+        return x

autogluon/tabular/models/mitra/_internal/models/base.py

@@ -0,0 +1,21 @@
+import torch
+import torch.nn as nn
+from abc import ABC, abstractmethod
+
+class BaseModel(nn.Module, ABC):
+    
+    def __init__(self):
+        super().__init__()
+    
+    def init_weights(self):
+        """Initialize model weights."""
+        pass
+    
+    @abstractmethod
+    def forward(self, 
+        x_support: torch.Tensor, 
+        y_support: torch.Tensor, 
+        x_query: torch.Tensor, 
+        **kwargs):
+        """Forward pass for the model."""
+        pass

autogluon/tabular/models/mitra/_internal/models/embedding.py

@@ -0,0 +1,182 @@
+
+import einops
+import einx
+import torch
+import torch.nn as nn
+
+
+class Tab2DEmbeddingX(torch.nn.Module):
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+
+        self.dim = dim
+        self.x_embedding = nn.Linear(1, dim)
+
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        
+        x = einx.rearrange('b s f -> b s f 1', x)
+        x = self.x_embedding(x)
+
+        return x
+
+        
+
+class Tab2DQuantileEmbeddingX(torch.nn.Module):
+
+    def __init__(
+            self,
+            dim: int,
+        ) -> None:
+        
+        super().__init__()
+
+        self.dim = dim
+
+        
+    def forward(
+        self, 
+        x_support: torch.Tensor, 
+        x_query__: torch.Tensor, 
+        padding_mask: torch.Tensor,
+        feature_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        
+        """
+        Syntax:
+        b = batch size
+        s = number of observations
+        f = number of features
+        q = number of quantiles
+        """
+
+        batch_size = padding_mask.shape[0]
+        seq_len = einx.sum('b [s]', ~padding_mask)
+        feature_count = einx.sum('b [f]', ~feature_mask)
+
+        # By setting the padded tokens to 9999 we ensure they don't participate in the quantile calculation
+        x_support[padding_mask] = 9999
+
+        q = torch.arange(1, 1000, dtype=torch.float, device=x_support.device) / 1000
+        quantiles = torch.quantile(x_support, q=q, dim=1)
+        quantiles = einx.rearrange('q b f -> (b f) q', quantiles)
+        x_support = einx.rearrange('b s f -> (b f) s', x_support).contiguous()
+        x_query__ = einx.rearrange('b s f -> (b f) s', x_query__).contiguous()
+
+        bucketize = torch.vmap(torch.bucketize, in_dims=(0, 0), out_dims=0)
+        x_support = bucketize(x_support, quantiles).float() 
+        x_query__ = bucketize(x_query__, quantiles).float()
+        x_support = einx.rearrange('(b f) s -> b s f', x_support, b=batch_size).contiguous()
+        x_query__ = einx.rearrange('(b f) s -> b s f', x_query__, b=batch_size).contiguous()
+
+        # If 30% is padded, the minimum will have quantile 0.0 and the maximum will have quantile 0.7 times max_length.
+        # Here we correct the quantiles so that the minimum has quantile 0.0 and the maximum has quantile 1.0.
+        x_support = x_support / seq_len[:, None, None] 
+        x_query__ = x_query__ / seq_len[:, None, None]
+
+        # Make sure that the padding is not used in the calculation of the mean
+        x_support[padding_mask] = 0
+        x_support_mean = einx.sum('b [s] f', x_support, keepdims=True) / seq_len[:, None, None]
+
+        x_support = x_support - x_support_mean
+        x_query__ = x_query__ - x_support_mean
+
+        # Make sure that the padding is not used in the calculation of the variance
+        x_support[padding_mask] = 0
+        x_support_var = einx.sum('b [s] f', x_support**2, keepdims=True) / seq_len[:, None, None]
+
+        x_support = x_support / x_support_var.sqrt()
+        x_query__ = x_query__ / x_support_var.sqrt()
+
+        # In case an x_support feature column contains one unique feature, set the feature to zero 
+        x_support = torch.where(x_support_var == 0, 0, x_support)
+        x_query__ = torch.where(x_support_var == 0, 0, x_query__)
+
+        return x_support, x_query__
+
+
+class Tab2DEmbeddingY(torch.nn.Module):
+
+    def __init__(self, dim: int, n_classes: int) -> None:
+        super().__init__()
+
+        self.dim = dim
+        self.n_classes = n_classes
+        self.y_embedding_support = nn.Linear(1, dim)
+        self.y_embedding_query = nn.Embedding(1, dim)
+
+
+    def forward(self, y_support: torch.Tensor, padding_obs_support: torch.Tensor, n_obs_query: int) -> tuple[torch.Tensor, torch.Tensor]:
+
+        batch_size = y_support.shape[0]
+
+        y_support = y_support.type(torch.float32)
+        y_support = y_support / self.n_classes - 0.5
+        y_support = einops.rearrange(y_support, 'b n -> b n 1')
+
+        y_support = self.y_embedding_support(y_support)
+        y_support[padding_obs_support] = 0
+
+        y_query = torch.zeros((batch_size, n_obs_query, 1), device=y_support.device, dtype=torch.int64)
+        y_query = self.y_embedding_query(y_query)
+
+        return y_support, y_query
+    
+
+class Tab2DEmbeddingYClasses(torch.nn.Module):
+
+    def __init__(
+            self,
+            dim: int,
+            n_classes: int, 
+        ) -> None:
+        
+        super().__init__()
+
+        self.n_classes = n_classes
+        self.dim = dim
+
+        self.y_embedding = nn.Embedding(n_classes, dim,)
+        self.y_mask = nn.Embedding(1, dim) # masking is also modeled as a separate class
+
+
+    def forward(self, y_support: torch.Tensor, padding_obs_support: torch.Tensor, n_obs_query: int) -> tuple[torch.Tensor, torch.Tensor]:
+
+        batch_size = y_support.shape[0]
+        n_obs_support = y_support.shape[1]
+
+        y_support = y_support.type(torch.int64)
+        y_support = einops.rearrange(y_support, 'b n -> b n 1')
+        y_support[padding_obs_support] = 0    # padded tokens are -100 -> set it to zero so nn.Embedding can handle it
+        y_support = self.y_embedding(y_support)
+        y_support[padding_obs_support] = 0    # just to make sure, set it to zero again
+         
+        y_query = torch.zeros((batch_size, n_obs_query, 1), device=y_support.device, dtype=torch.int64)
+        y_query = self.y_mask(y_query)
+
+        return y_support, y_query
+    
+
+class Tab2DEmbeddingYRegression(torch.nn.Module):
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+
+        self.dim = dim
+        self.y_embedding = nn.Linear(1, dim)
+        self.y_mask = nn.Embedding(1, dim)
+
+
+    def forward(self, y_support: torch.Tensor, padding_obs_support: torch.Tensor, n_obs_query: int) -> tuple[torch.Tensor, torch.Tensor]:
+
+        batch_size = y_support.shape[0]
+        y_support = y_support.type(torch.float32)
+        y_support = einops.rearrange(y_support, 'b n -> b n 1')
+        y_support = self.y_embedding(y_support)
+        y_support[padding_obs_support] = 0
+
+        y_query = torch.zeros((batch_size, n_obs_query, 1), device=y_support.device, dtype=torch.int64)
+        y_query = self.y_mask(y_query)
+
+        return y_support, y_query