autogluon.tabular 1.3.2b20250709__py3-none-any.whl → 1.3.2b20250710__py3-none-any.whl

autogluon/tabular/models/__init__.py

@@ -17,7 +17,10 @@ from .imodels.imodels_models import (
 from .knn.knn_model import KNNModel
 from .lgb.lgb_model import LGBModel
 from .lr.lr_model import LinearModel
+from .realmlp.realmlp_model import RealMLPModel
 from .rf.rf_model import RFModel
+from .tabicl.tabicl_model import TabICLModel
+from .tabm.tabm_model import TabMModel
 from .tabpfn.tabpfn_model import TabPFNModel
 from .tabpfnmix.tabpfnmix_model import TabPFNMixModel
 from .tabular_nn.torch.tabular_nn_torch import TabularNeuralNetTorchModel

autogluon/tabular/models/catboost/callbacks.py

@@ -170,14 +170,15 @@ class EarlyStoppingCallback:
 
         self.eval_metric_name = eval_metric_name
         self.is_max_optimal = is_max_optimal
-        self.is_quantile = self.eval_metric_name.startswith(CATBOOST_QUANTILE_PREFIX)
+        self.is_quantile = CATBOOST_QUANTILE_PREFIX in self.eval_metric_name
 
     def after_iteration(self, info):
         is_best_iter = False
         if self.is_quantile:
             # FIXME: CatBoost adds extra ',' in the metric name if quantile levels are not balanced
             # e.g., 'MultiQuantile:alpha=0.1,0.25,0.5,0.95' becomes 'MultiQuantile:alpha=0.1,,0.25,0.5,0.95'
-            eval_metric_name = [k for k in info.metrics[self.compare_key] if k.startswith(CATBOOST_QUANTILE_PREFIX)][0]
+            # `'Quantile:' in k` catches both multiquantile (MultiQuantile:) and single-quantile mode (Quantile:)
+            eval_metric_name = [k for k in info.metrics[self.compare_key] if CATBOOST_QUANTILE_PREFIX in k][0]
         else:
             eval_metric_name = self.eval_metric_name
         cur_score = info.metrics[self.compare_key][eval_metric_name][-1]

autogluon/tabular/models/catboost/catboost_model.py

@@ -350,8 +350,8 @@ class CatBoostModel(AbstractModel):
         return minimum_resources
 
     def _get_default_resources(self):
-        # logical=False is faster in training
-        num_cpus = ResourceManager.get_cpu_count_psutil(logical=False)
+        # only_physical_cores=True is faster in training
+        num_cpus = ResourceManager.get_cpu_count(only_physical_cores=True)
         num_gpus = 0
         return num_cpus, num_gpus
 

autogluon/tabular/models/catboost/catboost_utils.py

@@ -5,7 +5,7 @@ from autogluon.core.constants import BINARY, MULTICLASS, QUANTILE, REGRESSION, S
 logger = logging.getLogger(__name__)
 
 
-CATBOOST_QUANTILE_PREFIX = "MultiQuantile:"
+CATBOOST_QUANTILE_PREFIX = "Quantile:"
 
 
 # TODO: Add weight support?
@@ -74,8 +74,12 @@ def get_catboost_metric_from_ag_metric(metric, problem_type, quantile_levels=Non
             raise AssertionError(f"quantile_levels must be provided for problem_type = {problem_type}")
         if not all(0 < q < 1 for q in quantile_levels):
             raise AssertionError(f"quantile_levels must fulfill 0 < q < 1, provided quantile_levels: {quantile_levels}")
-        quantile_string = ",".join(str(q) for q in quantile_levels)
-        metric_class = f"{CATBOOST_QUANTILE_PREFIX}alpha={quantile_string}"
+        # Loss function MultiQuantile: can only be used if len(quantile_levels) >= 2, otherwise we must use Quantile:
+        if len(quantile_levels) == 1:
+            metric_class = f"{CATBOOST_QUANTILE_PREFIX}alpha={quantile_levels[0]}"
+        else:
+            quantile_string = ",".join(str(q) for q in quantile_levels)
+            metric_class = f"Multi{CATBOOST_QUANTILE_PREFIX}alpha={quantile_string}"
     else:
         raise AssertionError(f"CatBoost does not support {problem_type} problem type.")
 

autogluon/tabular/models/fastainn/tabular_nn_fastai.py

@@ -584,8 +584,8 @@ class NNFastAiTabularModel(AbstractModel):
         return default_auxiliary_params
 
     def _get_default_resources(self):
-        # logical=False is faster in training
-        num_cpus = ResourceManager.get_cpu_count_psutil(logical=False)
+        # only_physical_cores=True is faster in training
+        num_cpus = ResourceManager.get_cpu_count(only_physical_cores=True)
         num_gpus = 0
         return num_cpus, num_gpus
 
@@ -642,7 +642,7 @@ class NNFastAiTabularModel(AbstractModel):
 
     def _get_maximum_resources(self) -> dict[str, Union[int, float]]:
         # fastai model trains slower when utilizing virtual cores and this issue scale up when the number of cpu cores increases
-        return {"num_cpus": ResourceManager.get_cpu_count_psutil(logical=False)}
+        return {"num_cpus": ResourceManager.get_cpu_count(only_physical_cores=True)}
 
     def get_minimum_resources(self, is_gpu_available=False):
         minimum_resources = {

autogluon/tabular/models/lgb/lgb_model.py

@@ -532,8 +532,8 @@ class LGBModel(AbstractModel):
         return minimum_resources
 
     def _get_default_resources(self):
-        # logical=False is faster in training
-        num_cpus = ResourceManager.get_cpu_count_psutil(logical=False)
+        # only_physical_cores=True is faster in training
+        num_cpus = ResourceManager.get_cpu_count(only_physical_cores=True)
         num_gpus = 0
         return num_cpus, num_gpus
 

autogluon/tabular/models/realmlp/realmlp_model.py

@@ -0,0 +1,347 @@
+"""
+Code Adapted from TabArena: https://github.com/autogluon/tabrepo/blob/main/tabrepo/benchmark/models/ag/realmlp/realmlp_model.py
+
+Model: RealMLP
+Paper: Better by Default: Strong Pre-Tuned MLPs and Boosted Trees on Tabular Data
+Authors: David Holzmüller, Léo Grinsztajn, Ingo Steinwart
+Codebase: https://github.com/dholzmueller/pytabkit
+License: Apache-2.0
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+import time
+from contextlib import contextmanager
+from typing import Literal
+
+import numpy as np
+import pandas as pd
+from sklearn.impute import SimpleImputer
+
+from autogluon.common.utils.pandas_utils import get_approximate_df_mem_usage
+from autogluon.common.utils.resource_utils import ResourceManager
+from autogluon.core.models import AbstractModel
+from autogluon.tabular import __version__
+
+logger = logging.getLogger(__name__)
+
+
+@contextmanager
+def set_logger_level(logger_name: str, level: int):
+    _logger = logging.getLogger(logger_name)
+    old_level = _logger.level
+    _logger.setLevel(level)
+    try:
+        yield
+    finally:
+        _logger.setLevel(old_level)
+
+
+# pip install pytabkit
+class RealMLPModel(AbstractModel):
+    ag_key = "REALMLP"
+    ag_name = "RealMLP"
+    ag_priority = 75
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._imputer = None
+        self._features_to_impute = None
+        self._features_to_keep = None
+        self._indicator_columns = None
+        self._features_bool = None
+        self._bool_to_cat = None
+
+    def get_model_cls(self, default_hyperparameters: Literal["td", "td_s"] = "td"):
+        from pytabkit import RealMLP_TD_Classifier, RealMLP_TD_Regressor, RealMLP_TD_S_Classifier, RealMLP_TD_S_Regressor
+
+        assert default_hyperparameters in ["td", "td_s"]
+        if self.problem_type in ['binary', 'multiclass']:
+            if default_hyperparameters == "td":
+                model_cls = RealMLP_TD_Classifier
+            else:
+                model_cls = RealMLP_TD_S_Classifier
+        else:
+            if default_hyperparameters == "td":
+                model_cls = RealMLP_TD_Regressor
+            else:
+                model_cls = RealMLP_TD_S_Regressor
+        return model_cls
+
+    def _fit(
+        self,
+        X: pd.DataFrame,
+        y: pd.Series,
+        X_val: pd.DataFrame = None,
+        y_val: pd.Series = None,
+        time_limit: float = None,
+        num_cpus: int = 1,
+        num_gpus: float = 0,
+        verbosity: int = 2,
+        **kwargs,
+    ):
+        start_time = time.time()
+
+        try:
+            import pytabkit
+            import torch
+        except ImportError as err:
+            logger.log(
+                40,
+                f"\tFailed to import pytabkit/torch! To use the ReaLMLP model, "
+                f"do: `pip install autogluon.tabular[realmlp]=={__version__}`.",
+            )
+            raise err
+
+        if verbosity == 0:
+            _lightning_log_level = logging.ERROR
+        elif verbosity <= 2:
+            _lightning_log_level = logging.WARNING
+        else:
+            _lightning_log_level = logging.INFO
+
+        # FIXME: code assume we only see one GPU in the fit process.
+        device = "cpu" if num_gpus == 0 else "cuda:0"
+        if (device == "cuda:0") and (not torch.cuda.is_available()):
+            raise AssertionError(
+                "Fit specified to use GPU, but CUDA is not available on this machine. "
+                "Please switch to CPU usage instead.",
+            )
+
+        hyp = self._get_model_params()
+
+        default_hyperparameters = hyp.pop("default_hyperparameters", "td")
+
+        model_cls = self.get_model_cls(default_hyperparameters=default_hyperparameters)
+
+        metric_map = {
+            "roc_auc": "1-auc_ovr_alt",
+            "accuracy": "class_error",
+            "balanced_accuracy": "1-balanced_accuracy",
+            "log_loss": "cross_entropy",
+            "rmse": "rmse",
+            "root_mean_squared_error": "rmse",
+            "r2": "rmse",
+            "mae": "mae",
+            "mean_average_error": "mae",
+        }
+
+        val_metric_name = metric_map.get(self.stopping_metric.name, None)
+
+        init_kwargs = dict()
+
+        if val_metric_name is not None:
+            init_kwargs["val_metric_name"] = val_metric_name
+
+        # TODO: Make this smarter? Maybe use `eval_metric.needs_pred`
+        if hyp["use_ls"] is not None and isinstance(hyp["use_ls"], str) and hyp["use_ls"] == "auto":
+            if val_metric_name is None:
+                hyp["use_ls"] = False
+            elif val_metric_name in ["cross_entropy", "1-auc_ovr_alt"]:
+                hyp["use_ls"] = False
+            else:
+                hyp["use_ls"] = None
+
+        if X_val is None:
+            hyp["use_early_stopping"] = False
+            hyp["val_fraction"] = 0
+
+        bool_to_cat = hyp.pop("bool_to_cat", True)
+        impute_bool = hyp.pop("impute_bool", True)
+        name_categories = hyp.pop("name_categories", True)
+
+        n_features = len(X.columns)
+        if "predict_batch_size" in hyp and isinstance(hyp["predict_batch_size"], str) and hyp["predict_batch_size"] == "auto":
+            # simple heuristic to avoid OOM during inference time
+            # note: this isn't fool-proof, and ignores the actual memory availability of the machine.
+            # note: this is based on an assumption of 32 GB of memory available on the instance
+            # default is 1024
+            hyp["predict_batch_size"] = max(min(int(8192 * 200 / n_features), 8192), 64)
+
+        self.model = model_cls(
+            n_threads=num_cpus,
+            device=device,
+            **init_kwargs,
+            **hyp,
+        )
+
+        X = self.preprocess(X, is_train=True, bool_to_cat=bool_to_cat, impute_bool=impute_bool)
+
+        # FIXME: In rare cases can cause exceptions if name_categories=False, unknown why
+        extra_fit_kwargs = {}
+        if name_categories:
+            cat_col_names = X.select_dtypes(include='category').columns.tolist()
+            extra_fit_kwargs["cat_col_names"] = cat_col_names
+
+        if X_val is not None:
+            X_val = self.preprocess(X_val)
+
+        with set_logger_level("lightning.pytorch", _lightning_log_level):
+            self.model = self.model.fit(
+                X=X,
+                y=y,
+                X_val=X_val,
+                y_val=y_val,
+                time_to_fit_in_seconds=time_limit - (time.time() - start_time) if time_limit is not None else None,
+                **extra_fit_kwargs,
+            )
+
+    def _predict_proba(self, X, **kwargs) -> np.ndarray:
+        with set_logger_level("lightning.pytorch", logging.WARNING):
+            return super()._predict_proba(X=X, kwargs=kwargs)
+
+    # TODO: Move missing indicator + mean fill to a generic preprocess flag available to all models
+    # FIXME: bool_to_cat is a hack: Maybe move to abstract model?
+    def _preprocess(self, X: pd.DataFrame, is_train: bool = False, bool_to_cat: bool = False, impute_bool: bool = True, **kwargs) -> pd.DataFrame:
+        """
+        Imputes missing values via the mean and adds indicator columns for numerical features.
+        Converts indicator columns to categorical features to avoid them being treated as numerical by RealMLP.
+        """
+        X = super()._preprocess(X, **kwargs)
+
+        # FIXME: is copy needed?
+        X = X.copy(deep=True)
+        if is_train:
+            self._bool_to_cat = bool_to_cat
+            self._features_bool = self._feature_metadata.get_features(required_special_types=["bool"])
+            if impute_bool:  # Technically this should do nothing useful because bools will never have NaN
+                self._features_to_impute = self._feature_metadata.get_features(valid_raw_types=["int", "float"])
+                self._features_to_keep = self._feature_metadata.get_features(invalid_raw_types=["int", "float"])
+            else:
+                self._features_to_impute = self._feature_metadata.get_features(valid_raw_types=["int", "float"], invalid_special_types=["bool"])
+                self._features_to_keep = [f for f in self._feature_metadata.get_features() if f not in self._features_to_impute]
+            if self._features_to_impute:
+                self._imputer = SimpleImputer(strategy="mean", add_indicator=True)
+                self._imputer.fit(X=X[self._features_to_impute])
+                self._indicator_columns = [c for c in self._imputer.get_feature_names_out() if c not in self._features_to_impute]
+        if self._imputer is not None:
+            X_impute = self._imputer.transform(X=X[self._features_to_impute])
+            X_impute = pd.DataFrame(X_impute, index=X.index, columns=self._imputer.get_feature_names_out())
+            if self._indicator_columns:
+                # FIXME: Use CategoryFeatureGenerator? Or tell the model which is category
+                # TODO: Add to features_bool?
+                X_impute[self._indicator_columns] = X_impute[self._indicator_columns].astype("category")
+            X = pd.concat([X[self._features_to_keep], X_impute], axis=1)
+        if self._bool_to_cat and self._features_bool:
+            # FIXME: Use CategoryFeatureGenerator? Or tell the model which is category
+            X[self._features_bool] = X[self._features_bool].astype("category")
+        return X
+
+    def _set_default_params(self):
+        default_params = dict(
+            random_state=0,
+
+            # Don't use early stopping by default, seems to work well without
+            use_early_stopping=False,
+            early_stopping_additive_patience=40,
+            early_stopping_multiplicative_patience=3,
+
+            # verdict: use_ls="auto" is much better than None.
+            use_ls="auto",
+
+            # verdict: no impact, but makes more sense to be False.
+            impute_bool=False,
+
+            # verdict: name_categories=True avoids random exceptions being raised in rare cases
+            name_categories=True,
+
+            # verdict: bool_to_cat=True is equivalent to False in terms of quality, but can be slightly faster in training time
+            #  and slightly slower in inference time
+            bool_to_cat=True,
+
+            # verdict: "td" is better than "td_s"
+            default_hyperparameters="td",  # options ["td", "td_s"]
+
+            predict_batch_size="auto",  # if auto, uses AutoGluon's heuristic to set a value between 8192 and 64.
+        )
+        for param, val in default_params.items():
+            self._set_default_param_value(param, val)
+
+    @classmethod
+    def supported_problem_types(cls) -> list[str] | None:
+        return ["binary", "multiclass", "regression"]
+
+    def _get_default_stopping_metric(self):
+        return self.eval_metric
+
+    def _get_default_resources(self) -> tuple[int, int]:
+        # only_physical_cores=True is faster in training
+        num_cpus = ResourceManager.get_cpu_count(only_physical_cores=True)
+        num_gpus = min(ResourceManager.get_gpu_count_torch(), 1)
+        return num_cpus, num_gpus
+
+    def _estimate_memory_usage(self, X: pd.DataFrame, **kwargs) -> int:
+        hyperparameters = self._get_model_params()
+        return self.estimate_memory_usage_static(X=X, problem_type=self.problem_type, num_classes=self.num_classes, hyperparameters=hyperparameters, **kwargs)
+
+    @classmethod
+    def _estimate_memory_usage_static(
+        cls,
+        *,
+        X: pd.DataFrame,
+        hyperparameters: dict = None,
+        **kwargs,
+    ) -> int:
+        """
+        Heuristic memory estimate that correlates strongly with RealMLP's more sophisticated method
+
+        More comprehensive memory estimate logic:
+
+        ```python
+        from typing import Any
+
+        from pytabkit.models.alg_interfaces.nn_interfaces import NNAlgInterface
+        from pytabkit.models.data.data import DictDataset, TensorInfo
+        from pytabkit.models.sklearn.default_params import DefaultParams
+
+        def estimate_realmlp_cpu_ram_gb(hparams: dict[str, Any], n_numerical: int, cat_sizes: list[int], n_classes: int,
+                                        n_samples: int):
+            params = copy.copy(DefaultParams.RealMLP_TD_CLASS if n_classes > 0 else DefaultParams.RealMLP_TD_REG)
+            params.update(hparams)
+
+            ds = DictDataset(tensors=None, tensor_infos=dict(x_cont=TensorInfo(feat_shape=[n_numerical]),
+                                                             x_cat=TensorInfo(cat_sizes=cat_sizes),
+                                                             y=TensorInfo(cat_sizes=[n_classes])), device='cpu',
+                             n_samples=n_samples)
+
+            alg_interface = NNAlgInterface(**params)
+            res = alg_interface.get_required_resources(ds, n_cv=1, n_refit=0, n_splits=1, split_seeds=[0], n_train=n_samples)
+            return res.cpu_ram_gb
+        ```
+
+        """
+        if hyperparameters is None:
+            hyperparameters = {}
+        plr_hidden_1 = hyperparameters.get("plr_hidden_1", 16)
+        plr_hidden_2 = hyperparameters.get("plr_hidden_2", 4)
+        hidden_width = hyperparameters.get("hidden_width", 256)
+
+        num_features = len(X.columns)
+        columns_mem_est = num_features * 8e5
+
+        hidden_1_weight = 0.13
+        hidden_2_weight = 0.42
+        width_factor = math.sqrt(hidden_width / 256 + 0.6)
+
+        columns_mem_est_hidden_1 = columns_mem_est * hidden_1_weight * plr_hidden_1 / 16 * width_factor
+        columns_mem_est_hidden_2 = columns_mem_est * hidden_2_weight * plr_hidden_2 / 16 * width_factor
+        columns_mem_est = columns_mem_est_hidden_1 + columns_mem_est_hidden_2
+
+        dataset_size_mem_est = 5 * get_approximate_df_mem_usage(X).sum()  # roughly 5x DataFrame memory size
+        baseline_overhead_mem_est = 3e8  # 300 MB generic overhead
+
+        mem_estimate = dataset_size_mem_est + columns_mem_est + baseline_overhead_mem_est
+
+        return mem_estimate
+
+    @classmethod
+    def _class_tags(cls) -> dict:
+        return {"can_estimate_memory_usage_static": True}
+
+    def _more_tags(self) -> dict:
+        # TODO: Need to add train params support, track best epoch
+        #  How to mirror RealMLP learning rate scheduler while forcing stopping at a specific epoch?
+        tags = {"can_refit_full": False}
+        return tags

autogluon/tabular/models/rf/rf_model.py

@@ -309,8 +309,9 @@ class RFModel(AbstractModel):
                 if self.model.n_outputs_ == 1:
                     self.model.n_classes_ = [self.model.n_classes_]
             from sklearn.tree._tree import DOUBLE, DTYPE
+            from sklearn.utils.validation import check_X_y
 
-            X, y = self.model._validate_data(X, y, multi_output=True, accept_sparse="csc", dtype=DTYPE)
+            X, y = check_X_y(X, y, multi_output=True, accept_sparse="csc", dtype=DTYPE)
             if y.ndim == 1:
                 # reshape is necessary to preserve the data contiguity against vs
                 # [:, np.newaxis] that does not.

autogluon/tabular/models/tabicl/tabicl_model.py

@@ -0,0 +1,174 @@
+"""
+Code Adapted from TabArena: https://github.com/autogluon/tabrepo/blob/main/tabrepo/benchmark/models/ag/tabicl/tabicl_model.py
+Model: TabICL
+Paper: TabICL: A Tabular Foundation Model for In-Context Learning on Large Data
+Authors: Jingang Qu, David Holzmüller, Gaël Varoquaux, Marine Le Morvan
+Codebase: https://github.com/soda-inria/tabicl
+License: BSD-3-Clause
+"""
+
+from __future__ import annotations
+
+import logging
+
+import pandas as pd
+
+from autogluon.common.utils.pandas_utils import get_approximate_df_mem_usage
+from autogluon.common.utils.resource_utils import ResourceManager
+from autogluon.core.models import AbstractModel
+from autogluon.tabular import __version__
+
+logger = logging.getLogger(__name__)
+
+
+# TODO: Verify if crashes when weights are not yet downloaded and fit in parallel
+class TabICLModel(AbstractModel):
+    ag_key = "TABICL"
+    ag_name = "TabICL"
+    ag_priority = 65
+
+    def get_model_cls(self):
+        from tabicl import TabICLClassifier
+
+        if self.problem_type in ["binary", "multiclass"]:
+            model_cls = TabICLClassifier
+        else:
+            raise AssertionError(f"Unsupported problem_type: {self.problem_type}")
+        return model_cls
+
+    @staticmethod
+    def _get_batch_size(n_cells: int):
+        if n_cells <= 4_000_000:
+            return 8
+        elif n_cells <= 6_000_000:
+            return 4
+        else:
+            return 2
+
+    def _fit(
+        self,
+        X: pd.DataFrame,
+        y: pd.Series,
+        num_cpus: int = 1,
+        num_gpus: int = 0,
+        **kwargs,
+    ):
+        try:
+            import tabicl
+        except ImportError as err:
+            logger.log(
+                40,
+                f"\tFailed to import tabicl! To use the TabICL model, "
+                f"do: `pip install autogluon.tabular[tabicl]=={__version__}`.",
+            )
+            raise err
+
+        from torch.cuda import is_available
+
+        device = "cuda" if num_gpus != 0 else "cpu"
+        if (device == "cuda") and (not is_available()):
+            # FIXME: warn instead and switch to CPU.
+            raise AssertionError(
+                "Fit specified to use GPU, but CUDA is not available on this machine. "
+                "Please switch to CPU usage instead.",
+            )
+
+        model_cls = self.get_model_cls()
+        hyp = self._get_model_params()
+        hyp["batch_size"] = hyp.get("batch_size", self._get_batch_size(X.shape[0] * X.shape[1]))
+        self.model = model_cls(
+            **hyp,
+            device=device,
+            n_jobs=num_cpus,
+        )
+        X = self.preprocess(X)
+        self.model = self.model.fit(
+            X=X,
+            y=y,
+        )
+
+    def _set_default_params(self):
+        default_params = {
+            "random_state": 42,
+        }
+        for param, val in default_params.items():
+            self._set_default_param_value(param, val)
+
+    def _get_default_auxiliary_params(self) -> dict:
+        default_auxiliary_params = super()._get_default_auxiliary_params()
+        default_auxiliary_params.update(
+            {
+                "max_rows": 100000,
+                "max_features": 500,
+            }
+        )
+        return default_auxiliary_params
+
+    @classmethod
+    def supported_problem_types(cls) -> list[str] | None:
+        return ["binary", "multiclass"]
+
+    def _get_default_resources(self) -> tuple[int, int]:
+        num_cpus = ResourceManager.get_cpu_count(only_physical_cores=True)
+        num_gpus = min(ResourceManager.get_gpu_count_torch(), 1)
+        return num_cpus, num_gpus
+
+    def _estimate_memory_usage(self, X: pd.DataFrame, **kwargs) -> int:
+        hyperparameters = self._get_model_params()
+        return self.estimate_memory_usage_static(X=X, problem_type=self.problem_type, num_classes=self.num_classes, hyperparameters=hyperparameters, **kwargs)
+
+    @classmethod
+    def _estimate_memory_usage_static(
+        cls,
+        *,
+        X: pd.DataFrame,
+        hyperparameters: dict = None,
+        **kwargs,
+    ) -> int:
+        """
+        Heuristic memory estimate that is very primitive.
+        Can be vastly improved.
+        """
+        if hyperparameters is None:
+            hyperparameters = {}
+
+        dataset_size_mem_est = 3 * get_approximate_df_mem_usage(X).sum()  # roughly 3x DataFrame memory size
+        baseline_overhead_mem_est = 1e9  # 1 GB generic overhead
+
+        n_rows = X.shape[0]
+        n_features = X.shape[1]
+        batch_size = hyperparameters.get("batch_size", cls._get_batch_size(X.shape[0] * X.shape[1]))
+        embedding_dim = 128
+        bytes_per_float = 4
+        model_mem_estimate = 2 * batch_size * embedding_dim * bytes_per_float * (4 + n_rows) * n_features
+
+        model_mem_estimate *= 1.3  # add 30% buffer
+
+        # TODO: Observed memory spikes above expected values on large datasets, increasing mem estimate to compensate
+        model_mem_estimate *= 1.5
+
+        mem_estimate = model_mem_estimate + dataset_size_mem_est + baseline_overhead_mem_est
+
+        return mem_estimate
+
+    @classmethod
+    def _get_default_ag_args_ensemble(cls, **kwargs) -> dict:
+        """
+        Set fold_fitting_strategy to sequential_local,
+        as parallel folding crashes if model weights aren't pre-downloaded.
+        """
+        default_ag_args_ensemble = super()._get_default_ag_args_ensemble(**kwargs)
+        extra_ag_args_ensemble = {
+            # FIXME: If parallel, uses way more memory, seems to behave incorrectly, so we force sequential.
+            "fold_fitting_strategy": "sequential_local",
+            "refit_folds": True,  # Better to refit the model for faster inference and similar quality as the bag.
+        }
+        default_ag_args_ensemble.update(extra_ag_args_ensemble)
+        return default_ag_args_ensemble
+
+    @classmethod
+    def _class_tags(cls) -> dict:
+        return {"can_estimate_memory_usage_static": True}
+
+    def _more_tags(self) -> dict:
+        return {"can_refit_full": True}