validmind 2.0.7__py3-none-any.whl → 2.1.1__py3-none-any.whl

validmind/tests/model_validation/FeaturesAUC.py

@@ -0,0 +1,110 @@
+# 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 dataclasses import dataclass
+
+import numpy as np
+import pandas as pd
+import plotly.graph_objects as go
+from sklearn.metrics import roc_auc_score
+
+from validmind.errors import SkipTestError
+from validmind.logging import get_logger
+from validmind.vm_models import Figure, Metric
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class FeaturesAUC(Metric):
+    """
+    Evaluates the discriminatory power of each individual feature within a binary classification model by calculating the Area Under the Curve (AUC) for each feature separately.
+
+    **Purpose**: The central objective of this metric is to quantify how well each feature on its own can differentiate between the two classes in a binary classification problem. It serves as a univariate analysis tool that can help in pre-modeling feature selection or post-modeling interpretation.
+
+    **Test Mechanism**: For each feature, the metric treats the feature values as raw scores to compute the AUC against the actual binary outcomes. It provides an AUC value for each feature, offering a simple yet powerful indication of each feature's univariate classification strength.
+
+    **Signs of High Risk**:
+    - A feature with a low AUC score may not be contributing significantly to the differentiation between the two classes, which could be a concern if it is expected to be predictive.
+    - Conversely, a surprisingly high AUC for a feature not believed to be informative may suggest data leakage or other issues with the data.
+
+    **Strengths**:
+    - By isolating each feature, it highlights the individual contribution of features to the classification task without the influence of other variables.
+    - Useful for both initial feature evaluation and for providing insights into the model's reliance on individual features after model training.
+
+    **Limitations**:
+    - Does not reflect the combined effects of features or any interaction between them, which can be critical in certain models.
+    - The AUC values are calculated without considering the model's use of the features, which could lead to different interpretations of feature importance when considering the model holistically.
+    - This metric is applicable only to binary classification tasks and cannot be directly extended to multiclass classification or regression without modifications.
+    """
+
+    name = "features_auc"
+    required_inputs = ["model", "dataset"]
+    default_params = {
+        "fontsize": 12,
+        "figure_height": 500,
+    }
+    metadata = {
+        "task_types": ["classification"],
+        "tags": [
+            "feature_importance",
+            "AUC",
+            "visualization",
+        ],
+    }
+
+    def run(self):
+        x = self.inputs.dataset.x_df()
+        y = self.inputs.dataset.y_df()
+
+        if y.nunique() != 2:
+            raise SkipTestError("FeaturesAUC metric requires a binary target variable.")
+
+        aucs = pd.DataFrame(index=x.columns, columns=["AUC"])
+
+        for column in x.columns:
+            feature_values = x[column]
+            if feature_values.nunique() > 1:
+                auc_score = roc_auc_score(y, feature_values)
+                aucs.loc[column, "AUC"] = auc_score
+            else:
+                aucs.loc[
+                    column, "AUC"
+                ] = np.nan  # Not enough unique values to calculate AUC
+
+        # Sorting the AUC scores in descending order
+        sorted_indices = aucs["AUC"].dropna().sort_values(ascending=False).index
+
+        # Plotting the results
+        fig = go.Figure()
+        fig.add_trace(
+            go.Bar(
+                y=[column for column in sorted_indices],
+                x=[aucs.loc[column, "AUC"] for column in sorted_indices],
+                orientation="h",
+            )
+        )
+        fig.update_layout(
+            title_text="Feature AUC Scores",
+            yaxis=dict(
+                tickmode="linear",
+                dtick=1,
+                tickfont=dict(size=self.params["fontsize"]),
+                title="Features",
+                autorange="reversed",  # Ensure that the highest AUC is at the top
+            ),
+            xaxis=dict(title="AUC"),
+            height=self.params["figure_height"],
+        )
+
+        return self.cache_results(
+            metric_value=aucs.to_dict(),
+            figures=[
+                Figure(
+                    for_object=self,
+                    key="features_auc",
+                    figure=fig,
+                ),
+            ],
+        )

validmind/tests/model_validation/MeteorScore.py

@@ -54,7 +54,7 @@ class MeteorScore(Metric):
         meteor_scores = []
 
         for prediction, reference in zip(
-            self.inputs.dataset.y_pred(self.inputs.model.input_id),
+            self.inputs.dataset.y_pred(self.inputs.model),
             self.inputs.dataset.y,
         ):
             # Compute the METEOR score for the current prediction-reference pair

validmind/tests/model_validation/RegardHistogram.py

@@ -57,7 +57,7 @@ class RegardHistogram(Metric):
             raise AttributeError("The 'model' attribute is missing.")
 
         y_true = list(itertools.chain.from_iterable(self.inputs.dataset.y))
-        y_pred = self.inputs.dataset.y_pred(self.inputs.model.input_id)
+        y_pred = self.inputs.dataset.y_pred(self.inputs.model)
 
         if not len(y_true) == len(y_pred):
             raise ValueError(

validmind/tests/model_validation/RegardScore.py

@@ -58,7 +58,7 @@ class RegardScore(Metric):
             raise AttributeError("The 'model' attribute is missing.")
 
         y_true = list(itertools.chain.from_iterable(self.inputs.dataset.y))
-        y_pred = self.inputs.dataset.y_pred(self.inputs.model.input_id)
+        y_pred = self.inputs.dataset.y_pred(self.inputs.model)
 
         if not len(y_true) == len(y_pred):
             raise ValueError(

validmind/tests/model_validation/RegressionResidualsPlot.py

@@ -0,0 +1,127 @@
+# 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 dataclasses import dataclass
+
+import numpy as np
+import plotly.figure_factory as ff
+import plotly.graph_objects as go
+
+from validmind.vm_models import Figure, Metric
+
+
+@dataclass
+class RegressionResidualsPlot(Metric):
+    """
+    Evaluates regression model performance using residual distribution and actual vs. predicted plots.
+
+    **Purpose:**
+    The `RegressionResidualsPlot` metric aims to evaluate the performance of regression models. By generating and
+    analyzing two plots – a distribution of residuals and a scatter plot of actual versus predicted values – this tool
+    helps to visually appraise how well the model predicts and the nature of errors it makes.
+
+    **Test Mechanism:**
+    The process begins by extracting the true output values (`y_true`) and the model's predicted values (`y_pred`).
+    Residuals are computed by subtracting predicted from true values. These residuals are then visualized using a
+    histogram to display their distribution. Additionally, a scatter plot is derived to compare true values against
+    predicted values, together with a "Perfect Fit" line, which represents an ideal match (predicted values equal
+    actual values), facilitating the assessment of the model's predictive accuracy.
+
+    **Signs of High Risk:**
+    - Residuals showing a non-normal distribution, especially those with frequent extreme values.
+    - Significant deviations of predicted values from actual values in the scatter plot.
+    - Sparse density of data points near the "Perfect Fit" line in the scatter plot, indicating poor prediction
+    accuracy.
+    - Visible patterns or trends in the residuals plot, suggesting the model's failure to capture the underlying data
+    structure adequately.
+
+    **Strengths:**
+    - Provides a direct, visually intuitive assessment of a regression model’s accuracy and handling of data.
+    - Visual plots can highlight issues of underfitting or overfitting.
+    - Can reveal systematic deviations or trends that purely numerical metrics might miss.
+    - Applicable across various regression model types.
+
+    **Limitations:**
+    - Relies on visual interpretation, which can be subjective and less precise than numerical evaluations.
+    - May be difficult to interpret in cases with multi-dimensional outputs due to the plots’ two-dimensional nature.
+    - Overlapping data points in the residuals plot can complicate interpretation efforts.
+    - Does not summarize model performance into a single quantifiable metric, which might be needed for comparative or
+    summary analyses.
+    """
+
+    name = "regression_residuals_plot"
+    required_inputs = ["model", "dataset"]
+    metadata = {
+        "task_types": ["regression"],
+        "tags": [
+            "model_performance",
+        ],
+    }
+    default_params = {"bin_size": 0.1}
+
+    def run(self):
+        y_true = self.inputs.dataset.y
+        y_pred = self.inputs.dataset.y_pred(self.inputs.model)
+        # Calculate residuals
+        residuals = y_true.flatten() - y_pred.flatten()
+        # Create residuals plot
+        hist_data = [residuals]
+        group_labels = ["Residuals"]  # Names of the dataset
+        bin_size = self.params["bin_size"]
+        fig = ff.create_distplot(
+            hist_data, group_labels, bin_size=[bin_size], show_hist=True, show_rug=False
+        )
+        fig.update_layout(
+            title="Distribution of Residuals",
+            xaxis_title="Residuals",
+            yaxis_title="Density",
+        )
+        figures = [
+            Figure(
+                for_object=self,
+                key=self.key,
+                figure=fig,
+            )
+        ]
+        # Create a scatter plot of actual vs predicted values
+        scatter = go.Scatter(
+            x=y_true.flatten(),
+            y=y_pred.flatten(),
+            mode="markers",
+            name="True vs Predicted",
+            marker=dict(color="blue", opacity=0.5),
+        )
+
+        # Line of perfect prediction
+        max_val = np.nanmax([np.nanmax(y_true), np.nanmax(y_pred)])
+        min_val = np.nanmin([np.nanmin(y_true), np.nanmin(y_pred)])
+        line = go.Scatter(
+            x=[min_val, max_val],
+            y=[min_val, max_val],
+            mode="lines",
+            name="Perfect Fit",
+            line=dict(color="red", dash="dash"),
+        )
+
+        # Layout settings
+        layout = go.Layout(
+            title="True vs. Predicted Values",
+            xaxis_title="True Values",
+            yaxis_title="Predicted Values",
+            showlegend=True,
+        )
+
+        fig = go.Figure(data=[scatter, line], layout=layout)
+
+        figures.append(
+            Figure(
+                for_object=self,
+                key=self.key,
+                figure=fig,
+            )
+        )
+
+        return self.cache_results(
+            figures=figures,
+        )

validmind/tests/model_validation/RougeMetrics.py

@@ -85,7 +85,7 @@ class RougeMetrics(Metric):
             )
 
         y_true = list(itertools.chain.from_iterable(self.inputs.dataset.y))
-        y_pred = self.inputs.dataset.y_pred(self.inputs.model.input_id)
+        y_pred = self.inputs.dataset.y_pred(self.inputs.model)
 
         rouge = Rouge(metrics=r_metrics)
 

validmind/tests/model_validation/RougeMetricsAggregate.py

@@ -81,7 +81,7 @@ class RougeMetricsAggregate(Metric):
             )
 
         y_true = list(itertools.chain.from_iterable(self.inputs.dataset.y))
-        y_pred = self.inputs.dataset.y_pred(self.inputs.model.input_id)
+        y_pred = self.inputs.dataset.y_pred(self.inputs.model)
 
         rouge = Rouge(metrics=r_metrics)
 

validmind/tests/model_validation/SelfCheckNLIScore.py

@@ -45,7 +45,7 @@ class SelfCheckNLIScore(Metric):
     def run(self):
         # Assuming the dataset is structured with generated sentences and reference samples
         y_true = list(itertools.chain.from_iterable(self.inputs.dataset.y))
-        y_pred = self.inputs.dataset.y_pred(self.inputs.model.input_id)
+        y_pred = self.inputs.dataset.y_pred(self.inputs.model)
 
         hallucination_scores = self.compute_hallucination_scores(y_pred, y_true)
 

validmind/tests/model_validation/TokenDisparity.py

@@ -62,7 +62,7 @@ class TokenDisparity(Metric):
 
     def run(self):
         y_true = list(itertools.chain.from_iterable(self.inputs.dataset.y))
-        y_pred = self.inputs.dataset.y_pred(self.inputs.model.input_id)
+        y_pred = self.inputs.dataset.y_pred(self.inputs.model)
 
         df = pd.DataFrame({"reference_column": y_true, "generated_column": y_pred})
 

validmind/tests/model_validation/ToxicityHistogram.py

@@ -57,7 +57,7 @@ class ToxicityHistogram(Metric):
             raise AttributeError("The 'model' attribute is missing.")
 
         y_true = list(itertools.chain.from_iterable(self.inputs.dataset.y))
-        y_pred = self.inputs.dataset.y_pred(self.inputs.model.input_id)
+        y_pred = self.inputs.dataset.y_pred(self.inputs.model)
         input_text = self.inputs.dataset.df[self.inputs.dataset.text_column]
 
         # Ensure consistency in lengths

validmind/tests/model_validation/ToxicityScore.py

@@ -60,7 +60,7 @@ class ToxicityScore(Metric):
             raise AttributeError("The 'model' attribute is missing.")
 
         y_true = list(itertools.chain.from_iterable(self.inputs.dataset.y))
-        y_pred = self.inputs.dataset.y_pred(self.inputs.model.input_id)
+        y_pred = self.inputs.dataset.y_pred(self.inputs.model)
         input_text = self.inputs.dataset.df[self.inputs.dataset.text_column]
 
         # Ensure consistency in lengths

validmind/tests/model_validation/embeddings/ClusterDistribution.py

@@ -59,7 +59,7 @@ class ClusterDistribution(Metric):
     def run(self):
         # run kmeans clustering on embeddings
         kmeans = KMeans(n_clusters=self.params["num_clusters"]).fit(
-            self.inputs.dataset.y_pred(self.inputs.model.input_id)
+            self.inputs.dataset.y_pred(self.inputs.model)
         )
 
         # plot the distribution

validmind/tests/model_validation/embeddings/CosineSimilarityDistribution.py

@@ -57,9 +57,7 @@ class CosineSimilarityDistribution(Metric):
 
     def run(self):
         # Compute cosine similarity
-        similarities = cosine_similarity(
-            self.inputs.dataset.y_pred(self.inputs.model.input_id)
-        )
+        similarities = cosine_similarity(self.inputs.dataset.y_pred(self.inputs.model))
 
         # plot the distribution
         fig = px.histogram(

validmind/tests/model_validation/embeddings/DescriptiveAnalytics.py

@@ -60,7 +60,7 @@ class DescriptiveAnalytics(Metric):
 
     def run(self):
         # Assuming y_pred returns a 2D array of embeddings [samples, features]
-        preds = self.inputs.dataset.y_pred(self.inputs.model.input_id)
+        preds = self.inputs.dataset.y_pred(self.inputs.model)
 
         # Calculate statistics across the embedding dimensions, not across all embeddings
         means = np.mean(preds, axis=0)  # Mean of each feature across all samples

validmind/tests/model_validation/embeddings/EmbeddingsVisualization2D.py

@@ -67,7 +67,7 @@ class EmbeddingsVisualization2D(Metric):
             )
 
         # use TSNE to reduce dimensionality of embeddings
-        num_samples = len(self.inputs.dataset.y_pred(self.inputs.model.input_id))
+        num_samples = len(self.inputs.dataset.y_pred(self.inputs.model))
 
         if self.params["perplexity"] >= num_samples:
             perplexity = num_samples - 1

validmind/tests/model_validation/sklearn/ClassifierPerformance.py

@@ -5,17 +5,17 @@
 from dataclasses import dataclass
 
 from numpy import unique
-from sklearn import metrics, preprocessing
+from sklearn.metrics import classification_report, roc_auc_score
+from sklearn.preprocessing import LabelBinarizer
 
 from validmind.vm_models import Metric, ResultSummary, ResultTable, ResultTableMetadata
 
 
 def multiclass_roc_auc_score(y_test, y_pred, average="macro"):
-    lb = preprocessing.LabelBinarizer()
+    lb = LabelBinarizer()
     lb.fit(y_test)
-    y_test = lb.transform(y_test)
-    y_pred = lb.transform(y_pred)
-    return metrics.roc_auc_score(y_test, y_pred, average=average)
+
+    return roc_auc_score(lb.transform(y_test), lb.transform(y_pred), average=average)
 
 
 @dataclass
@@ -73,7 +73,7 @@ class ClassifierPerformance(Metric):
         When building a multi-class summary we need to calculate weighted average,
         macro average and per class metrics.
         """
-        classes = {str(i) for i in unique(self.y_true())}
+        classes = {str(i) for i in unique(self.inputs.dataset.y)}
         pr_f1_table = [
             {
                 "Class": class_name,
@@ -121,19 +121,16 @@ class ClassifierPerformance(Metric):
             ]
         )
 
-    def y_true(self):
-        return self.inputs.dataset.y
-
-    def y_pred(self):
-        return self.inputs.dataset.y_pred(model_id=self.inputs.model.input_id)
-
     def run(self):
-        y_true = self.y_true()
-        class_pred = self.y_pred()
-
-        report = metrics.classification_report(
-            y_true, class_pred, output_dict=True, zero_division=0
+        report = classification_report(
+            self.inputs.dataset.y,
+            self.inputs.dataset.y_pred(self.inputs.model),
+            output_dict=True,
+            zero_division=0,
+        )
+        report["roc_auc"] = multiclass_roc_auc_score(
+            self.inputs.dataset.y,
+            self.inputs.dataset.y_pred(self.inputs.model),
         )
-        report["roc_auc"] = multiclass_roc_auc_score(y_true, class_pred)
 
         return self.cache_results(report)

validmind/tests/model_validation/sklearn/ClusterCosineSimilarity.py

@@ -67,7 +67,7 @@ class ClusterCosineSimilarity(Metric):
 
     def run(self):
         y_true_train = self.inputs.dataset.y
-        y_pred_train = self.inputs.dataset.y_pred(self.inputs.model.input_id)
+        y_pred_train = self.inputs.dataset.y_pred(self.inputs.model)
         y_true_train = y_true_train.astype(y_pred_train.dtype).flatten()
         num_clusters = len(np.unique(y_pred_train))
         # Calculate cosine similarity for each cluster

validmind/tests/model_validation/sklearn/ClusterPerformance.py

@@ -101,11 +101,11 @@ class ClusterPerformance(Metric):
 
     def run(self):
         y_true_train = self.inputs.datasets[0].y
-        class_pred_train = self.inputs.datasets[0].y_pred(self.inputs.model.input_id)
+        class_pred_train = self.inputs.datasets[0].y_pred(self.inputs.model)
         y_true_train = y_true_train.astype(class_pred_train.dtype)
 
         y_true_test = self.inputs.datasets[1].y
-        class_pred_test = self.inputs.datasets[1].y_pred(self.inputs.model.input_id)
+        class_pred_test = self.inputs.datasets[1].y_pred(self.inputs.model)
         y_true_test = y_true_test.astype(class_pred_test.dtype)
 
         samples = ["train", "test"]

validmind/tests/model_validation/sklearn/ConfusionMatrix.py

@@ -72,15 +72,33 @@ class ConfusionMatrix(Metric):
         labels.sort()
         labels = np.array(labels).T.tolist()
 
-        class_pred = self.inputs.dataset.y_pred(model_id=self.inputs.model.input_id)
-        y_true = y_true.astype(class_pred.dtype)
-        cm = metrics.confusion_matrix(y_true, class_pred, labels=labels)
+        y_pred = self.inputs.dataset.y_pred(self.inputs.model)
+        y_true = y_true.astype(y_pred.dtype)
+
+        cm = metrics.confusion_matrix(y_true, y_pred, labels=labels)
+
+        text = None
+        if len(labels) == 2:
+            tn, fp, fn, tp = cm.ravel()
+
+            # Custom text to display on the heatmap cells
+            text = [
+                [
+                    f"<b>True Negatives (TN)</b><br />{tn}",
+                    f"<b>False Positives (FP)</b><br />{fp}",
+                ],
+                [
+                    f"<b>False Negatives (FN)</b><br />{fn}",
+                    f"<b>True Positives (TP)</b><br />{tp}",
+                ],
+            ]
 
         fig = ff.create_annotated_heatmap(
             z=cm,
             colorscale="Blues",
             x=labels,
             y=labels,
+            annotation_text=text,
         )
 
         fig["data"][0][

validmind/tests/model_validation/sklearn/MinimumAccuracy.py

@@ -96,7 +96,7 @@ class MinimumAccuracy(ThresholdTest):
 
     def run(self):
         y_true = self.inputs.dataset.y
-        class_pred = self.inputs.dataset.y_pred(self.inputs.model.input_id)
+        class_pred = self.inputs.dataset.y_pred(self.inputs.model)
         y_true = y_true.astype(class_pred.dtype)
 
         accuracy_score = metrics.accuracy_score(y_true, class_pred)

validmind/tests/model_validation/sklearn/MinimumF1Score.py

@@ -97,7 +97,7 @@ class MinimumF1Score(ThresholdTest):
 
     def run(self):
         y_true = self.inputs.dataset.y
-        class_pred = self.inputs.dataset.y_pred(self.inputs.model.input_id)
+        class_pred = self.inputs.dataset.y_pred(self.inputs.model)
         y_true = y_true.astype(class_pred.dtype)
 
         if len(unique(y_true)) > 2:

validmind/tests/model_validation/sklearn/MinimumROCAUCScore.py

@@ -101,7 +101,7 @@ class MinimumROCAUCScore(ThresholdTest):
 
     def run(self):
         y_true = self.inputs.dataset.y
-        class_pred = self.inputs.dataset.y_pred(self.inputs.model.input_id)
+        class_pred = self.inputs.dataset.y_pred(self.inputs.model)
         y_true = y_true.astype(class_pred.dtype)
         roc_auc = self.multiclass_roc_auc_score(y_true, class_pred)
 

validmind/tests/model_validation/sklearn/ModelsPerformanceComparison.py

@@ -5,7 +5,7 @@
 from dataclasses import dataclass
 
 from numpy import unique
-from sklearn import metrics
+from sklearn.metrics import classification_report
 
 from validmind.errors import SkipTestError
 from validmind.vm_models import ResultSummary, ResultTable, ResultTableMetadata
@@ -129,8 +129,9 @@ class ModelsPerformanceComparison(ClassifierPerformance):
         results = {}
         for idx, model in enumerate(all_models):
             y_true = self.inputs.dataset.y
-            class_pred = self.inputs.dataset.y_pred(model.input_id)
-            report = metrics.classification_report(y_true, class_pred, output_dict=True)
-            report["roc_auc"] = multiclass_roc_auc_score(y_true, class_pred)
+            y_pred = self.inputs.dataset.y_pred(model)
+            report = classification_report(y_true, y_pred, output_dict=True)
+            report["roc_auc"] = multiclass_roc_auc_score(y_true, y_pred)
             results["model_" + str(idx)] = report
+
         return self.cache_results(results)

validmind/tests/model_validation/sklearn/OverfitDiagnosis.py

@@ -119,12 +119,12 @@ class OverfitDiagnosis(ThresholdTest):
 
         # Add prediction column in the training dataset
         train_df = self.inputs.datasets[0].df.copy()
-        train_class_pred = self.inputs.datasets[0].y_pred(self.inputs.model.input_id)
+        train_class_pred = self.inputs.datasets[0].y_pred(self.inputs.model)
         train_df[prediction_column] = train_class_pred
 
         # Add prediction column in the test dataset
         test_df = self.inputs.datasets[1].df.copy()
-        test_class_pred = self.inputs.datasets[1].y_pred(self.inputs.model.input_id)
+        test_class_pred = self.inputs.datasets[1].y_pred(self.inputs.model)
         test_df[prediction_column] = test_class_pred
 
         test_results = []

validmind/tests/model_validation/sklearn/ROCCurve.py

@@ -73,15 +73,8 @@ class ROCCurve(Metric):
         if self.inputs.model.model_library() == "FoundationModel":
             raise SkipTestError("Skipping ROCCurve for Foundation models")
 
-        # Extract the actual model
-        model = (
-            self.inputs.model[0]
-            if isinstance(self.inputs.model, list)
-            else self.inputs.model
-        )
-
         y_true = self.inputs.dataset.y
-        y_pred = model.predict_proba(self.inputs.dataset.x)
+        y_prob = self.inputs.dataset.y_prob(self.inputs.model)
 
         # ROC curve is only supported for binary classification
         if len(np.unique(y_true)) > 2:
@@ -89,14 +82,15 @@ class ROCCurve(Metric):
                 "ROC Curve is only supported for binary classification models"
             )
 
-        y_true = y_true.astype(y_pred.dtype).flatten()
-        assert np.all((y_pred >= 0) & (y_pred <= 1)), "Invalid probabilities in y_pred."
+        y_true = y_true.astype(y_prob.dtype).flatten()
+        assert np.all((y_prob >= 0) & (y_prob <= 1)), "Invalid probabilities in y_prob."
+
+        fpr, tpr, roc_thresholds = roc_curve(y_true, y_prob, drop_intermediate=False)
 
-        fpr, tpr, roc_thresholds = roc_curve(y_true, y_pred, drop_intermediate=False)
         # Remove Inf values from roc_thresholds
         valid_thresholds_mask = np.isfinite(roc_thresholds)
         roc_thresholds = roc_thresholds[valid_thresholds_mask]
-        auc = roc_auc_score(y_true, y_pred)
+        auc = roc_auc_score(y_true, y_prob)
 
         trace0 = go.Scatter(
             x=fpr,

validmind/tests/model_validation/sklearn/RegressionErrors.py

@@ -130,11 +130,11 @@ class RegressionErrors(Metric):
 
     def run(self):
         y_train_true = self.inputs.datasets[0].y
-        y_train_pred = self.inputs.datasets[0].y_pred(self.inputs.model.input_id)
+        y_train_pred = self.inputs.datasets[0].y_pred(self.inputs.model)
         y_train_true = y_train_true.astype(y_train_pred.dtype)
 
         y_test_true = self.inputs.datasets[1].y
-        y_test_pred = self.inputs.datasets[1].y_pred(self.inputs.model.input_id)
+        y_test_pred = self.inputs.datasets[1].y_pred(self.inputs.model)
         y_test_true = y_test_true.astype(y_test_pred.dtype)
 
         results = self.regression_errors(

validmind/tests/model_validation/sklearn/RegressionModelsPerformanceComparison.py

@@ -6,7 +6,7 @@ import re
 from dataclasses import dataclass
 
 import numpy as np
-from sklearn import metrics
+from sklearn.metrics import mean_absolute_error, mean_squared_error
 
 from validmind.errors import SkipTestError
 from validmind.vm_models import Metric, ResultSummary, ResultTable, ResultTableMetadata
@@ -67,12 +67,12 @@ class RegressionModelsPerformanceComparison(Metric):
     }
 
     def regression_errors(self, y_true_test, y_pred_test):
-        mae_test = metrics.mean_absolute_error(y_true_test, y_pred_test)
+        mae_test = mean_absolute_error(y_true_test, y_pred_test)
 
         results = {}
         results["Mean Absolute Error (MAE)"] = mae_test
 
-        mse_test = metrics.mean_squared_error(y_true_test, y_pred_test)
+        mse_test = mean_squared_error(y_true_test, y_pred_test)
         results["Mean Squared Error (MSE)"] = mse_test
         results["Root Mean Squared Error (RMSE)"] = np.sqrt(mse_test)
 
@@ -121,12 +121,14 @@ class RegressionModelsPerformanceComparison(Metric):
 
         if self.inputs.models is not None:
             all_models.extend(self.inputs.models)
+
         results = {}
 
         for idx, model in enumerate(all_models):
             result = self.regression_errors(
                 y_true_test=self.inputs.dataset.y,
-                y_pred_test=self.inputs.dataset.y_pred(model.input_id),
+                y_pred_test=self.inputs.dataset.y_pred(model),
             )
             results["model_" + str(idx)] = result
+
         return self.cache_results(results)