validmind 2.5.1__tar.gz → 2.5.6__tar.gz

{validmind-2.5.1 → validmind-2.5.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: validmind
-Version: 2.5.1
+Version: 2.5.6
 Summary: ValidMind Developer Framework
 License: Commercial License
 Author: Andres Rodriguez
@@ -49,7 +49,7 @@ Requires-Dist: scorecardpy (>=0.1.9.6,<0.2.0.0)
 Requires-Dist: seaborn
 Requires-Dist: sentencepiece (>=0.2.0,<0.3.0) ; extra == "all" or extra == "huggingface" or extra == "llm"
 Requires-Dist: sentry-sdk (>=1.24.0,<2.0.0)
-Requires-Dist: shap (>=0.42.0,<0.43.0)
+Requires-Dist: shap (==0.44.1)
 Requires-Dist: statsmodels
 Requires-Dist: tabulate (>=0.8.9,<0.9.0)
 Requires-Dist: textblob (>=0.18.0.post0,<0.19.0)

{validmind-2.5.1 → validmind-2.5.6}/pyproject.toml

@@ -10,7 +10,7 @@ description = "ValidMind Developer Framework"
 license = "Commercial License"
 name = "validmind"
 readme = "README.pypi.md"
-version = "2.5.1"
+version = "2.5.6"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<3.12"
@@ -47,7 +47,7 @@ scorecardpy = "^0.1.9.6"
 seaborn = "*"
 sentencepiece = {version = "^0.2.0", optional = true}
 sentry-sdk = "^1.24.0"
-shap = "^0.42.0"
+shap = "0.44.1"
 statsmodels = "*"
 tabulate = "^0.8.9"
 textblob = "^0.18.0.post0"

validmind-2.5.6/validmind/__version__.py

@@ -0,0 +1 @@
+__version__ = "2.5.6"

{validmind-2.5.1 → validmind-2.5.6}/validmind/client.py

@@ -240,6 +240,11 @@ def init_model(
         vm_model = class_obj(
             pipeline=model,
             input_id=input_id,
+            attributes=(
+                ModelAttributes.from_dict(attributes)
+                if attributes
+                else ModelAttributes()
+            ),
         )
         # TODO: Add metadata for pipeline model
         metadata = get_model_info(vm_model)
@@ -248,6 +253,7 @@ def init_model(
             input_id=input_id,
             model=model,  # Trained model instance
             predict_fn=predict_fn,
+            attributes=ModelAttributes.from_dict(attributes) if attributes else None,
             **kwargs,
         )
         metadata = get_model_info(vm_model)

{validmind-2.5.1 → validmind-2.5.6}/validmind/tests/__types__.py

@@ -56,7 +56,6 @@ TestID = Literal[
     "validmind.model_validation.ragas.AnswerSimilarity",
     "validmind.model_validation.ragas.AnswerCorrectness",
     "validmind.model_validation.ragas.ContextRecall",
-    "validmind.model_validation.ragas.ContextRelevancy",
     "validmind.model_validation.ragas.ContextPrecision",
     "validmind.model_validation.ragas.AnswerRelevance",
     "validmind.model_validation.sklearn.RegressionModelsPerformanceComparison",

{validmind-2.5.1 → validmind-2.5.6}/validmind/tests/model_validation/ragas/AnswerCorrectness.py

@@ -105,7 +105,7 @@ def AnswerCorrectness(
         "ground_truth": ground_truth_column,
     }
 
-    df = get_renamed_columns(dataset.df, required_columns)
+    df = get_renamed_columns(dataset._df, required_columns)
 
     result_df = evaluate(
         Dataset.from_pandas(df), metrics=[answer_correctness], **get_ragas_config()

{validmind-2.5.1 → validmind-2.5.6}/validmind/tests/model_validation/ragas/AnswerRelevance.py

@@ -109,7 +109,7 @@ def AnswerRelevance(
         "contexts": contexts_column,
     }
 
-    df = get_renamed_columns(dataset.df, required_columns)
+    df = get_renamed_columns(dataset._df, required_columns)
 
     result_df = evaluate(
         Dataset.from_pandas(df), metrics=[answer_relevancy], **get_ragas_config()

{validmind-2.5.1 → validmind-2.5.6}/validmind/tests/model_validation/ragas/AnswerSimilarity.py

@@ -94,7 +94,7 @@ def AnswerSimilarity(
         "ground_truth": ground_truth_column,
     }
 
-    df = get_renamed_columns(dataset.df, required_columns)
+    df = get_renamed_columns(dataset._df, required_columns)
 
     result_df = evaluate(
         Dataset.from_pandas(df), metrics=[answer_similarity], **get_ragas_config()

{validmind-2.5.1 → validmind-2.5.6}/validmind/tests/model_validation/ragas/AspectCritique.py

@@ -132,7 +132,7 @@ def AspectCritique(
         "contexts": contexts_column,
     }
 
-    df = get_renamed_columns(dataset.df, required_columns)
+    df = get_renamed_columns(dataset._df, required_columns)
 
     built_in_aspects = [aspect_map[aspect] for aspect in aspects]
     custom_aspects = (

{validmind-2.5.1 → validmind-2.5.6}/validmind/tests/model_validation/ragas/ContextEntityRecall.py

@@ -100,7 +100,7 @@ def ContextEntityRecall(
         "contexts": contexts_column,
     }
 
-    df = get_renamed_columns(dataset.df, required_columns)
+    df = get_renamed_columns(dataset._df, required_columns)
 
     result_df = evaluate(
         Dataset.from_pandas(df), metrics=[context_entity_recall], **get_ragas_config()

{validmind-2.5.1 → validmind-2.5.6}/validmind/tests/model_validation/ragas/ContextPrecision.py

@@ -96,7 +96,7 @@ def ContextPrecision(
         "ground_truth": ground_truth_column,
     }
 
-    df = get_renamed_columns(dataset.df, required_columns)
+    df = get_renamed_columns(dataset._df, required_columns)
 
     result_df = evaluate(
         Dataset.from_pandas(df), metrics=[context_precision], **get_ragas_config()

{validmind-2.5.1 → validmind-2.5.6}/validmind/tests/model_validation/ragas/ContextRecall.py

@@ -96,7 +96,7 @@ def ContextRecall(
         "ground_truth": ground_truth_column,
     }
 
-    df = get_renamed_columns(dataset.df, required_columns)
+    df = get_renamed_columns(dataset._df, required_columns)
 
     result_df = evaluate(
         Dataset.from_pandas(df), metrics=[context_recall], **get_ragas_config()

{validmind-2.5.1 → validmind-2.5.6}/validmind/tests/model_validation/ragas/Faithfulness.py

@@ -94,7 +94,7 @@ def Faithfulness(
         "contexts": contexts_column,
     }
 
-    df = get_renamed_columns(dataset.df, required_columns)
+    df = get_renamed_columns(dataset._df, required_columns)
 
     result_df = evaluate(
         Dataset.from_pandas(df), metrics=[faithfulness], **get_ragas_config()

{validmind-2.5.1 → validmind-2.5.6}/validmind/tests/model_validation/sklearn/OverfitDiagnosis.py

@@ -25,51 +25,48 @@ from validmind.vm_models import (
 
 logger = get_logger(__name__)
 
+# TODO: A couple of improvements here could be to:
+# 1. Allow the test to use multiple metrics at once
+# 2. Allow custom functions for computing performance
+
 DEFAULT_THRESHOLD = 0.04
+DEFAULT_CLASSIFICATION_METRIC = "auc"
+DEFAULT_REGRESSION_METRIC = "mse"
 PERFORMANCE_METRICS = {
     "accuracy": {
         "function": metrics.accuracy_score,
-        "is_classification": True,
         "is_lower_better": False,
     },
     "auc": {
         "function": metrics.roc_auc_score,
-        "is_classification": True,
         "is_lower_better": False,
     },
     "f1": {
         "function": metrics.f1_score,
-        "is_classification": True,
         "is_lower_better": False,
     },
     "precision": {
         "function": metrics.precision_score,
-        "is_classification": True,
         "is_lower_better": False,
     },
     "recall": {
         "function": metrics.recall_score,
-        "is_classification": True,
         "is_lower_better": False,
     },
     "mse": {
         "function": metrics.mean_squared_error,
-        "is_classification": False,
         "is_lower_better": True,
     },
     "mae": {
         "function": metrics.mean_absolute_error,
-        "is_classification": False,
         "is_lower_better": True,
     },
     "r2": {
         "function": metrics.r2_score,
-        "is_classification": False,
         "is_lower_better": False,
     },
     "mape": {
         "function": metrics.mean_absolute_percentage_error,
-        "is_classification": False,
         "is_lower_better": True,
     },
 }
@@ -123,20 +120,13 @@ def _compute_metrics(
     if is_classification and metric == "auc":
         # if only one class is present in the data, return 0
         if len(np.unique(y_true)) == 1:
-            results[metric].append(0)
-            return
-
-        score = metric_func(y_true, df_region[prob_column].values)
-
-    # All other classification metrics
-    elif is_classification:
-        score = metric_func(y_true, df_region[pred_column].values)
+            return results[metric].append(0)
 
-    # Regression metrics
-    else:
-        score = metric_func(y_true, df_region[pred_column].values)
+        return results[metric].append(
+            metric_func(y_true, df_region[prob_column].values)
+        )
 
-    results[metric].append(score)
+    return results[metric].append(metric_func(y_true, df_region[pred_column].values))
 
 
 def _plot_overfit_regions(
@@ -219,8 +209,12 @@ def overfit_diagnosis(  # noqa: C901
     is_classification = bool(datasets[0].probability_column(model))
 
     # Set default metric if not provided
-    if metric is None:
-        metric = "auc" if is_classification else "mse"
+    if not metric:
+        metric = (
+            DEFAULT_CLASSIFICATION_METRIC
+            if is_classification
+            else DEFAULT_REGRESSION_METRIC
+        )
         logger.info(
             f"Using default {'classification' if is_classification else 'regression'} metric: {metric}"
         )
@@ -228,19 +222,6 @@ def overfit_diagnosis(  # noqa: C901
     if id(cut_off_threshold) == id(DEFAULT_THRESHOLD):
         logger.info("Using default cut-off threshold of 0.04")
 
-    metric = metric.lower()
-    try:
-        _metric = PERFORMANCE_METRICS[metric.lower()]
-    except KeyError:
-        raise ValueError(
-            f"Invalid metric. Choose from: {', '.join(PERFORMANCE_METRICS.keys())}"
-        )
-
-    if is_classification and not _metric["is_classification"]:
-        raise ValueError(f"Cannot use regression metric ({metric}) for classification.")
-    elif not is_classification and _metric["is_classification"]:
-        raise ValueError(f"Cannot use classification metric ({metric}) for regression.")
-
     train_df = datasets[0].df
     test_df = datasets[1].df
 

validmind-2.5.6/validmind/tests/model_validation/sklearn/RobustnessDiagnosis.py

@@ -0,0 +1,379 @@
+# Copyright © 2023-2024 ValidMind Inc. All rights reserved.
+# See the LICENSE file in the root of this repository for details.
+# SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
+
+from collections import defaultdict
+from dataclasses import dataclass
+from operator import add
+from typing import List, Tuple
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import seaborn as sns
+from sklearn import metrics
+
+from validmind.errors import MissingOrInvalidModelPredictFnError
+from validmind.logging import get_logger
+from validmind.vm_models import (
+    Figure,
+    ResultSummary,
+    ResultTable,
+    ResultTableMetadata,
+    ThresholdTest,
+    ThresholdTestResult,
+    VMDataset,
+    VMModel,
+)
+
+logger = get_logger(__name__)
+
+DEFAULT_DECAY_THRESHOLD = 0.05
+DEFAULT_STD_DEV_LIST = [0.1, 0.2, 0.3, 0.4, 0.5]
+DEFAULT_CLASSIFICATION_METRIC = "auc"
+DEFAULT_REGRESSION_METRIC = "mse"
+PERFORMANCE_METRICS = {
+    "accuracy": {
+        "function": metrics.accuracy_score,
+        "is_lower_better": False,
+    },
+    "auc": {
+        "function": metrics.roc_auc_score,
+        "is_lower_better": False,
+    },
+    "f1": {
+        "function": metrics.f1_score,
+        "is_lower_better": False,
+    },
+    "precision": {
+        "function": metrics.precision_score,
+        "is_lower_better": False,
+    },
+    "recall": {
+        "function": metrics.recall_score,
+        "is_lower_better": False,
+    },
+    "mse": {
+        "function": metrics.mean_squared_error,
+        "is_lower_better": True,
+    },
+    "mae": {
+        "function": metrics.mean_absolute_error,
+        "is_lower_better": True,
+    },
+    "r2": {
+        "function": metrics.r2_score,
+        "is_lower_better": False,
+    },
+    "mape": {
+        "function": metrics.mean_absolute_percentage_error,
+        "is_lower_better": True,
+    },
+}
+
+
+def _add_noise_std_dev(
+    values: List[float], x_std_dev: float
+) -> Tuple[List[float], float]:
+    """
+    Adds Gaussian noise to a list of values.
+    Args:
+        values (list[float]): A list of numerical values to which noise is added.
+        x_std_dev (float): A scaling factor for the standard deviation of the noise.
+    Returns:
+        tuple[list[float], float]: A tuple containing:
+            - A list of noisy values, where each value is the sum of the corresponding value
+            in the input list and a randomly generated value sampled from a Gaussian distribution
+            with mean 0 and standard deviation x_std_dev times the standard deviation of the input list.
+            - The standard deviation of the input list of values.
+    """
+    std_dev = np.std(values)
+    noise_list = np.random.normal(0, x_std_dev * std_dev, size=len(values))
+    noisy_values = list(map(add, noise_list, values))
+
+    return noisy_values
+
+
+def _compute_metric(
+    dataset: VMDataset, model: VMModel, X: pd.DataFrame, metric: str
+) -> float:
+    if metric not in PERFORMANCE_METRICS:
+        raise ValueError(
+            f"Invalid metric: {metric}, expected one of {PERFORMANCE_METRICS.keys()}"
+        )
+
+    if metric == "auc":
+        try:
+            y_proba = model.predict_proba(X)
+        except MissingOrInvalidModelPredictFnError:
+            y_proba = model.predict(X)
+        return metrics.roc_auc_score(dataset.y, y_proba)
+
+    return PERFORMANCE_METRICS[metric]["function"](dataset.y, model.predict(X))
+
+
+def _compute_gap(result: dict, metric: str) -> float:
+    if PERFORMANCE_METRICS[metric]["is_lower_better"]:
+        return result[metric.upper()][-1] - result[metric.upper()][0]
+
+    return result[metric.upper()][0] - result[metric.upper()][-1]
+
+
+def _combine_results(results: List[dict]):
+    final_results = defaultdict(list)
+
+    # Interleave rows from each dictionary
+    for i in range(len(results[0]["Perturbation Size"])):
+        for result in results:
+            for key in result.keys():
+                final_results[key].append(result[key][i])
+
+    return pd.DataFrame(final_results)
+
+
+def _plot_robustness(
+    results: pd.DataFrame, metric: str, threshold: float, columns: List[str]
+):
+    fig, ax = plt.subplots()
+
+    pallete = sns.color_palette("muted", len(results["Dataset"].unique()))
+    sns.lineplot(
+        data=results,
+        x="Perturbation Size",
+        y=metric.upper(),
+        hue="Dataset",
+        style="Dataset",
+        linewidth=3,
+        markers=True,
+        markersize=10,
+        dashes=False,
+        palette=pallete,
+        ax=ax,
+    )
+
+    if PERFORMANCE_METRICS[metric]["is_lower_better"]:
+        y_label = f"{metric.upper()} (lower is better)"
+    else:
+        threshold = -threshold
+        y_label = f"{metric.upper()} (higher is better)"
+
+    # add dotted threshold line
+    for i in range(len(results["Dataset"].unique())):
+        baseline = results[results["Dataset"] == results["Dataset"].unique()[i]][
+            metric.upper()
+        ].iloc[0]
+        ax.axhline(
+            y=baseline + threshold,
+            color=pallete[i],
+            linestyle="dotted",
+        )
+
+    ax.tick_params(axis="x")
+    ax.set_ylabel(y_label, weight="bold", fontsize=18)
+    ax.legend(fontsize=18)
+    ax.set_xlabel(
+        "Perturbation Size (X * Standard Deviation)", weight="bold", fontsize=18
+    )
+    ax.set_title(
+        f"Perturbed Features: {', '.join(columns)}",
+        weight="bold",
+        fontsize=20,
+        wrap=True,
+    )
+
+    # prevent the figure from being displayed
+    plt.close("all")
+
+    return fig
+
+
+# TODO: make this a functional test instead of class-based when appropriate
+# simply have to remove the class and rename this func to OverfitDiagnosis
+def robustness_diagnosis(
+    model: VMModel,
+    datasets: List[VMDataset],
+    metric: str = None,
+    scaling_factor_std_dev_list: List[float] = DEFAULT_STD_DEV_LIST,
+    performance_decay_threshold: float = DEFAULT_DECAY_THRESHOLD,
+):
+    if not metric:
+        metric = (
+            DEFAULT_CLASSIFICATION_METRIC
+            if datasets[0].probability_column(model)
+            else DEFAULT_REGRESSION_METRIC
+        )
+        logger.info(f"Using default metric ({metric.upper()}) for robustness diagnosis")
+
+    if id(scaling_factor_std_dev_list) == id(DEFAULT_STD_DEV_LIST):
+        logger.info(
+            f"Using default scaling factors for the standard deviation of the noise: {DEFAULT_STD_DEV_LIST}"
+        )
+
+    if id(performance_decay_threshold) == id(DEFAULT_DECAY_THRESHOLD):
+        logger.info(
+            f"Using default performance decay threshold of {DEFAULT_DECAY_THRESHOLD}"
+        )
+
+    results = [{} for _ in range(len(datasets))]
+
+    # add baseline results (no perturbation)
+    for dataset, result in zip(datasets, results):
+        result["Perturbation Size"] = [0.0]
+        result["Dataset"] = [f"{dataset.input_id}"]
+        result["Row Count"] = [dataset._df.shape[0]]
+
+        result[metric.upper()] = [
+            _compute_metric(
+                dataset=dataset,
+                model=model,
+                X=dataset.x_df(),
+                metric=metric,
+            )
+        ]
+        result["Performance Decay"] = [0.0]
+        result["Passed"] = [True]
+
+    # Iterate scaling factor for the standard deviation list
+    for x_std_dev in scaling_factor_std_dev_list:
+        for dataset, result in zip(datasets, results):
+
+            result["Perturbation Size"].append(x_std_dev)
+            result["Dataset"].append(result["Dataset"][0])
+            result["Row Count"].append(result["Row Count"][0])
+
+            temp_df = dataset.x_df().copy()
+            for feature in dataset.feature_columns_numeric:
+                temp_df[feature] = _add_noise_std_dev(
+                    values=temp_df[feature].to_list(),
+                    x_std_dev=x_std_dev,
+                )
+
+            result[metric.upper()].append(
+                _compute_metric(
+                    dataset=dataset,
+                    model=model,
+                    X=temp_df,
+                    metric=metric,
+                )
+            )
+            result["Performance Decay"].append(_compute_gap(result, metric))
+            result["Passed"].append(
+                result["Performance Decay"][-1] < performance_decay_threshold
+            )
+
+    results_df = _combine_results(results)
+    fig = _plot_robustness(
+        results=results_df,
+        metric=metric,
+        threshold=performance_decay_threshold,
+        columns=datasets[0].feature_columns_numeric,
+    )
+
+    # rename perturbation size for baseline
+    results_df["Perturbation Size"][
+        results_df["Perturbation Size"] == 0.0
+    ] = "Baseline (0.0)"
+
+    return results_df, fig
+
+
+@dataclass
+class RobustnessDiagnosis(ThresholdTest):
+    """Evaluate the robustness of a machine learning model to noise
+
+    Robustness refers to a model's ability to maintain a high level of performance in
+    the face of perturbations or changes (particularly noise) added to its input data.
+    This test is designed to help gauge how well the model can handle potential real-
+    world scenarios where the input data might be incomplete or corrupted.
+
+    ## Test Methodology
+    This test is conducted by adding Gaussian noise, proportional to a particular standard
+    deviation scale, to numeric input features of the input datasets. The model's
+    performance on the perturbed data is then evaluated using a user-defined metric or the
+    default metric of AUC for classification tasks and MSE for regression tasks. The results
+    are then plotted to visualize the model's performance decay as the perturbation size
+    increases.
+
+    When using this test, it is highly recommended to tailor the performance metric, list
+    of scaling factors for the standard deviation of the noise, and the performance decay
+    threshold to the specific use case of the model being evaluated.
+
+    **Inputs**:
+    - model (VMModel): The trained model to be evaluated.
+    - datasets (List[VMDataset]): A list of datasets to evaluate the model against.
+
+    ## Parameters
+    - metric (str, optional): The performance metric to be used for evaluation. If not
+        provided, the default metric is used based on the task of the model. Default values
+        are "auc" for classification tasks and "mse" for regression tasks.
+    - scaling_factor_std_dev_list (List[float], optional): A list of scaling factors for
+        the standard deviation of the noise to be added to the input features. The default
+        values are [0.1, 0.2, 0.3, 0.4, 0.5].
+    - performance_decay_threshold (float, optional): The threshold for the performance
+        decay of the model. The default value is 0.05.
+    """
+
+    name = "robustness"
+    required_inputs = ["model", "datasets"]
+    default_params = {
+        "metric": None,
+        "scaling_factor_std_dev_list": DEFAULT_STD_DEV_LIST,
+        "performance_decay_threshold": DEFAULT_DECAY_THRESHOLD,
+    }
+    tasks = ["classification", "regression"]
+    tags = [
+        "sklearn",
+        "model_diagnosis",
+        "visualization",
+    ]
+
+    def run(self):
+        results, fig = robustness_diagnosis(
+            model=self.inputs.model,
+            datasets=self.inputs.datasets,
+            metric=self.params["metric"],
+            scaling_factor_std_dev_list=self.params["scaling_factor_std_dev_list"],
+            performance_decay_threshold=self.params["performance_decay_threshold"],
+        )
+
+        return self.cache_results(
+            passed=results["Passed"].all(),
+            test_results_list=[
+                ThresholdTestResult(
+                    test_name=self.params["metric"],
+                    passed=results["Passed"].all(),
+                    values=results.to_dict(orient="records"),
+                )
+            ],
+            figures=[
+                Figure(
+                    for_object=self,
+                    key=f"{self.name}:{self.params['metric']}",
+                    figure=fig,
+                )
+            ],
+        )
+
+    def summary(self, results: List[ThresholdTestResult], _):
+        return ResultSummary(
+            results=[
+                ResultTable(
+                    data=results[0].values,
+                    metadata=ResultTableMetadata(title="Robustness Diagnosis Results"),
+                )
+            ]
+        )
+
+    def test(self):
+        """Unit Test for Robustness Diagnosis Threshold Test"""
+        # Verify the result object is present
+        assert self.result is not None
+
+        # Verify test results and their type
+        assert isinstance(self.result.test_results.results, list)
+
+        # Check for presence and validity of 'values' and 'passed' flag in each result
+        for test_result in self.result.test_results.results:
+            assert "values" in test_result.__dict__
+            assert "passed" in test_result.__dict__
+            assert isinstance(test_result.values, list)

{validmind-2.5.1 → validmind-2.5.6}/validmind/tests/run.py

@@ -405,7 +405,7 @@ def run_test(
 
     if unit_metrics:
         metric_id_name = "".join(word[0].upper() + word[1:] for word in name.split())
-        test_id = f"validmind.composite_test.{metric_id_name}"
+        test_id = f"validmind.composite_metric.{metric_id_name}"
 
         error, TestClass = load_composite_metric(
             unit_metrics=unit_metrics, metric_name=metric_id_name

{validmind-2.5.1 → validmind-2.5.6}/validmind/vm_models/dataset/utils.py

@@ -10,6 +10,7 @@ import pandas as pd
 
 from validmind.errors import MissingOrInvalidModelPredictFnError
 from validmind.logging import get_logger
+from validmind.vm_models.model import ModelTask
 
 logger = get_logger(__name__)
 
@@ -118,8 +119,14 @@ def compute_predictions(model, X, **kwargs) -> tuple:
             "You can pass `prediction_values` or `prediction_columns` to use precomputed predictions"
         )
 
-    # TODO: this is really not ideal/robust and should not be handled by dataset class
-    if probability_values is None and _is_probabilties(prediction_values):
+    if model.attributes.task is ModelTask.REGRESSION:
+        logger.info("Model is configured for regression.")
+        return probability_values, prediction_values
+
+    if probability_values is None and (
+        model.attributes.task is ModelTask.CLASSIFICATION
+        or _is_probabilties(prediction_values)
+    ):
         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."