validmind 2.8.28__py3-none-any.whl → 2.8.29__py3-none-any.whl

validmind/vm_models/dataset/dataset.py

@@ -8,7 +8,7 @@ Dataset class wrapper
 
 import warnings
 from copy import deepcopy
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, Optional
 
 import numpy as np
 import pandas as pd
@@ -258,69 +258,91 @@ class VMDataset(VMInput):
                 f"Options {kwargs} are not supported for this input"
             )
 
-    def assign_predictions(
-        self,
-        model: VMModel,
-        prediction_column: Optional[str] = None,
-        prediction_values: Optional[List[Any]] = None,
-        probability_column: Optional[str] = None,
-        probability_values: Optional[List[float]] = None,
-        prediction_probabilities: Optional[
-            List[float]
-        ] = None,  # DEPRECATED: use probability_values
-        **kwargs: Dict[str, Any],
-    ) -> None:
-        """Assign predictions and probabilities to the dataset.
-
-        Args:
-            model (VMModel): The model used to generate the predictions.
-            prediction_column (Optional[str]): The name of the column containing the predictions.
-            prediction_values (Optional[List[Any]]): The values of the predictions.
-            probability_column (Optional[str]): The name of the column containing the probabilities.
-            probability_values (Optional[List[float]]): The values of the probabilities.
-            prediction_probabilities (Optional[List[float]]): DEPRECATED: The values of the probabilities.
-            **kwargs: Additional keyword arguments that will get passed through to the model's `predict` method.
-        """
+    def _handle_deprecated_parameters(
+        self, prediction_probabilities, probability_values
+    ):
+        """Handle deprecated parameters and return the correct probability values."""
         if prediction_probabilities is not None:
             warnings.warn(
                 "The `prediction_probabilities` argument is deprecated. Use `probability_values` instead.",
                 DeprecationWarning,
             )
-            probability_values = prediction_probabilities
-
-        self._validate_assign_predictions(
-            model,
-            prediction_column,
-            prediction_values,
-            probability_column,
-            probability_values,
-        )
+            return prediction_probabilities
+        return probability_values
 
+    def _check_existing_predictions(self, model):
+        """Check for existing predictions and probabilities, warn if overwriting."""
         if self.prediction_column(model):
             logger.warning("Model predictions already assigned... Overwriting.")
 
         if self.probability_column(model):
             logger.warning("Model probabilities already assigned... Overwriting.")
 
-        # if the user passes a column name, we assume it has precomputed predictions
+    def _get_precomputed_values(self, prediction_column, probability_column):
+        """Get precomputed prediction and probability values from existing columns."""
+        prediction_values = None
+        probability_values = None
+
         if prediction_column:
             prediction_values = self._df[prediction_column].values
 
             if probability_column:
                 probability_values = self._df[probability_column].values
 
+        return prediction_values, probability_values
+
+    def _compute_predictions_if_needed(self, model, prediction_values, **kwargs):
+        """Compute predictions if not provided."""
         if prediction_values is None:
             X = self.df if isinstance(model, (FunctionModel, PipelineModel)) else self.x
-            probability_values, prediction_values = compute_predictions(
-                model, X, **kwargs
+            return compute_predictions(model, X, **kwargs)
+        return None, prediction_values
+
+    def _handle_dictionary_predictions(self, model, prediction_values):
+        """Handle dictionary predictions by converting to separate columns."""
+        if (
+            prediction_values is not None
+            and len(prediction_values) > 0
+            and isinstance(prediction_values[0], dict)
+        ):
+            df_prediction_values = pd.DataFrame.from_dict(
+                prediction_values, orient="columns"
             )
 
-        prediction_column = prediction_column or f"{model.input_id}_prediction"
+            for column_name in df_prediction_values.columns.tolist():
+                values = df_prediction_values[column_name].values
+
+                if column_name == "prediction":
+                    prediction_column = f"{model.input_id}_prediction"
+                    self._add_column(prediction_column, values)
+                    self.prediction_column(model, prediction_column)
+                else:
+                    self._add_column(f"{model.input_id}_{column_name}", values)
+
+            return (
+                True,
+                None,
+            )  # Return True to indicate dictionary handled, None for prediction_column
+        return False, None
+
+    def _add_prediction_columns(
+        self,
+        model,
+        prediction_column,
+        prediction_values,
+        probability_column,
+        probability_values,
+    ):
+        """Add prediction and probability columns to the dataset."""
+        if prediction_column is None:
+            prediction_column = f"{model.input_id}_prediction"
+
         self._add_column(prediction_column, prediction_values)
         self.prediction_column(model, prediction_column)
 
         if probability_values is not None:
-            probability_column = probability_column or f"{model.input_id}_probabilities"
+            if probability_column is None:
+                probability_column = f"{model.input_id}_probabilities"
             self._add_column(probability_column, probability_values)
             self.probability_column(model, probability_column)
         else:
@@ -329,6 +351,91 @@ class VMDataset(VMInput):
                 "Not adding probability column to the dataset."
             )
 
+    def assign_predictions(
+        self,
+        model: VMModel,
+        prediction_column: Optional[str] = None,
+        prediction_values: Optional[Any] = None,
+        probability_column: Optional[str] = None,
+        probability_values: Optional[Any] = None,
+        prediction_probabilities: Optional[
+            Any
+        ] = None,  # DEPRECATED: use probability_values
+        **kwargs: Dict[str, Any],
+    ) -> None:
+        """Assign predictions and probabilities to the dataset.
+
+        Args:
+            model (VMModel): The model used to generate the predictions.
+            prediction_column (Optional[str]): The name of the column containing the predictions.
+            prediction_values (Optional[Any]): The values of the predictions. Can be array-like (list, numpy array, pandas Series, etc.).
+            probability_column (Optional[str]): The name of the column containing the probabilities.
+            probability_values (Optional[Any]): The values of the probabilities. Can be array-like (list, numpy array, pandas Series, etc.).
+            prediction_probabilities (Optional[Any]): DEPRECATED: The values of the probabilities. Use probability_values instead.
+            **kwargs: Additional keyword arguments that will get passed through to the model's `predict` method.
+        """
+        # Handle deprecated parameters
+        probability_values = self._handle_deprecated_parameters(
+            prediction_probabilities, probability_values
+        )
+
+        # Convert pandas Series to numpy array for prediction_values
+        if (
+            hasattr(prediction_values, "values")
+            and hasattr(prediction_values, "index")
+            and hasattr(prediction_values, "dtype")
+        ):
+            prediction_values = prediction_values.values
+
+        # Convert pandas Series to numpy array for probability_values
+        if (
+            hasattr(probability_values, "values")
+            and hasattr(probability_values, "index")
+            and hasattr(probability_values, "dtype")
+        ):
+            probability_values = probability_values.values
+
+        # Validate input parameters
+        self._validate_assign_predictions(
+            model,
+            prediction_column,
+            prediction_values,
+            probability_column,
+            probability_values,
+        )
+
+        # Check for existing predictions and warn if overwriting
+        self._check_existing_predictions(model)
+
+        # Get precomputed values if column names are provided
+        if prediction_column or probability_column:
+            prediction_values, prob_values_from_column = self._get_precomputed_values(
+                prediction_column, probability_column
+            )
+            if prob_values_from_column is not None:
+                probability_values = prob_values_from_column
+
+        # Compute predictions if not provided
+        if prediction_values is None:
+            probability_values, prediction_values = self._compute_predictions_if_needed(
+                model, prediction_values, **kwargs
+            )
+
+        # Handle dictionary predictions
+        is_dict_handled, _ = self._handle_dictionary_predictions(
+            model, prediction_values
+        )
+
+        # Add prediction and probability columns (skip if dictionary was handled)
+        if not is_dict_handled:
+            self._add_prediction_columns(
+                model,
+                prediction_column,
+                prediction_values,
+                probability_column,
+                probability_values,
+            )
+
     def prediction_column(self, model: VMModel, column_name: str = None) -> str:
         """Get or set the prediction column for a model."""
         if column_name and column_name not in self.columns:

{validmind-2.8.28.dist-info → validmind-2.8.29.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: validmind
-Version: 2.8.28
+Version: 2.8.29
 Summary: ValidMind Library
 License: Commercial License
 Author: Andres Rodriguez