validmind 2.0.1__py3-none-any.whl → 2.1.0__py3-none-any.whl

validmind/vm_models/dataset.py

@@ -6,12 +6,15 @@
 Dataset class wrapper
 """
 
+import warnings
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 
 import numpy as np
 import pandas as pd
+import polars as pl
 
+from validmind.errors import MissingOrInvalidModelPredictFnError
 from validmind.logging import get_logger
 from validmind.vm_models.model import VMModel
 
@@ -35,9 +38,246 @@ class VMDataset(ABC):
         pass
 
     @abstractmethod
-    def serialize(self):
+    def assign_predictions(
+        self,
+        model,
+        prediction_values: list = None,
+        prediction_probabilities: list = None,
+        prediction_column=None,
+        probability_column=None,
+    ):
         """
-        Serializes the dataset to a dictionary.
+        Assigns predictions to the dataset for a given model or prediction values.
+        The dataset is updated with a new column containing the predictions.
+        """
+        pass
+
+    @abstractmethod
+    def get_extra_column(self, column_name):
+        """
+        Returns the values of the specified extra column.
+
+        Args:
+            column_name (str): The name of the extra column.
+
+        Returns:
+            np.ndarray: The values of the extra column.
+        """
+        pass
+
+    @abstractmethod
+    def add_extra_column(self, column_name, column_values=None):
+        """
+        Adds an extra column to the dataset without modifying the dataset `features` and `target` columns.
+
+        Args:
+            column_name (str): The name of the extra column.
+            column_values (np.ndarray, optional): The values of the extra column.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def input_id(self) -> str:
+        """
+        Returns input id of dataset.
+
+        Returns:
+            str: input_id.
+        """
+        return self.input_id
+
+    @property
+    @abstractmethod
+    def columns(self) -> list:
+        """
+        Returns the the list of columns in the dataset.
+
+        Returns:
+            List[str]: The columns list.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def target_column(self) -> str:
+        """
+        Returns the target column name of the dataset.
+
+        Returns:
+            str: The target column name.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def feature_columns(self) -> list:
+        """
+        Returns the feature columns of the dataset. If _feature_columns is None,
+        it returns all columns except the target column.
+
+        Returns:
+            list: The list of feature column names.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def text_column(self) -> str:
+        """
+        Returns the text column of the dataset.
+
+        Returns:
+            str: The text column name.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def x(self) -> np.ndarray:
+        """
+        Returns the input features (X) of the dataset.
+
+        Returns:
+            np.ndarray: The input features.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def y(self) -> np.ndarray:
+        """
+        Returns the target variables (y) of the dataset.
+
+        Returns:
+            np.ndarray: The target variables.
+        """
+        pass
+
+    @abstractmethod
+    def y_pred(self, model) -> np.ndarray:
+        """
+        Returns the prediction values (y_pred) of the dataset for a given model.
+
+        Returns:
+            np.ndarray: The prediction values.
+        """
+        pass
+
+    def y_prob(self, model) -> np.ndarray:
+        """
+        Returns the prediction probabilities (y_prob) of the dataset for a given model.
+
+        Returns:
+            np.ndarray: The prediction probabilities.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def df(self):
+        """
+        Returns the dataset as a pandas DataFrame.
+
+        Returns:
+            pd.DataFrame: The dataset as a DataFrame.
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def copy(self):
+        """
+        Returns a copy of the raw_dataset dataframe.
+        """
+        pass
+
+    @abstractmethod
+    def x_df(self):
+        """
+        Returns the non target and prediction columns.
+
+        Returns:
+            pd.DataFrame: The non target and prediction columns .
+        """
+        pass
+
+    @abstractmethod
+    def y_df(self):
+        """
+        Returns the target columns (y) of the dataset.
+
+        Returns:
+            pd.DataFrame: The target columns.
+        """
+        pass
+
+    @abstractmethod
+    def y_pred_df(self, model):
+        """
+        Returns the target columns (y) of the dataset.
+
+        Returns:
+            pd.DataFrame: The target columns.
+        """
+        pass
+
+    @abstractmethod
+    def y_prob_df(self, model):
+        """
+        Returns the target columns (y) of the dataset.
+
+        Returns:
+            pd.DataFrame: The target columns.
+        """
+        pass
+
+    @abstractmethod
+    def prediction_column(self, model) -> str:
+        """
+        Returns the prediction column name of the dataset.
+
+        Returns:
+            str: The prediction column name.
+        """
+        pass
+
+    def probability_column(self, model) -> str:
+        """
+        Returns the probability column name of the dataset.
+
+        Returns:
+            str: The probability column name.
+        """
+        pass
+
+    @abstractmethod
+    def get_features_columns(self):
+        """
+        Returns the column names of the feature variables.
+
+        Returns:
+            List[str]: The column names of the feature variables.
+        """
+        pass
+
+    @abstractmethod
+    def get_numeric_features_columns(self):
+        """
+        Returns the column names of the numeric feature variables.
+
+        Returns:
+            List[str]: The column names of the numeric feature variables.
+        """
+        pass
+
+    @abstractmethod
+    def get_categorical_features_columns(self):
+        """
+        Returns the column names of the categorical feature variables.
+
+        Returns:
+            List[str]: The column names of the categorical feature variables.
         """
         pass
 
@@ -62,6 +302,7 @@ class NumpyDataset(VMDataset):
     _extra_columns: dict = field(
         default_factory=lambda: {
             "prediction_columns": {},
+            "probability_columns": {},
             "group_by_column": None,
         }
     )
@@ -134,7 +375,7 @@ class NumpyDataset(VMDataset):
         # initialize target column
         self._target_column = target_column
         # initialize extra columns
-        self.__set_extra_columns(extra_columns, model)
+        self.__set_extra_columns(extra_columns)
         # initialize feature columns
         self.__set_feature_columns(feature_columns)
         # initialize text column, target class labels and options
@@ -144,10 +385,11 @@ class NumpyDataset(VMDataset):
         if model:
             self.assign_predictions(model)
 
-    def __set_extra_columns(self, extra_columns, model):
+    def __set_extra_columns(self, extra_columns):
         if extra_columns is None:
             extra_columns = {
                 "prediction_columns": {},
+                "probability_columns": {},
                 "group_by_column": None,
             }
         self._extra_columns = extra_columns
@@ -187,30 +429,91 @@ class NumpyDataset(VMDataset):
 
         return df
 
+    def __model_id_in_probability_columns(self, model, probability_column):
+        return model.input_id in self._extra_columns.get("probability_columns", {})
+
     def __model_id_in_prediction_columns(self, model, prediction_column):
         return model.input_id in self._extra_columns.get("prediction_columns", {})
 
     def __assign_prediction_values(self, model, pred_column, prediction_values):
+        # Link the prediction column with the model
         self._extra_columns.setdefault("prediction_columns", {})[
             model.input_id
         ] = pred_column
-        self._raw_dataset = np.hstack(
-            (self._raw_dataset, np.array(prediction_values).reshape(-1, 1))
+
+        # Check if the predictions are multi-dimensional (e.g., embeddings)
+        is_multi_dimensional = (
+            isinstance(prediction_values, np.ndarray) and prediction_values.ndim > 1
         )
-        self._columns.append(pred_column)
-        self._df[pred_column] = prediction_values
+
+        if is_multi_dimensional:
+            # For multi-dimensional outputs, convert to a list of lists to store in DataFrame
+            self._df[pred_column] = list(map(list, prediction_values))
+        else:
+            # If not multi-dimensional or a standard numpy array, reshape for compatibility
+            self._raw_dataset = np.hstack(
+                (self._raw_dataset, np.array(prediction_values).reshape(-1, 1))
+            )
+            self._df[pred_column] = prediction_values
+
+        # Update the dataset columns list
+        if pred_column not in self._columns:
+            self._columns.append(pred_column)
+
+    def __assign_prediction_probabilities(
+        self, model, prob_column, prediction_probabilities
+    ):
+        # Link the prediction column with the model
+        self._extra_columns.setdefault("probability_columns", {})[
+            model.input_id
+        ] = prob_column
+
+        # Check if the predictions are multi-dimensional (e.g., embeddings)
+        is_multi_dimensional = (
+            isinstance(prediction_probabilities, np.ndarray)
+            and prediction_probabilities.ndim > 1
+        )
+
+        if is_multi_dimensional:
+            # For multi-dimensional outputs, convert to a list of lists to store in DataFrame
+            self._df[prob_column] = list(map(list, prediction_probabilities))
+        else:
+            # If not multi-dimensional or a standard numpy array, reshape for compatibility
+            self._raw_dataset = np.hstack(
+                (self._raw_dataset, np.array(prediction_probabilities).reshape(-1, 1))
+            )
+            self._df[prob_column] = prediction_probabilities
+
+        # Update the dataset columns list
+        if prob_column not in self._columns:
+            self._columns.append(prob_column)
 
     def assign_predictions(  # noqa: C901 - we need to simplify this method
         self,
         model,
         prediction_values: list = None,
+        prediction_probabilities: list = None,
         prediction_column=None,
+        probability_column=None,
     ):
+        def _is_probability(output):
+            """Check if the output from the predict method is probabilities."""
+            # This is a simple check that assumes output is probabilities if they lie between 0 and 1
+            if np.all((output >= 0) & (output <= 1)):
+                # Check if there is at least one element that is neither 0 nor 1
+                if np.any((output > 0) & (output < 1)):
+                    return True
+            return np.all((output >= 0) & (output <= 1)) and np.any(
+                (output > 0) & (output < 1)
+            )
+
+        # Step 1: Check for Model Presence
         if not model:
             raise ValueError(
                 "Model must be provided to link prediction column with the dataset"
             )
 
+        # Step 2: Prediction Column Provided
         if prediction_column:
             if prediction_column not in self.columns:
                 raise ValueError(
@@ -225,6 +528,8 @@ class NumpyDataset(VMDataset):
             self._extra_columns.setdefault("prediction_columns", {})[
                 model.input_id
             ] = prediction_column
+
+        # Step 4: Prediction Values Provided without Specific Column
         elif prediction_values is not None:
             if len(prediction_values) != self.df.shape[0]:
                 raise ValueError(
@@ -232,13 +537,58 @@ class NumpyDataset(VMDataset):
                 )
             pred_column = f"{model.input_id}_prediction"
             if pred_column in self.columns:
-                raise ValueError(
-                    f"Prediction column {pred_column} already exists in the dataset"
+                warnings.warn(
+                    f"Prediction column {pred_column} already exists in the dataset, overwriting the existing predictions",
+                    UserWarning,
                 )
+
+            logger.info(
+                f"Assigning prediction values to column '{pred_column}' and linked to model '{model.input_id}'"
+            )
             self.__assign_prediction_values(model, pred_column, prediction_values)
+
+        # Step 3: Probability Column Provided
+        if probability_column:
+            if probability_column not in self.columns:
+                raise ValueError(
+                    f"Probability column {probability_column} doesn't exist in the dataset"
+                )
+            if self.__model_id_in_probability_columns(
+                model=model, probability_column=probability_column
+            ):
+                raise ValueError(
+                    f"Probability column {probability_column} already linked to the VM model"
+                )
+            self._extra_columns.setdefault("probability_columns", {})[
+                model.input_id
+            ] = probability_column
+
+        # Step 5: Prediction Probabilities Provided without Specific Column
+        elif prediction_probabilities is not None:
+            if len(prediction_probabilities) != self.df.shape[0]:
+                raise ValueError(
+                    "Length of prediction probabilities doesn't match number of rows of the dataset"
+                )
+            prob_column = f"{model.input_id}_probabilities"
+            if prob_column in self.columns:
+                warnings.warn(
+                    f"Probability column {prob_column} already exists in the dataset, overwriting the existing probabilities",
+                    UserWarning,
+                )
+
+            logger.info(
+                f"Assigning prediction probabilities to column '{prob_column}' and linked to model '{model.input_id}'"
+            )
+            self.__assign_prediction_probabilities(
+                model, prob_column, prediction_probabilities
+            )
+
+        # Step 6: Neither Specific Column Nor Values Provided
         elif not self.__model_id_in_prediction_columns(
             model=model, prediction_column=prediction_column
         ):
+
+            # Compute prediction values directly from the VM model
             pred_column = f"{model.input_id}_prediction"
             if pred_column in self.columns:
                 logger.info(
@@ -256,12 +606,107 @@ class NumpyDataset(VMDataset):
             )
 
             prediction_values = np.array(model.predict(x_only))
-            self.__assign_prediction_values(model, pred_column, prediction_values)
+
+            # Check if the prediction values are probabilities
+            if _is_probability(prediction_values):
+
+                threshold = 0.5
+
+                logger.info(
+                    "Predict method returned probabilities instead of direct labels or regression values. "
+                    + "This implies the model is likely configured for a classification task with probability output."
+                )
+                prob_column = f"{model.input_id}_probabilities"
+                logger.info(
+                    f"Assigning probabilities to column '{prob_column}' and computing class labels using a threshold of {threshold}."
+                )
+                self.__assign_prediction_probabilities(
+                    model, prob_column, prediction_values
+                )
+
+                # Convert probabilities to class labels based on the threshold
+                prediction_classes = (prediction_values > threshold).astype(int)
+                self.__assign_prediction_values(model, pred_column, prediction_classes)
+
+            else:
+
+                # If not assign the prediction values directly
+                pred_column = f"{model.input_id}_prediction"
+                self.__assign_prediction_values(model, pred_column, prediction_values)
+
+                try:
+                    logger.info("Running predict_proba()... This may take a while")
+                    prediction_probabilities = np.array(model.predict_proba(x_only))
+                    prob_column = f"{model.input_id}_probabilities"
+                    self.__assign_prediction_probabilities(
+                        model, prob_column, prediction_probabilities
+                    )
+                except MissingOrInvalidModelPredictFnError:
+                    # Log that predict_proba is not available or failed
+                    logger.warn(
+                        f"Model class '{model.__class__}' does not have a compatible predict_proba implementation."
+                        + " Please assign predictions directly with vm_dataset.assign_predictions(model, prediction_values)"
+                    )
+
+        # Step 7: Prediction Column Already Linked
         else:
             logger.info(
                 f"Prediction column {self._extra_columns['prediction_columns'][model.input_id]} already linked to the {model.input_id}"
             )
 
+    def get_extra_column(self, column_name):
+        """
+        Returns the values of the specified extra column.
+
+        Args:
+            column_name (str): The name of the extra column.
+
+        Returns:
+            np.ndarray: The values of the extra column.
+        """
+        if column_name not in self.extra_columns:
+            raise ValueError(f"Column {column_name} is not an extra column")
+
+        return self._df[column_name]
+
+    def add_extra_column(self, column_name, column_values=None):
+        """
+        Adds an extra column to the dataset without modifying the dataset `features` and `target` columns.
+
+        Args:
+            column_name (str): The name of the extra column.
+            column_values (np.ndarray, optional): The values of the extra column.
+        """
+        if column_name in self.extra_columns:
+            logger.info(f"Column {column_name} already registered as an extra column")
+            return
+
+        # The column name already exists in the dataset so we just assign the extra column
+        if column_name in self.columns:
+            self._extra_columns[column_name] = column_name
+            logger.info(
+                f"Column {column_name} exists in the dataset, registering as an extra column"
+            )
+            return
+
+        if column_values is None:
+            raise ValueError(
+                "Column values must be provided when the column doesn't exist in the dataset"
+            )
+
+        if len(column_values) != self.df.shape[0]:
+            raise ValueError(
+                "Length of column values doesn't match number of rows of the dataset"
+            )
+
+        self._raw_dataset = np.hstack(
+            (self._raw_dataset, np.array(column_values).reshape(-1, 1))
+        )
+        self._columns.append(column_name)
+        self._df[column_name] = column_values
+        self._extra_columns[column_name] = column_name
+        logger.info(f"Column {column_name} added as an extra column")
+
     @property
     def raw_dataset(self) -> np.ndarray:
         """
@@ -397,21 +842,83 @@ class NumpyDataset(VMDataset):
             ],
         ]
 
-    def y_pred(self, model_id) -> np.ndarray:
+    def y_pred(self, model) -> np.ndarray:
         """
-        Returns the prediction variable (y_pred) of the dataset.
+        Returns the prediction variables for a given model, accommodating
+        both scalar predictions and multi-dimensional outputs such as embeddings.
+
+        Args:
+            model (VMModel): The model whose predictions are sought.
 
         Returns:
-            np.ndarray: The prediction variables.
+            np.ndarray: The prediction variables, either as a flattened array for
+            scalar predictions or as an array of arrays for multi-dimensional outputs.
+        """
+        pred_column = self.prediction_column(model)
+
+        # First, attempt to retrieve the prediction data from the DataFrame
+        if hasattr(self, "_df") and pred_column in self._df.columns:
+            predictions = self._df[pred_column].to_numpy()
+
+            # Check if the predictions are stored as objects (e.g., lists for embeddings)
+            if self._df[pred_column].dtype == object:
+                # Attempt to convert lists to a numpy array
+                try:
+                    predictions = np.stack(predictions)
+                except ValueError as e:
+                    # Handling cases where predictions cannot be directly stacked
+                    raise ValueError(f"Error stacking prediction arrays: {e}")
+        else:
+            # Fallback to using the raw numpy dataset if DataFrame is not available or suitable
+            try:
+                predictions = self.raw_dataset[
+                    :, self.columns.index(pred_column)
+                ].flatten()
+            except IndexError as e:
+                raise ValueError(
+                    f"Prediction column '{pred_column}' not found in raw dataset: {e}"
+                )
+
+        return predictions
+
+    def y_prob(self, model) -> np.ndarray:
         """
-        return self.raw_dataset[
-            :,
-            [
-                self.columns.index(name)
-                for name in self.columns
-                if name == self.prediction_column(model_id=model_id)
-            ],
-        ].flatten()
+        Returns the prediction variables for a given model, accommodating
+        both scalar predictions and multi-dimensional outputs such as embeddings.
+
+        Args:
+            model (str): The ID of the model whose predictions are sought.
+
+        Returns:
+            np.ndarray: The prediction variables, either as a flattened array for
+            scalar predictions or as an array of arrays for multi-dimensional outputs.
+        """
+        prob_column = self.probability_column(model)
+
+        # First, attempt to retrieve the prediction data from the DataFrame
+        if hasattr(self, "_df") and prob_column in self._df.columns:
+            probabilities = self._df[prob_column].to_numpy()
+
+            # Check if the predictions are stored as objects (e.g., lists for embeddings)
+            if self._df[prob_column].dtype == object:
+                # Attempt to convert lists to a numpy array
+                try:
+                    probabilities = np.stack(probabilities)
+                except ValueError as e:
+                    # Handling cases where predictions cannot be directly stacked
+                    raise ValueError(f"Error stacking prediction arrays: {e}")
+        else:
+            # Fallback to using the raw numpy dataset if DataFrame is not available or suitable
+            try:
+                probabilities = self.raw_dataset[
+                    :, self.columns.index(prob_column)
+                ].flatten()
+            except IndexError as e:
+                raise ValueError(
+                    f"Prediction column '{prob_column}' not found in raw dataset: {e}"
+                )
+
+        return probabilities
 
     @property
     def type(self) -> str:
@@ -458,22 +965,32 @@ class NumpyDataset(VMDataset):
         """
         return self._df[self.target_column]
 
-    def y_pred_df(self, model_id):
+    def y_pred_df(self, model):
         """
         Returns the target columns (y) of the dataset.
 
         Returns:
             pd.DataFrame: The target columns.
         """
-        return self._df[self.prediction_column(model_id=model_id)]
+        return self._df[self.prediction_column(model)]
 
-    def prediction_column(self, model_id) -> str:
+    def y_prob_df(self, model):
+        """
+        Returns the target columns (y) of the dataset.
+
+        Returns:
+            pd.DataFrame: The target columns.
+        """
+        return self._df[self.probability_column(model)]
+
+    def prediction_column(self, model) -> str:
         """
         Returns the prediction column name of the dataset.
 
         Returns:
             str: The prediction column name.
         """
+        model_id = model.input_id
         pred_column = self._extra_columns.get("prediction_columns", {}).get(model_id)
         if pred_column is None:
             raise ValueError(
@@ -481,6 +998,21 @@ class NumpyDataset(VMDataset):
             )
         return pred_column
 
+    def probability_column(self, model) -> str:
+        """
+        Returns the prediction column name of the dataset.
+
+        Returns:
+            str: The prediction column name.
+        """
+        model_id = model.input_id
+        prob_column = self._extra_columns.get("probability_columns", {}).get(model_id)
+        if prob_column is None:
+            raise ValueError(
+                f"Probability column is not linked with the given {model_id}"
+            )
+        return prob_column
+
     def serialize(self):
         """
         Serializes the dataset to a dictionary.
@@ -608,10 +1140,15 @@ class DataFrameDataset(NumpyDataset):
 
         Args:
             raw_dataset (pd.DataFrame): The raw dataset as a pandas DataFrame.
+            input_id (str, optional): Identifier for the dataset. Defaults to None.
+            model (VMModel, optional): Model associated with the dataset. Defaults to None.
             target_column (str, optional): The target column of the dataset. Defaults to None.
+            extra_columns (dict, optional): Extra columns to include in the dataset. Defaults to None.
             feature_columns (list, optional): The feature columns of the dataset. Defaults to None.
-            text_column (str, optional): The text column name of the dataset for nlp tasks. Defaults to None.
-            target_class_labels (Dict, optional): The class labels for the target columns. Defaults to None.
+            text_column (str, optional): The text column name of the dataset for NLP tasks. Defaults to None.
+            target_class_labels (dict, optional): The class labels for the target columns. Defaults to None.
+            options (dict, optional): Additional options for the dataset. Defaults to None.
+            date_time_index (bool, optional): Whether to use date-time index. Defaults to False.
         """
         index = None
         if isinstance(raw_dataset.index, pd.Index):
@@ -634,6 +1171,57 @@ class DataFrameDataset(NumpyDataset):
         )
 
 
+@dataclass
+class PolarsDataset(NumpyDataset):
+    """
+    VM dataset implementation for Polars DataFrame.
+    """
+
+    def __init__(
+        self,
+        raw_dataset: pl.DataFrame,
+        input_id: str = None,
+        model: VMModel = None,
+        target_column: str = None,
+        extra_columns: dict = None,
+        feature_columns: list = None,
+        text_column: str = None,
+        target_class_labels: dict = None,
+        options: dict = None,
+        date_time_index: bool = False,
+    ):
+        """
+        Initializes a PolarsDataset instance.
+
+        Args:
+            raw_dataset (pl.DataFrame): The raw dataset as a Polars DataFrame.
+            input_id (str, optional): Identifier for the dataset. Defaults to None.
+            model (VMModel, optional): Model associated with the dataset. Defaults to None.
+            target_column (str, optional): The target column of the dataset. Defaults to None.
+            extra_columns (dict, optional): Extra columns to include in the dataset. Defaults to None.
+            feature_columns (list, optional): The feature columns of the dataset. Defaults to None.
+            text_column (str, optional): The text column name of the dataset for NLP tasks. Defaults to None.
+            target_class_labels (dict, optional): The class labels for the target columns. Defaults to None.
+            options (dict, optional): Additional options for the dataset. Defaults to None.
+            date_time_index (bool, optional): Whether to use date-time index. Defaults to False.
+        """
+        super().__init__(
+            raw_dataset=raw_dataset.to_numpy(),
+            input_id=input_id,
+            model=model,
+            index_name=None,
+            index=None,
+            columns=raw_dataset.columns,
+            target_column=target_column,
+            extra_columns=extra_columns,
+            feature_columns=feature_columns,
+            text_column=text_column,
+            target_class_labels=target_class_labels,
+            options=options,
+            date_time_index=date_time_index,
+        )
+
+
 @dataclass
 class TorchDataset(NumpyDataset):
     """
@@ -668,12 +1256,16 @@ class TorchDataset(NumpyDataset):
             text_column (str, optional): The text column name of the dataset for nlp tasks. Defaults to None.
             target_class_labels (Dict, optional): The class labels for the target columns. Defaults to None.
         """
-        # if we can't import torch, then it's not a PyTorch model
+
         try:
             import torch
         except ImportError:
-            return False
+            raise ImportError(
+                "PyTorch is not installed, please run `pip install validmind[pytorch]`"
+            )
+
         columns = []
+
         for id, tens in zip(range(0, len(raw_dataset.tensors)), raw_dataset.tensors):
             if id == 0 and feature_columns is None:
                 n_cols = tens.shape[1]
@@ -684,9 +1276,11 @@ class TorchDataset(NumpyDataset):
                     ).astype(str)
                 ]
                 columns.append(feature_columns)
+
             elif id == 1 and target_column is None:
                 target_column = "y"
                 columns.append(target_column)
+
             elif id == 2 and extra_columns is None:
                 extra_columns.prediction_column = "y_pred"
                 columns.append(extra_columns.prediction_column)