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

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

@@ -76,7 +76,6 @@ class RougeMetrics(Metric):
         if r_metrics is None:
             raise ValueError("rouge_metrics must be provided in params")
 
-        # With all
         if not (
             set(self.default_params.get("rouge_metrics")).intersection(r_metrics)
             == set(r_metrics)
@@ -86,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)
 
@@ -97,12 +96,13 @@ class RougeMetrics(Metric):
 
         metrics_df = pd.DataFrame(score_list)
         figures = []
+
         for m in metrics_df.columns:
             df_scores = pd.DataFrame(metrics_df[m].tolist())
             # Visualization part
             fig = go.Figure()
 
-            # Adding the line plots
+            # Adding the line plots for precision, recall, and F1-score with lines and markers
             fig.add_trace(
                 go.Scatter(
                     x=df_scores.index,
@@ -129,11 +129,13 @@ class RougeMetrics(Metric):
             )
 
             fig.update_layout(
-                title="ROUGE Scores for Each Row",
+                title=f"ROUGE Scores for {m}",
                 xaxis_title="Row Index",
                 yaxis_title="Score",
             )
-            k = m.replace("-", "")
+
+            # Ensure a unique key for each metric
+            k = f"{m.replace('-', '')}_{len(figures)}"
             figures.append(
                 Figure(
                     for_object=self,

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

@@ -0,0 +1,112 @@
+# 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
+
+import itertools
+from dataclasses import dataclass
+
+import pandas as pd
+import plotly.graph_objects as go
+import torch
+from selfcheckgpt.modeling_selfcheck import SelfCheckNLI
+from tqdm import tqdm
+
+from validmind.vm_models import Figure, Metric
+
+
+@dataclass
+class SelfCheckNLIScore(Metric):
+    """
+    Evaluates text generation models' performance by quantifying the level of hallucination in generated texts compared to reference texts.
+
+    **Purpose**: The HallucinationScore metric is designed to assess the factual accuracy and reliability of text generated by models, focusing on the detection and quantification of hallucinations—instances where generated content deviates from factual or expected outputs. By comparing generated texts against reference texts, this metric highlights discrepancies indicative of hallucinations, offering insights into the model's ability to produce contextually and factually coherent content.
+
+    **Test Mechanism**: To compute the HallucinationScore, the metric employs a comparison between the generated texts (model predictions) and the provided reference texts (true values). Using the SelfCheckNLI model, it evaluates each generated text's level of factual congruence with the reference, assigning a hallucination score based on the semantic coherence and factual accuracy. The scores for each text instance are then visualized in a line plot, allowing for the examination of hallucination trends across the dataset.
+
+    **Signs of High Risk**:
+    - High hallucination scores across a significant portion of the dataset, indicating a prevalence of factually inaccurate or irrelevant content generation.
+    - Patterns of consistent hallucination in specific contexts or subjects, suggesting gaps in the model's understanding or knowledge.
+    - Sharp fluctuations in hallucination scores, which may reveal inconsistencies in the model's performance or sensitivity to certain types of input.
+
+    **Strengths**:
+    - Directly addresses the critical aspect of factual accuracy in generated text, beyond mere linguistic or stylistic coherence.
+    - Provides a granular, instance-by-instance analysis of model performance, allowing for targeted improvements and diagnostics.
+    - Facilitates a deeper understanding of a model's capabilities and limitations in producing reliable and accurate content.
+
+    **Limitations**:
+    - Reliance on the SelfCheckNLI model means the accuracy and effectiveness of the HallucinationScore are contingent upon the performance and suitability of the underlying NLI model.
+    - May not fully capture the subtleties of certain factual inaccuracies or the contextual relevance of reference texts, especially in complex or nuanced domains.
+    - Potentially resource-intensive, given the computational demands of running advanced NLI models for large datasets.
+    """
+
+    name = "self_check_nli_score"
+    required_inputs = ["model", "dataset"]
+
+    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)
+
+        hallucination_scores = self.compute_hallucination_scores(y_pred, y_true)
+
+        # Visualization of scores
+        figures = self.visualize_scores(hallucination_scores)
+
+        return self.cache_results(figures=figures)
+
+    def compute_hallucination_scores(self, predictions, references):
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        selfcheck_nli = SelfCheckNLI(device=device)
+        hallucination_scores = []
+
+        print("Starting hallucination score computation...")
+
+        for index, (sentences, samples) in enumerate(
+            tqdm(zip(predictions, references), total=len(predictions))
+        ):
+            sent_scores_nli = selfcheck_nli.predict(
+                sentences=sentences, sampled_passages=samples
+            )
+
+            # Compute the mean of the hallucination scores for this row
+            average_score = sent_scores_nli.mean()
+            hallucination_scores.append(average_score)
+
+            # Print a progress update for each row
+            print(
+                f"Row {index + 1}/{len(predictions)}: Average hallucination score: {average_score}"
+            )
+
+        print("Completed hallucination score computation.")
+
+        return hallucination_scores
+
+    def visualize_scores(self, scores):
+        scores_df = pd.DataFrame(scores, columns=["Hallucination Score"])
+
+        fig = go.Figure()
+        fig.add_trace(
+            go.Scatter(
+                x=scores_df.index,
+                y=scores_df["Hallucination Score"],
+                mode="lines+markers",
+                name="Hallucination Score",
+            )
+        )
+
+        fig.update_layout(
+            title="Hallucination Scores Across Text Instances",
+            xaxis_title="Text Instance Index",
+            yaxis_title="Hallucination Score",
+        )
+
+        # Wrapping the plotly figure for compatibility with your framework might be needed
+        figures = [
+            Figure(
+                for_object=self,
+                key=self.key,
+                figure=fig,
+            )
+        ]
+
+        return figures

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

@@ -59,30 +59,25 @@ class DescriptiveAnalytics(Metric):
     }
 
     def run(self):
-        mean = np.mean(self.inputs.dataset.y_pred(self.inputs.model.input_id))
-        median = np.median(self.inputs.dataset.y_pred(self.inputs.model.input_id))
-        std = np.std(self.inputs.dataset.y_pred(self.inputs.model.input_id))
+        # Assuming y_pred returns a 2D array of embeddings [samples, features]
+        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
+        medians = np.median(preds, axis=0)  # Median of each feature across all samples
+        stds = np.std(preds, axis=0)  # Std. dev. of each feature across all samples
+
+        # Plot histograms of the calculated statistics
+        mean_fig = px.histogram(x=means, title="Distribution of Embedding Means")
+        median_fig = px.histogram(x=medians, title="Distribution of Embedding Medians")
+        std_fig = px.histogram(
+            x=stds, title="Distribution of Embedding Standard Deviations"
+        )
 
         return self.cache_results(
             figures=[
-                Figure(
-                    for_object=self,
-                    key=self.key,
-                    figure=px.histogram(mean, title="Distribution of Embedding Means"),
-                ),
-                Figure(
-                    for_object=self,
-                    key=self.key,
-                    figure=px.histogram(
-                        median, title="Distribution of Embedding Medians"
-                    ),
-                ),
-                Figure(
-                    for_object=self,
-                    key=self.key,
-                    figure=px.histogram(
-                        std, title="Distribution of Embedding Standard Deviations"
-                    ),
-                ),
+                Figure(for_object=self, key=f"{self.key}_mean", figure=mean_fig),
+                Figure(for_object=self, key=f"{self.key}_median", figure=median_fig),
+                Figure(for_object=self, key=f"{self.key}_std", figure=std_fig),
             ],
         )

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,17 +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)
-        report["roc_auc"] = multiclass_roc_auc_score(y_true, class_pred)
+        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),
+        )
 
         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(