General_improvements and Doc changes (#351)

* general improvements and doc changes

* fix

Author: manujosephv

docs/tabular_model.md

@@ -129,6 +129,14 @@ For self-supervised learning, there is a different API because the process is di
     options:
         show_root_heading: yes
         heading_level: 4
+::: pytorch_tabular.TabularModel.cross_validate
+    options:
+        show_root_heading: yes
+        heading_level: 4
+::: pytorch_tabular.TabularModel.bagging_predict
+    options:
+        show_root_heading: yes
+        heading_level: 4
 
 # Artifact Saving and Loading
 

src/pytorch_tabular/categorical_encoders.py

@@ -22,6 +22,19 @@
 
 class BaseEncoder:
     def __init__(self, cols, handle_unseen, min_samples, imputed, handle_missing):
+        """Base class for categorical encoders.
+        Args:
+            cols (list): list of columns to encode, or None (then all dataset columns will be encoded at fitting time)
+            handle_unseen (str):
+                'error'  - raise an error if a category unseen at fitting time is found
+                'ignore' - skip unseen categories
+                'impute' - impute new categories to a predefined value, which is same as NAN_CATEGORY
+            min_samples (int): minimum samples to take category as valid
+            imputed (float): value to impute unseen categories
+            handle_missing (str):
+                'error'  - raise an error if missing values are found in columns to encode
+                'impute' - impute missing values to a predefined value, which is same as NAN_CATEGORY
+        """
         self.cols = cols
         self.handle_unseen = handle_unseen
         self.handle_missing = handle_missing
@@ -87,11 +100,21 @@ def _before_fit_check(self, X, y):
             assert X.shape[0] == y.shape[0]
 
     def save_as_object_file(self, path):
+        """Save the encoder as a pickle file.
+
+        Args:
+            path (str): path to save the encoder
+        """
         if not self._mapping:
             raise ValueError("`fit` method must be called before `save_as_object_file`.")
         pickle.dump(self.__dict__, open(path, "wb"))
 
     def load_from_object_file(self, path):
+        """Load the encoder from a pickle file.
+
+        Args:
+            path (str): path to load the encoder
+        """
         for k, v in pickle.load(open(path, "rb")).items():
             setattr(self, k, v)
 

src/pytorch_tabular/config/config.py

@@ -4,7 +4,6 @@
 """Config."""
 import os
 import re
-import warnings
 from dataclasses import MISSING, dataclass, field
 from typing import Any, Dict, Iterable, List, Optional
 
@@ -242,7 +241,8 @@ def __post_init__(self):
 
 @dataclass
 class TrainerConfig:
-    """Trainer configuration
+    """Trainer configuration.
+
     Args:
         batch_size (int): Number of samples in each batch of training
 
@@ -539,7 +539,6 @@ def __post_init__(self):
         if self.accelerator is None:
             self.accelerator = "cpu"
         if self.devices_list is not None:
-            warnings.warn("Ignoring devices in favor of devices_list")
             self.devices = self.devices_list
         delattr(self, "devices_list")
         for key in self.early_stopping_kwargs.keys():

src/pytorch_tabular/feature_extractor.py

@@ -24,6 +24,8 @@ def __init__(self, tabular_model, extract_keys=["backbone_features"], drop_origi
 
         Args:
             tabular_model (TabularModel): The trained TabularModel object
+            extract_keys (list, optional): The keys of the features to extract. Defaults to ["backbone_features"].
+            drop_original (bool, optional): Whether to drop the original columns. Defaults to True.
         """
         assert not (
             isinstance(tabular_model.model, NODEModel)
@@ -102,10 +104,20 @@ def fit_transform(self, X: pd.DataFrame, y=None) -> pd.DataFrame:
         return self.transform(X)
 
     def save_as_object_file(self, path):
+        """Saves the feature extractor as a pickle file.
+
+        Args:
+            path (str): The path to save the file
+        """
         if not self._mapping:
             raise ValueError("`fit` method must be called before `save_as_object_file`.")
         pickle.dump(self.__dict__, open(path, "wb"))
 
     def load_from_object_file(self, path):
+        """Loads the feature extractor from a pickle file.
+
+        Args:
+            path (str): The path to load the file from
+        """
         for k, v in pickle.load(open(path, "rb")).items():
             setattr(self, k, v)

src/pytorch_tabular/models/autoint/autoint.py

@@ -15,6 +15,11 @@
 
 class AutoIntBackbone(nn.Module):
     def __init__(self, config: DictConfig):
+        """Automatic Feature Interaction Network.
+
+        Args:
+            config (DictConfig): config of the model
+        """
         super().__init__()
         self.hparams = config
         self._build_network()

src/pytorch_tabular/models/base_model.py

@@ -157,14 +157,14 @@ def __init__(
             warnings.warn(
                 "Wandb is not installed. Please install wandb to log logits. "
                 "You can install wandb using pip install wandb or install PyTorch Tabular"
-                " using pip install pytorch-tabular[all]"
+                " using pip install pytorch-tabular[extra]"
             )
         if not PLOTLY_INSTALLED:
             self.do_log_logits = False
             warnings.warn(
                 "Plotly is not installed. Please install plotly to log logits. "
                 "You can install plotly using pip install plotly or install PyTorch Tabular"
-                " using pip install pytorch-tabular[all]"
+                " using pip install pytorch-tabular[extra]"
             )
 
     @abstractmethod
@@ -376,7 +376,7 @@ def pack_output(self, y_hat: torch.Tensor, backbone_features: torch.tensor) -> D
         """
         # if self.head is the Identity function it means that we cannot extract backbone features,
         # because the model cannot be divide in backbone and head (i.e. TabNet)
-        if type(self.head) == nn.Identity:
+        if type(self.head) is nn.Identity:
             return {"logits": y_hat}
         return {"logits": y_hat, "backbone_features": backbone_features}
 

src/pytorch_tabular/models/mixture_density/mdn.py

@@ -2,7 +2,6 @@
 # Author: Manu Joseph <manujoseph@gmail.com>
 # For license information, see LICENSE.TXT
 """Mixture Density Models."""
-import warnings
 from typing import Dict, Optional, Union
 
 import torch
@@ -13,7 +12,9 @@
 from pytorch_tabular import models
 from pytorch_tabular.config.config import ModelConfig
 from pytorch_tabular.models.common.heads import blocks
-from pytorch_tabular.models.tab_transformer.tab_transformer import TabTransformerBackbone
+from pytorch_tabular.models.tab_transformer.tab_transformer import (
+    TabTransformerBackbone,
+)
 from pytorch_tabular.tabular_model import getattr_nested
 from pytorch_tabular.utils import get_logger
 
@@ -22,7 +23,7 @@
 try:
     import wandb
 except ImportError:
-    warnings.warn("Wandb not installed. WandbLogger will not work.")
+    pass
 
 logger = get_logger(__name__)
 

src/pytorch_tabular/models/node/config.py

@@ -1,3 +1,4 @@
+import warnings
 from dataclasses import dataclass, field
 from typing import Optional
 
@@ -195,11 +196,27 @@ class NodeConfig(ModelConfig):
         },
     )
 
+    head: Optional[str] = field(
+        default=None,
+    )
+
     _module_src: str = field(default="models.node")
     _model_name: str = field(default="NODEModel")
     _backbone_name: str = field(default="NODEBackbone")
     _config_name: str = field(default="NodeConfig")
 
+    def __post_init__(self):
+        if self.head is not None:
+            warnings.warn(
+                "`head` and `head_config` is ignored as NODE has a specific"
+                " head which subsets the tree outputs. Set `head=None`"
+                " to turn off the warning"
+            )
+        else:
+            # Setting Head to LinearHead for compatibility
+            self.head = "LinearHead"
+        return super().__post_init__()
+
 
 # if __name__ == "__main__":
 #     from pytorch_tabular.utils import generate_doc_dataclass

src/pytorch_tabular/ssl_models/base_model.py

@@ -44,6 +44,16 @@ def __init__(
         custom_optimizer_params: Dict = {},
         **kwargs,
     ):
+        """Base Model for all SSL Models.
+
+        Args:
+            config (DictConfig): Configuration defined by the user
+            mode (str, optional): Mode of the model. Defaults to "pretrain".
+            encoder (Optional[nn.Module], optional): Encoder of the model. Defaults to None.
+            decoder (Optional[nn.Module], optional): Decoder of the model. Defaults to None.
+            custom_optimizer (Optional[torch.optim.Optimizer], optional): Custom optimizer to use. Defaults to None.
+            custom_optimizer_params (Dict, optional): Custom optimizer parameters to use. Defaults to {}.
+        """
         super().__init__()
         assert "inferred_config" in kwargs, "inferred_config not found in initialization arguments"
         inferred_config = kwargs["inferred_config"]
@@ -167,7 +177,9 @@ def test_step(self, batch, batch_idx):
 
     def on_validation_epoch_end(self) -> None:
         if hasattr(self.hparams, "log_logits") and self.hparams.log_logits:
-            warnings.warn("Logging Logits is disabled for SSL tasks")
+            warnings.warn(
+                "Logging Logits is disabled for SSL tasks. Set `log_logits` to False" " to turn off this warning"
+            )
         super().on_validation_epoch_end()
 
     def configure_optimizers(self):

src/pytorch_tabular/tabular_model.py

@@ -714,8 +714,9 @@ def fit(
         else:
             if train is not None:
                 warnings.warn(
-                    "train data is provided but datamodule is provided."
-                    " Ignoring the train data and using the datamodule"
+                    "train data and datamodule is provided."
+                    " Ignoring the train data and using the datamodule."
+                    " Set either one of them to None to avoid this warning."
                 )
         model = self.prepare_model(
             datamodule,
@@ -791,8 +792,9 @@ def pretrain(
         else:
             if train is not None:
                 warnings.warn(
-                    "train data is provided but datamodule is provided."
-                    " Ignoring the train data and using the datamodule"
+                    "train data and datamodule is provided."
+                    " Ignoring the train data and using the datamodule."
+                    " Set either one of them to None to avoid this warning."
                 )
         model = self.prepare_model(
             datamodule,
@@ -1050,8 +1052,9 @@ def finetune(
         else:
             if train is not None:
                 warnings.warn(
-                    "train data is provided but datamodule is provided."
-                    " Ignoring the train data and using the datamodule"
+                    "train data and datamodule is provided."
+                    " Ignoring the train data and using the datamodule."
+                    " Set either one of them to None to avoid this warning."
                 )
         if freeze_backbone:
             for param in self.model.backbone.parameters():
@@ -1197,7 +1200,9 @@ def predict(
                 If classification, it returns probabilities and final prediction
         """
         warnings.warn(
-            "`include_input_features` will be deprecated in the next release.",
+            "`include_input_features` will be deprecated in the next release."
+            " Please add index columns to the test dataframe if you want to"
+            " retain some features like the key or id",
             DeprecationWarning,
         )
         assert all(q <= 1 and q >= 0 for q in quantiles), "Quantiles should be a decimal between 0 and 1"
@@ -1286,6 +1291,11 @@ def predict(
             pred_df["prediction"] = self.datamodule.label_encoder.inverse_transform(
                 np.argmax(point_predictions, axis=1)
             )
+            warnings.warn(
+                "Classification prediction column will be renamed to `{target_col}_prediction` "
+                "in the next release to maintain consistency with regression.",
+                DeprecationWarning,
+            )
         if ret_logits:
             for k, v in logits_predictions.items():
                 v = torch.cat(v, dim=0).numpy()
@@ -1558,6 +1568,7 @@ def explain(
                 Defaults to None.
 
             **kwargs: Additional keyword arguments to be passed to the Captum method `attribute` function.
+
         Returns:
             DataFrame: The dataframe with the feature importance
         """
@@ -1587,7 +1598,7 @@ def explain(
             if len(data) <= 100:
                 warnings.warn(
                     f"{method} gives better results when the number of samples is"
-                    " large. For better results, try usingmore samples or some other"
+                    " large. For better results, try using more samples or some other"
                     " methods like GradientShap which works well on single examples."
                 )
         is_full_baselines = method in ["GradientShap", "DeepLiftShap"]
@@ -1742,6 +1753,7 @@ def cross_validate(
                 fold, they will be valid for all the other folds. Defaults to True.
 
             **kwargs: Additional keyword arguments to be passed to the `fit` method of the model.
+
         Returns:
             DataFrame: The dataframe with the cross validation results
         """
@@ -1900,6 +1912,7 @@ def bagging_predict(
                 Defaults to None.
 
             **kwargs: Additional keyword arguments to be passed to the `fit` method of the model.
+
         Returns:
             DataFrame: The dataframe with the bagged predictions.
         """