From 53965c0c606e35c6b309b9057dbeacfbcaca56f6 Mon Sep 17 00:00:00 2001
From: Juan <juan@validmind.ai>
Date: Tue, 18 Feb 2025 13:29:35 +0100
Subject: [PATCH 01/10] Add input ids to NLP tests

---
 .../how_to/understand_utilize_rawdata.ipynb   | 158 +++++++++++++++++-
 scripts/bulk_ai_test_updates.py               |   9 +-
 .../tests/data_validation/nlp/CommonWords.py  |   2 +-
 .../tests/data_validation/nlp/Hashtags.py     |   2 +-
 .../data_validation/nlp/LanguageDetection.py  |   2 +-
 .../tests/data_validation/nlp/Mentions.py     |   2 +-
 .../nlp/PolarityAndSubjectivity.py            |   6 +-
 .../tests/data_validation/nlp/Punctuations.py |   2 +-
 .../tests/data_validation/nlp/Sentiment.py    |   4 +-
 .../data_validation/nlp/TextDescription.py    |   2 +-
 .../tests/data_validation/nlp/Toxicity.py     |   2 +-
 .../sklearn/ConfusionMatrix.py                |  25 ++-
 12 files changed, 197 insertions(+), 19 deletions(-)

diff --git a/notebooks/how_to/understand_utilize_rawdata.ipynb b/notebooks/how_to/understand_utilize_rawdata.ipynb
index 7354eae85..b5c1dd6f6 100644
--- a/notebooks/how_to/understand_utilize_rawdata.ipynb
+++ b/notebooks/how_to/understand_utilize_rawdata.ipynb
@@ -31,7 +31,7 @@
     "  - [Pearson Correlation Matrix](#toc2_2_)    \n",
     "  - [Precision-Recall Curve](#toc2_3_)    \n",
     "  - [Using `RawData` in custom tests](#toc2_4_)    \n",
-    "\n",
+    "  - [Using `RawData` in comparison tests](#toc2_5_)    \n",
     ":::\n",
     "<!-- jn-toc-notebook-config\n",
     "\tnumbering=false\n",
@@ -213,7 +213,8 @@
     "  - [Using `RawData` from the ROC Curve Test](#toc2_1_)    \n",
     "  - [Pearson Correlation Matrix](#toc2_2_)    \n",
     "  - [Precision-Recall Curve](#toc2_3_)    \n",
-    "  - [Using `RawData` in custom tests](#toc2_4_)    "
+    "  - [Using `RawData` in custom tests](#toc2_4_)  \n",
+    "  - [Using `RawData` in comparison tests](#toc2_5_)  "
    ]
   },
   {
@@ -553,17 +554,164 @@
     "    generate_description=False,\n",
     ")"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "53084493",
+   "metadata": {},
+   "source": [
+    "<a id='toc2_5_'></a>\n",
+    "\n",
+    "### Using `RawData` in comparison tests\n",
+    "\n",
+    "TBC."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "bcbbe9f4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def cm_table(result: TestResult):\n",
+    "    # For individual results\n",
+    "    if not isinstance(result.raw_data.confusion_matrix, list):\n",
+    "        # Extract values from single confusion matrix\n",
+    "        cm = result.raw_data.confusion_matrix\n",
+    "        tn, fp = cm[0, 0], cm[0, 1]\n",
+    "        fn, tp = cm[1, 0], cm[1, 1]\n",
+    "        \n",
+    "        # Create DataFrame for individual matrix\n",
+    "        cm_df = pd.DataFrame({\n",
+    "            'TN': [tn],\n",
+    "            'FP': [fp],\n",
+    "            'FN': [fn],\n",
+    "            'TP': [tp]\n",
+    "        })\n",
+    "        \n",
+    "        # Add individual table\n",
+    "        result.add_table(cm_df, title=\"Confusion Matrix\")\n",
+    "        \n",
+    "    # For comparison results\n",
+    "    else:\n",
+    "        cms = result.raw_data.confusion_matrix\n",
+    "        cm1, cm2 = cms[0], cms[1]\n",
+    "        \n",
+    "        # Create individual results table\n",
+    "        rows = []\n",
+    "        for i, cm in enumerate(cms):\n",
+    "            rows.append({\n",
+    "                'dataset': result.raw_data.dataset[i],\n",
+    "                'model': result.raw_data.model[i],\n",
+    "                'TN': cm[0, 0],\n",
+    "                'FP': cm[0, 1],\n",
+    "                'FN': cm[1, 0],\n",
+    "                'TP': cm[1, 1]\n",
+    "            })\n",
+    "        individual_df = pd.DataFrame(rows)\n",
+    "        \n",
+    "        # Calculate percentage differences\n",
+    "        diff_df = pd.DataFrame({\n",
+    "            'TN_drift (%)': [(cm2[0, 0] - cm1[0, 0]) / cm1[0, 0] * 100],\n",
+    "            'FP_drift (%)': [(cm2[0, 1] - cm1[0, 1]) / cm1[0, 1] * 100],\n",
+    "            'FN_drift (%)': [(cm2[1, 0] - cm1[1, 0]) / cm1[1, 0] * 100],\n",
+    "            'TP_drift (%)': [(cm2[1, 1] - cm1[1, 1]) / cm1[1, 1] * 100]\n",
+    "        }).round(2)\n",
+    "        \n",
+    "        # Add both tables\n",
+    "        result.add_table(individual_df, title=\"Individual Confusion Matrices\")\n",
+    "        result.add_table(diff_df, title=\"Confusion Matrix Drift\")\n",
+    "        \n",
+    "    return result"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cf3c47fe",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from validmind.tests import run_test\n",
+    "\n",
+    "result_cm = run_test(\n",
+    "    \"validmind.model_validation.sklearn.ConfusionMatrix\",\n",
+    "    inputs={\n",
+    "        \"dataset\": vm_test_ds,\n",
+    "        \"model\": vm_model,\n",
+    "    },\n",
+    "    post_process_fn=cm_table,\n",
+    "    generate_description=False,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6a1b4388",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "result_cm = run_test(\n",
+    "    \"validmind.model_validation.sklearn.ConfusionMatrix\",\n",
+    "    input_grid={\n",
+    "        \"dataset\": [vm_test_ds, vm_train_ds],\n",
+    "        \"model\": [vm_model]\n",
+    "    },\n",
+    "    post_process_fn=cm_table,\n",
+    "    generate_description=False,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "012ec495",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "result_cm.raw_data.confusion_matrix"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4626bb96",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "result_cm.raw_data.model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "07825ec9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "result_cm.raw_data.dataset"
+   ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "ValidMind Library",
    "language": "python",
-   "name": "python3"
+   "name": "validmind"
   },
   "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
    "name": "python",
-   "version": "3.10.13"
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.15"
   }
  },
  "nbformat": 4,
diff --git a/scripts/bulk_ai_test_updates.py b/scripts/bulk_ai_test_updates.py
index 33506a1b6..d44a7136f 100644
--- a/scripts/bulk_ai_test_updates.py
+++ b/scripts/bulk_ai_test_updates.py
@@ -118,7 +118,7 @@ def list_to_str(lst):
 
 **Purpose**:
 
-The Feature Drift test aims to evaluate how much the distribution of features has shifted over time between two datasets, typically training and monitoring datasets. It uses the Population Stability Index (PSI) to quantify this change, providing insights into the model’s robustness and the necessity for retraining or feature engineering.
+The Feature Drift test aims to evaluate how much the distribution of features has shifted over time between two datasets, typically training and monitoring datasets. It uses the Population Stability Index (PSI) to quantify this change, providing insights into the model's robustness and the necessity for retraining or feature engineering.
 
 **Test Mechanism**:
 
@@ -181,6 +181,11 @@ def list_to_str(lst):
 Its a class that can be initialized with any number of any type of objects using a key-value like interface where the key in the constructor is the name of the object and the value is the object itself.
 It should only be used to store data that is not already returned as part of the test result (i.e. in a table) but could be useful to re-generate any of the test result objects (tables, figures).
 
+When adding raw data, you should always include:
+- If the test has access to a model parameter (VMModel), include its input_id as model=model.input_id
+- If the test has access to a dataset parameter (VMDataset), include its input_id as dataset=dataset.input_id
+Only include these if they are available in the test function parameters - don't force both if only one is accessible.
+
 You will be provided with the source code for a "test" that is run against an ML model or dataset.
 You will analyze the code to determine the details and implementation of the test.
 Then you will use the below example to implement changes to the test to make it use the new raw data mechanism offered by the ValidMind SDK.
@@ -228,7 +233,7 @@ def ExampleConfusionMatrix(model: VMModel, dataset: VMDataset):
     fig = ff.create_annotated_heatmap()
     ..
 
-    return fig, RawData(confusion_matrix=cm)
+    return fig, RawData(confusion_matrix=cm, model=model.input_id, dataset=dataset.input_id)
 ```
 
 Notice that the test now returns a tuple of the figure and the raw data.
diff --git a/validmind/tests/data_validation/nlp/CommonWords.py b/validmind/tests/data_validation/nlp/CommonWords.py
index 21f21bd40..52fe819f8 100644
--- a/validmind/tests/data_validation/nlp/CommonWords.py
+++ b/validmind/tests/data_validation/nlp/CommonWords.py
@@ -94,4 +94,4 @@ def create_corpus(df, text_column):
         xaxis_tickangle=-45,
     )
 
-    return fig, RawData(words=x, frequencies=y)
+    return fig, RawData(words=x, frequencies=y, dataset=dataset.input_id)
diff --git a/validmind/tests/data_validation/nlp/Hashtags.py b/validmind/tests/data_validation/nlp/Hashtags.py
index 4d292ee12..54d4d6c64 100644
--- a/validmind/tests/data_validation/nlp/Hashtags.py
+++ b/validmind/tests/data_validation/nlp/Hashtags.py
@@ -76,4 +76,4 @@ def Hashtags(dataset: VMDataset, top_hashtags: int = 25):
         xaxis_tickangle=-45,
     )
 
-    return fig, RawData(top_hashtag_counts=top_hashtag_counts)
+    return fig, RawData(top_hashtag_counts=top_hashtag_counts, dataset=dataset.input_id)
diff --git a/validmind/tests/data_validation/nlp/LanguageDetection.py b/validmind/tests/data_validation/nlp/LanguageDetection.py
index a4300288f..9b6608240 100644
--- a/validmind/tests/data_validation/nlp/LanguageDetection.py
+++ b/validmind/tests/data_validation/nlp/LanguageDetection.py
@@ -71,5 +71,5 @@ def detect_language(text):
             title="Language Distribution",
             labels={"x": "Language Codes"},
         ),
-        RawData(detected_languages=languages),
+        RawData(detected_languages=languages, dataset=dataset.input_id),
     )
diff --git a/validmind/tests/data_validation/nlp/Mentions.py b/validmind/tests/data_validation/nlp/Mentions.py
index 8eb63855f..9d69a05fc 100644
--- a/validmind/tests/data_validation/nlp/Mentions.py
+++ b/validmind/tests/data_validation/nlp/Mentions.py
@@ -82,5 +82,5 @@ def Mentions(dataset: VMDataset, top_mentions: int = 25):
             values="Percentage",
             title="Tree of Mentions",
         ),
-        RawData(mention_counts=mention_counts),
+        RawData(mention_counts=mention_counts, dataset=dataset.input_id),
     )
diff --git a/validmind/tests/data_validation/nlp/PolarityAndSubjectivity.py b/validmind/tests/data_validation/nlp/PolarityAndSubjectivity.py
index 596f2c386..da4f816f6 100644
--- a/validmind/tests/data_validation/nlp/PolarityAndSubjectivity.py
+++ b/validmind/tests/data_validation/nlp/PolarityAndSubjectivity.py
@@ -144,4 +144,8 @@ def PolarityAndSubjectivity(dataset, threshold_subjectivity=0.5, threshold_polar
 
     statistics_tables = {"Quadrant Distribution": quadrant_df, "Statistics": stats_df}
 
-    return fig, statistics_tables, RawData(sentiment_data=data)
+    return (
+        fig,
+        statistics_tables,
+        RawData(sentiment_data=data, dataset=dataset.input_id),
+    )
diff --git a/validmind/tests/data_validation/nlp/Punctuations.py b/validmind/tests/data_validation/nlp/Punctuations.py
index f846aa74c..352b28a0e 100644
--- a/validmind/tests/data_validation/nlp/Punctuations.py
+++ b/validmind/tests/data_validation/nlp/Punctuations.py
@@ -65,7 +65,7 @@ def Punctuations(dataset, count_mode="token"):
     punctuation_counts = _count_punctuations(corpus, count_mode)
     fig = _create_punctuation_plot(punctuation_counts)
 
-    return fig, RawData(punctuation_counts=punctuation_counts)
+    return fig, RawData(punctuation_counts=punctuation_counts, dataset=dataset.input_id)
 
 
 def _create_punctuation_plot(punctuation_counts):
diff --git a/validmind/tests/data_validation/nlp/Sentiment.py b/validmind/tests/data_validation/nlp/Sentiment.py
index 6abb052e2..1c8ef1932 100644
--- a/validmind/tests/data_validation/nlp/Sentiment.py
+++ b/validmind/tests/data_validation/nlp/Sentiment.py
@@ -77,4 +77,6 @@ def get_vader_sentiment(text):
 
     plt.close("all")
 
-    return fig, RawData(sentiment_scores=vader_sentiment.tolist())
+    return fig, RawData(
+        sentiment_scores=vader_sentiment.tolist(), dataset=dataset.input_id
+    )
diff --git a/validmind/tests/data_validation/nlp/TextDescription.py b/validmind/tests/data_validation/nlp/TextDescription.py
index 352aba282..69f3c17eb 100644
--- a/validmind/tests/data_validation/nlp/TextDescription.py
+++ b/validmind/tests/data_validation/nlp/TextDescription.py
@@ -173,4 +173,4 @@ def TextDescription(
             )
         )
 
-    return (*figures, RawData(metrics_dataframe=metrics_df))
+    return (*figures, RawData(metrics_dataframe=metrics_df, dataset=dataset.input_id))
diff --git a/validmind/tests/data_validation/nlp/Toxicity.py b/validmind/tests/data_validation/nlp/Toxicity.py
index ebb8fa92a..cae2a2ff3 100644
--- a/validmind/tests/data_validation/nlp/Toxicity.py
+++ b/validmind/tests/data_validation/nlp/Toxicity.py
@@ -73,4 +73,4 @@ def Toxicity(dataset):
 
     plt.close()
 
-    return fig, RawData(toxicity_scores=toxicity_scores)
+    return fig, RawData(toxicity_scores=toxicity_scores, dataset=dataset.input_id)
diff --git a/validmind/tests/model_validation/sklearn/ConfusionMatrix.py b/validmind/tests/model_validation/sklearn/ConfusionMatrix.py
index 050b70f54..81017c87b 100644
--- a/validmind/tests/model_validation/sklearn/ConfusionMatrix.py
+++ b/validmind/tests/model_validation/sklearn/ConfusionMatrix.py
@@ -19,7 +19,11 @@
     "visualization",
 )
 @tasks("classification", "text_classification")
-def ConfusionMatrix(dataset: VMDataset, model: VMModel):
+def ConfusionMatrix(
+    dataset: VMDataset,
+    model: VMModel,
+    threshold: float = 0.5,
+):
     """
     Evaluates and visually represents the classification ML model's predictive performance using a Confusion Matrix
     heatmap.
@@ -66,7 +70,17 @@ def ConfusionMatrix(dataset: VMDataset, model: VMModel):
     - Risks of misinterpretation exist because the matrix doesn't directly provide precision, recall, or F1-score data.
     These metrics have to be computed separately.
     """
-    y_pred = dataset.y_pred(model)
+    # Get predictions using threshold for binary classification if possible
+    if hasattr(model.model, "predict_proba"):
+        y_prob = dataset.y_prob(model)
+        # Handle both 1D and 2D probability arrays
+        if y_prob.ndim == 2:
+            y_pred = (y_prob[:, 1] > threshold).astype(int)
+        else:
+            y_pred = (y_prob > threshold).astype(int)
+    else:
+        y_pred = dataset.y_pred(model)
+
     y_true = dataset.y.astype(y_pred.dtype)
 
     labels = np.unique(y_true)
@@ -119,4 +133,9 @@ def ConfusionMatrix(dataset: VMDataset, model: VMModel):
         font=dict(size=14),
     )
 
-    return fig, RawData(confusion_matrix=cm)
+    return fig, RawData(
+        confusion_matrix=cm,
+        threshold=threshold,
+        dataset=dataset.input_id,
+        model=model.input_id,
+    )

From 75f000c667b9c065b023169703da03c36c92c689 Mon Sep 17 00:00:00 2001
From: Juan <juan@validmind.ai>
Date: Tue, 18 Feb 2025 14:28:01 +0100
Subject: [PATCH 02/10] Update RawData for data validation tests

---
 .../tests/data_validation/ACFandPACFPlot.py   |  5 +++-
 validmind/tests/data_validation/AutoMA.py     |  2 +-
 .../data_validation/BivariateScatterPlots.py  |  6 ++++-
 validmind/tests/data_validation/BoxPierce.py  |  4 +++-
 .../tests/data_validation/ClassImbalance.py   |  2 +-
 .../data_validation/DatasetDescription.py     |  2 +-
 .../tests/data_validation/DickeyFullerGLS.py  |  2 +-
 .../FeatureTargetCorrelationPlot.py           |  2 +-
 .../tests/data_validation/HighCardinality.py  |  6 ++++-
 .../data_validation/HighPearsonCorrelation.py |  2 +-
 .../data_validation/IQROutliersBarPlot.py     |  8 ++++---
 .../tests/data_validation/IQROutliersTable.py |  7 ++++--
 .../IsolationForestOutliers.py                |  9 +++----
 validmind/tests/data_validation/JarqueBera.py |  4 ++--
 validmind/tests/data_validation/LJungBox.py   |  4 ++--
 .../LaggedCorrelationHeatmap.py               |  2 +-
 .../tests/data_validation/MissingValues.py    | 24 +++++++++++--------
 .../data_validation/MissingValuesBarPlot.py   |  4 +++-
 .../data_validation/MutualInformation.py      |  3 ++-
 .../PearsonCorrelationMatrix.py               |  2 +-
 .../ProtectedClassesCombination.py            |  2 ++
 .../ProtectedClassesDescription.py            |  4 ++--
 .../ProtectedClassesDisparity.py              | 14 +++++++----
 .../ProtectedClassesThresholdOptimizer.py     | 12 ++++++++--
 .../tests/data_validation/RollingStatsPlot.py |  5 +++-
 .../data_validation/ScoreBandDefaultRates.py  |  6 +++--
 .../data_validation/SeasonalDecompose.py      |  2 +-
 .../tests/data_validation/ShapiroWilk.py      |  4 ++--
 validmind/tests/data_validation/SpreadPlot.py |  2 +-
 .../TabularCategoricalBarPlots.py             |  2 +-
 .../TabularDateTimeHistograms.py              |  2 +-
 .../data_validation/TargetRateBarPlots.py     |  5 +++-
 .../data_validation/TimeSeriesFrequency.py    |  2 +-
 .../data_validation/TimeSeriesOutliers.py     |  9 +++++--
 .../tests/data_validation/WOEBinPlots.py      |  2 +-
 .../tests/data_validation/WOEBinTable.py      |  2 +-
 .../tests/data_validation/ZivotAndrewsArch.py |  7 ++++--
 37 files changed, 119 insertions(+), 63 deletions(-)

diff --git a/validmind/tests/data_validation/ACFandPACFPlot.py b/validmind/tests/data_validation/ACFandPACFPlot.py
index 804960efe..c2b6adbc1 100644
--- a/validmind/tests/data_validation/ACFandPACFPlot.py
+++ b/validmind/tests/data_validation/ACFandPACFPlot.py
@@ -94,4 +94,7 @@ def ACFandPACFPlot(dataset: VMDataset):
         figures.append(pacf_fig)
         pacf_store[col] = pacf_values
 
-    return (*figures, RawData(acf_values=acf_store, pacf_values=pacf_store))
+    return (
+        *figures,
+        RawData(acf_values=acf_store, pacf_values=pacf_store, dataset=dataset.input_id),
+    )
diff --git a/validmind/tests/data_validation/AutoMA.py b/validmind/tests/data_validation/AutoMA.py
index a1cf2f34c..1c28ee7de 100644
--- a/validmind/tests/data_validation/AutoMA.py
+++ b/validmind/tests/data_validation/AutoMA.py
@@ -116,4 +116,4 @@ def AutoMA(dataset: VMDataset, max_ma_order: int = 3):
     return {
         "Auto MA Analysis Results": summary_ma_analysis,
         "Best MA Order Results": best_ma_order,
-    }, RawData(raw_series_data=df)
+    }, RawData(raw_series_data=df, dataset=dataset.input_id)
diff --git a/validmind/tests/data_validation/BivariateScatterPlots.py b/validmind/tests/data_validation/BivariateScatterPlots.py
index 9c29a96fc..c22af160e 100644
--- a/validmind/tests/data_validation/BivariateScatterPlots.py
+++ b/validmind/tests/data_validation/BivariateScatterPlots.py
@@ -80,5 +80,9 @@ def BivariateScatterPlots(dataset):
         figures.append(fig)
 
     return tuple(figures) + (
-        RawData(selected_numerical_df=df, feature_pairs=features_pairs),
+        RawData(
+            selected_numerical_df=df,
+            feature_pairs=features_pairs,
+            dataset=dataset.input_id,
+        ),
     )
diff --git a/validmind/tests/data_validation/BoxPierce.py b/validmind/tests/data_validation/BoxPierce.py
index cf9b8ac35..a85e57a8e 100644
--- a/validmind/tests/data_validation/BoxPierce.py
+++ b/validmind/tests/data_validation/BoxPierce.py
@@ -68,4 +68,6 @@ def BoxPierce(dataset):
     box_pierce_df.reset_index(inplace=True)
     box_pierce_df.columns = ["column", "stat", "pvalue"]
 
-    return box_pierce_df, RawData(box_pierce_values=box_pierce_values)
+    return box_pierce_df, RawData(
+        box_pierce_values=box_pierce_values, dataset=dataset.input_id
+    )
diff --git a/validmind/tests/data_validation/ClassImbalance.py b/validmind/tests/data_validation/ClassImbalance.py
index 24968231d..bab192e74 100644
--- a/validmind/tests/data_validation/ClassImbalance.py
+++ b/validmind/tests/data_validation/ClassImbalance.py
@@ -104,5 +104,5 @@ def ClassImbalance(
         },
         go.Figure(data=[trace], layout=layout),
         all(row["Pass/Fail"] == "Pass" for row in imbalanced_classes),
-        RawData(imbalance_percentages=imbalance_percentages),
+        RawData(imbalance_percentages=imbalance_percentages, dataset=dataset.input_id),
     )
diff --git a/validmind/tests/data_validation/DatasetDescription.py b/validmind/tests/data_validation/DatasetDescription.py
index cec442297..f8f3d0699 100644
--- a/validmind/tests/data_validation/DatasetDescription.py
+++ b/validmind/tests/data_validation/DatasetDescription.py
@@ -242,4 +242,4 @@ def DatasetDescription(dataset: VMDataset):
             }
             for column in results
         ]
-    }, RawData(raw_data=raw_data)
+    }, RawData(raw_data=raw_data, dataset=dataset.input_id)
diff --git a/validmind/tests/data_validation/DickeyFullerGLS.py b/validmind/tests/data_validation/DickeyFullerGLS.py
index d53a9aa76..0a9dda2c6 100644
--- a/validmind/tests/data_validation/DickeyFullerGLS.py
+++ b/validmind/tests/data_validation/DickeyFullerGLS.py
@@ -97,4 +97,4 @@ def DickeyFullerGLS(dataset: VMDataset):
 
     return {
         "DFGLS Test Results": dfgls_values,
-    }, RawData(df=df)
+    }, RawData(df=df, dataset=dataset.input_id)
diff --git a/validmind/tests/data_validation/FeatureTargetCorrelationPlot.py b/validmind/tests/data_validation/FeatureTargetCorrelationPlot.py
index 5fb06c998..83c2523db 100644
--- a/validmind/tests/data_validation/FeatureTargetCorrelationPlot.py
+++ b/validmind/tests/data_validation/FeatureTargetCorrelationPlot.py
@@ -58,7 +58,7 @@ def FeatureTargetCorrelationPlot(dataset, fig_height=600):
         df, dataset.target_column, fig_height
     )
 
-    return fig, RawData(correlation_data=correlations)
+    return fig, RawData(correlation_data=correlations, dataset=dataset.input_id)
 
 
 def _visualize_feature_target_correlation(df, target_column, fig_height):
diff --git a/validmind/tests/data_validation/HighCardinality.py b/validmind/tests/data_validation/HighCardinality.py
index 445fe1be8..078b7f034 100644
--- a/validmind/tests/data_validation/HighCardinality.py
+++ b/validmind/tests/data_validation/HighCardinality.py
@@ -83,4 +83,8 @@ def HighCardinality(
         if not passed:
             all_passed = False
 
-    return table, all_passed, RawData(raw_cardinality_details=raw_data)
+    return (
+        table,
+        all_passed,
+        RawData(raw_cardinality_details=raw_data, dataset=dataset.input_id),
+    )
diff --git a/validmind/tests/data_validation/HighPearsonCorrelation.py b/validmind/tests/data_validation/HighPearsonCorrelation.py
index 1b849cd4d..d533ca3c1 100644
--- a/validmind/tests/data_validation/HighPearsonCorrelation.py
+++ b/validmind/tests/data_validation/HighPearsonCorrelation.py
@@ -84,5 +84,5 @@ def HighPearsonCorrelation(
     return (
         pairs,
         all(p["Pass/Fail"] == "Pass" for p in pairs),
-        RawData(correlation_matrix=corr),
+        RawData(correlation_matrix=corr, dataset=dataset.input_id),
     )
diff --git a/validmind/tests/data_validation/IQROutliersBarPlot.py b/validmind/tests/data_validation/IQROutliersBarPlot.py
index 0efe50334..c8c639750 100644
--- a/validmind/tests/data_validation/IQROutliersBarPlot.py
+++ b/validmind/tests/data_validation/IQROutliersBarPlot.py
@@ -118,11 +118,13 @@ def IQROutliersBarPlot(
         )
         figures.append(fig)
 
+    outliers_by_feature = df[dataset.feature_columns_numeric].apply(
+        lambda col: compute_outliers(col, threshold)
+    )
+
     return (
         *figures,
         RawData(
-            outlier_counts_by_feature=df[dataset.feature_columns_numeric].apply(
-                lambda col: compute_outliers(col, threshold)
-            )
+            outlier_counts_by_feature=outliers_by_feature, dataset=dataset.input_id
         ),
     )
diff --git a/validmind/tests/data_validation/IQROutliersTable.py b/validmind/tests/data_validation/IQROutliersTable.py
index 1b7f0e6bf..d9e154ca4 100644
--- a/validmind/tests/data_validation/IQROutliersTable.py
+++ b/validmind/tests/data_validation/IQROutliersTable.py
@@ -2,7 +2,7 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset
 
 
@@ -64,6 +64,7 @@ def IQROutliersTable(dataset: VMDataset, threshold: float = 1.5):
     df = dataset.df
 
     outliers_table = []
+    all_outliers = {}
 
     for col in dataset.feature_columns_numeric:
         # Skip binary features
@@ -71,6 +72,8 @@ def IQROutliersTable(dataset: VMDataset, threshold: float = 1.5):
             continue
 
         outliers = compute_outliers(df[col], threshold)
+        all_outliers[col] = outliers
+
         if outliers.empty:
             continue
 
@@ -89,4 +92,4 @@ def IQROutliersTable(dataset: VMDataset, threshold: float = 1.5):
 
     return {
         "Summary of Outliers Detected by IQR Method": outliers_table,
-    }
+    }, RawData(all_outliers=all_outliers, dataset=dataset.input_id)
diff --git a/validmind/tests/data_validation/IsolationForestOutliers.py b/validmind/tests/data_validation/IsolationForestOutliers.py
index 6b6b191fc..5120a97c9 100644
--- a/validmind/tests/data_validation/IsolationForestOutliers.py
+++ b/validmind/tests/data_validation/IsolationForestOutliers.py
@@ -8,7 +8,7 @@
 import seaborn as sns
 from sklearn.ensemble import IsolationForest
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset
 
 
@@ -91,6 +91,7 @@ def IsolationForestOutliers(
 
         figures.append(fig)
 
-        plt.close()
-
-    return tuple(figures)
+    return (
+        *figures,
+        RawData(predictions=y_pred, dataset=dataset.input_id),
+    )
diff --git a/validmind/tests/data_validation/JarqueBera.py b/validmind/tests/data_validation/JarqueBera.py
index 927474a5b..d9857f8f5 100644
--- a/validmind/tests/data_validation/JarqueBera.py
+++ b/validmind/tests/data_validation/JarqueBera.py
@@ -5,7 +5,7 @@
 import pandas as pd
 from statsmodels.stats.stattools import jarque_bera
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 
 
 @tasks("classification", "regression")
@@ -67,4 +67,4 @@ def JarqueBera(dataset):
     jb_df.reset_index(inplace=True)
     jb_df.columns = ["column", "stat", "pvalue", "skew", "kurtosis"]
 
-    return jb_df
+    return jb_df, RawData(jb_values=jb_values, dataset=dataset.input_id)
diff --git a/validmind/tests/data_validation/LJungBox.py b/validmind/tests/data_validation/LJungBox.py
index f746379bb..6ffb966d7 100644
--- a/validmind/tests/data_validation/LJungBox.py
+++ b/validmind/tests/data_validation/LJungBox.py
@@ -5,7 +5,7 @@
 import pandas as pd
 from statsmodels.stats.diagnostic import acorr_ljungbox
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 
 
 @tasks("regression")
@@ -63,4 +63,4 @@ def LJungBox(dataset):
     ljung_box_df.reset_index(inplace=True)
     ljung_box_df.columns = ["column", "stat", "pvalue"]
 
-    return ljung_box_df
+    return ljung_box_df, RawData(ljung_box_df=ljung_box_df, dataset=dataset.input_id)
diff --git a/validmind/tests/data_validation/LaggedCorrelationHeatmap.py b/validmind/tests/data_validation/LaggedCorrelationHeatmap.py
index 61c4d561c..d295b1caa 100644
--- a/validmind/tests/data_validation/LaggedCorrelationHeatmap.py
+++ b/validmind/tests/data_validation/LaggedCorrelationHeatmap.py
@@ -101,4 +101,4 @@ def LaggedCorrelationHeatmap(dataset: VMDataset, num_lags: int = 10):
         xaxis_title="Lags",
     )
 
-    return fig, RawData(correlation_matrix=correlation_df)
+    return fig, RawData(correlation_matrix=correlation_df, dataset=dataset.input_id)
diff --git a/validmind/tests/data_validation/MissingValues.py b/validmind/tests/data_validation/MissingValues.py
index 93ba4be2a..7aae5a14d 100644
--- a/validmind/tests/data_validation/MissingValues.py
+++ b/validmind/tests/data_validation/MissingValues.py
@@ -2,7 +2,7 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset
 
 
@@ -49,12 +49,16 @@ def MissingValues(dataset: VMDataset, min_threshold: int = 1):
     df = dataset.df
     missing = df.isna().sum()
 
-    return [
-        {
-            "Column": col,
-            "Number of Missing Values": missing[col],
-            "Percentage of Missing Values (%)": missing[col] / df.shape[0] * 100,
-            "Pass/Fail": "Pass" if missing[col] < min_threshold else "Fail",
-        }
-        for col in missing.index
-    ], all(missing[col] < min_threshold for col in missing.index)
+    return (
+        [
+            {
+                "Column": col,
+                "Number of Missing Values": missing[col],
+                "Percentage of Missing Values (%)": missing[col] / df.shape[0] * 100,
+                "Pass/Fail": "Pass" if missing[col] < min_threshold else "Fail",
+            }
+            for col in missing.index
+        ],
+        all(missing[col] < min_threshold for col in missing.index),
+        RawData(missing_values=missing, dataset=dataset.input_id),
+    )
diff --git a/validmind/tests/data_validation/MissingValuesBarPlot.py b/validmind/tests/data_validation/MissingValuesBarPlot.py
index 9267f1a16..8ec66519d 100644
--- a/validmind/tests/data_validation/MissingValuesBarPlot.py
+++ b/validmind/tests/data_validation/MissingValuesBarPlot.py
@@ -117,5 +117,7 @@ def MissingValuesBarPlot(
                 height=fig_height,
             ),
         ),
-        RawData(missing_percentages=missing_percentages_sorted),
+        RawData(
+            missing_percentages=missing_percentages_sorted, dataset=dataset.input_id
+        ),
     )
diff --git a/validmind/tests/data_validation/MutualInformation.py b/validmind/tests/data_validation/MutualInformation.py
index 7c230576f..03a2d259c 100644
--- a/validmind/tests/data_validation/MutualInformation.py
+++ b/validmind/tests/data_validation/MutualInformation.py
@@ -123,5 +123,6 @@ def MutualInformation(
     return fig, RawData(
         mutual_information_scores={
             feature: score for feature, score in zip(sorted_features, sorted_scores)
-        }
+        },
+        dataset=dataset.input_id,
     )
diff --git a/validmind/tests/data_validation/PearsonCorrelationMatrix.py b/validmind/tests/data_validation/PearsonCorrelationMatrix.py
index 9cfb6a284..5b8d756a1 100644
--- a/validmind/tests/data_validation/PearsonCorrelationMatrix.py
+++ b/validmind/tests/data_validation/PearsonCorrelationMatrix.py
@@ -88,4 +88,4 @@ def PearsonCorrelationMatrix(dataset):
 
     fig = go.Figure(data=[heatmap], layout=layout)
 
-    return fig, RawData(correlation_matrix=corr_matrix)
+    return fig, RawData(correlation_matrix=corr_matrix, dataset=dataset.input_id)
diff --git a/validmind/tests/data_validation/ProtectedClassesCombination.py b/validmind/tests/data_validation/ProtectedClassesCombination.py
index bbd48c8c8..606144d45 100644
--- a/validmind/tests/data_validation/ProtectedClassesCombination.py
+++ b/validmind/tests/data_validation/ProtectedClassesCombination.py
@@ -206,5 +206,7 @@ def ProtectedClassesCombination(dataset, model, protected_classes=None):
             metrics_frame=mf,
             demographic_parity_ratios=m_dpr,
             equalized_odds_ratios=m_eqo,
+            model=model.input_id,
+            dataset=dataset.input_id,
         ),
     )
diff --git a/validmind/tests/data_validation/ProtectedClassesDescription.py b/validmind/tests/data_validation/ProtectedClassesDescription.py
index 53f763705..cb4e631ac 100644
--- a/validmind/tests/data_validation/ProtectedClassesDescription.py
+++ b/validmind/tests/data_validation/ProtectedClassesDescription.py
@@ -6,7 +6,7 @@
 import pandas as pd
 import plotly.graph_objects as go
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.logging import get_logger
 
 logger = get_logger(__name__)
@@ -127,4 +127,4 @@ def ProtectedClassesDescription(dataset, protected_classes=None):
         ["Protected Class", "Count"], ascending=[True, False]
     )
 
-    return (stats_df, *figures)
+    return (stats_df, *figures, RawData(dataset=dataset.input_id))
diff --git a/validmind/tests/data_validation/ProtectedClassesDisparity.py b/validmind/tests/data_validation/ProtectedClassesDisparity.py
index 61a7d9b1e..ac9b88251 100644
--- a/validmind/tests/data_validation/ProtectedClassesDisparity.py
+++ b/validmind/tests/data_validation/ProtectedClassesDisparity.py
@@ -7,7 +7,7 @@
 
 import pandas as pd
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.errors import MissingDependencyError
 from validmind.logging import get_logger
 
@@ -119,7 +119,7 @@ def ProtectedClassesDisparity(
         mask_significance=True,
     )
 
-    figures = []
+    returns = []  # Renamed to 'returns' for clarity
     for protected_class in protected_classes:
         plot = ap.disparity(
             bdf, metrics, protected_class, fairness_threshold=disparity_tolerance
@@ -129,12 +129,16 @@ def ProtectedClassesDisparity(
         plot.save(
             buf, format="png"
         )  # as long as the above library is installed, this will work
-        figures.append(buf.getvalue())
+        returns.append(buf.getvalue())
 
     string = "_disparity"
     metrics_adj = [x + string for x in metrics]
 
     table = bdf[["attribute_name", "attribute_value"] + b.list_disparities(bdf)]
-    figures.append(aqp.plot_disparity_all(bdf, metrics=metrics_adj))
+    returns.append(aqp.plot_disparity_all(bdf, metrics=metrics_adj))
 
-    return (table, *figures)
+    return (
+        table,
+        *returns,
+        RawData(model=model.input_id, dataset=dataset.input_id, disparity_data=bdf),
+    )
diff --git a/validmind/tests/data_validation/ProtectedClassesThresholdOptimizer.py b/validmind/tests/data_validation/ProtectedClassesThresholdOptimizer.py
index 4ed285b8f..31124a621 100644
--- a/validmind/tests/data_validation/ProtectedClassesThresholdOptimizer.py
+++ b/validmind/tests/data_validation/ProtectedClassesThresholdOptimizer.py
@@ -8,7 +8,7 @@
 import matplotlib.pyplot as plt
 import pandas as pd
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.errors import MissingDependencyError
 from validmind.logging import get_logger
 
@@ -103,7 +103,15 @@ def ProtectedClassesThresholdOptimizer(
         test_df, target, y_pred_opt, protected_classes
     )
 
-    return {"DPR and EOR Table": fairness_metrics.reset_index()}, fig
+    return (
+        {"DPR and EOR Table": fairness_metrics.reset_index()},
+        fig,
+        RawData(
+            y_predictions=y_pred_opt.tolist(),
+            dataset=dataset.input_id,
+            protected_classes=protected_classes,
+        ),
+    )
 
 
 def initialize_and_fit_optimizer(pipeline, X_train, y_train, protected_classes_df):
diff --git a/validmind/tests/data_validation/RollingStatsPlot.py b/validmind/tests/data_validation/RollingStatsPlot.py
index 2d2e2d8f8..3ef99dd76 100644
--- a/validmind/tests/data_validation/RollingStatsPlot.py
+++ b/validmind/tests/data_validation/RollingStatsPlot.py
@@ -113,6 +113,9 @@ def RollingStatsPlot(dataset: VMDataset, window_size: int = 12):
                     "rolling_std": dataset.df[col].rolling(window=window_size).std(),
                 }
                 for col in dataset.feature_columns
-            }
+            },
+            dataset=dataset.input_id,
         ),
     )
+
+    return tuple(returns)
diff --git a/validmind/tests/data_validation/ScoreBandDefaultRates.py b/validmind/tests/data_validation/ScoreBandDefaultRates.py
index 1f96b12a6..ba995ad1c 100644
--- a/validmind/tests/data_validation/ScoreBandDefaultRates.py
+++ b/validmind/tests/data_validation/ScoreBandDefaultRates.py
@@ -5,7 +5,7 @@
 import numpy as np
 import pandas as pd
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
 
@@ -137,4 +137,6 @@ def ScoreBandDefaultRates(
         }
     )
 
-    return pd.DataFrame(results)
+    return pd.DataFrame(results), RawData(
+        results=results, model=model.input_id, dataset=dataset.input_id
+    )
diff --git a/validmind/tests/data_validation/SeasonalDecompose.py b/validmind/tests/data_validation/SeasonalDecompose.py
index a764dbfba..e72879a5e 100644
--- a/validmind/tests/data_validation/SeasonalDecompose.py
+++ b/validmind/tests/data_validation/SeasonalDecompose.py
@@ -166,4 +166,4 @@ def SeasonalDecompose(dataset: VMDataset, seasonal_model: str = "additive"):
     if not figures:
         raise SkipTestError("No valid features found for seasonal decomposition")
 
-    return (*figures, RawData(decomposed_components=raw_data))
+    return (*figures, RawData(decomposed_components=raw_data, dataset=dataset.input_id))
diff --git a/validmind/tests/data_validation/ShapiroWilk.py b/validmind/tests/data_validation/ShapiroWilk.py
index b41d0cae3..b299040c7 100644
--- a/validmind/tests/data_validation/ShapiroWilk.py
+++ b/validmind/tests/data_validation/ShapiroWilk.py
@@ -5,7 +5,7 @@
 import pandas as pd
 from scipy import stats
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 
 
 @tasks("classification", "regression")
@@ -66,4 +66,4 @@ def ShapiroWilk(dataset):
     sw_df.reset_index(inplace=True)
     sw_df.columns = ["column", "stat", "pvalue"]
 
-    return sw_df
+    return sw_df, RawData(shapiro_results=sw_values, dataset=dataset.input_id)
diff --git a/validmind/tests/data_validation/SpreadPlot.py b/validmind/tests/data_validation/SpreadPlot.py
index 835fbdf1d..d4f40df06 100644
--- a/validmind/tests/data_validation/SpreadPlot.py
+++ b/validmind/tests/data_validation/SpreadPlot.py
@@ -95,4 +95,4 @@ def SpreadPlot(dataset: VMDataset):
             name=f"spread_{var1}_{var2}"
         )
 
-    return (*figures, RawData(spread_data=spread_data))
+    return (*figures, RawData(spread_data=spread_data, dataset=dataset.input_id))
diff --git a/validmind/tests/data_validation/TabularCategoricalBarPlots.py b/validmind/tests/data_validation/TabularCategoricalBarPlots.py
index ca3c16930..4a2419f34 100644
--- a/validmind/tests/data_validation/TabularCategoricalBarPlots.py
+++ b/validmind/tests/data_validation/TabularCategoricalBarPlots.py
@@ -92,4 +92,4 @@ def TabularCategoricalBarPlots(dataset: VMDataset):
         )
         figures.append(fig)
 
-    return (*figures, RawData(category_counts=counts_dict))
+    return (*figures, RawData(category_counts=counts_dict, dataset=dataset.input_id))
diff --git a/validmind/tests/data_validation/TabularDateTimeHistograms.py b/validmind/tests/data_validation/TabularDateTimeHistograms.py
index bf56f0d82..9d2a6ffca 100644
--- a/validmind/tests/data_validation/TabularDateTimeHistograms.py
+++ b/validmind/tests/data_validation/TabularDateTimeHistograms.py
@@ -72,4 +72,4 @@ def TabularDateTimeHistograms(dataset: VMDataset):
         font=dict(size=18),
     )
 
-    return fig, RawData(date_differences=date_diffs)
+    return fig, RawData(date_differences=date_diffs, dataset=dataset.input_id)
diff --git a/validmind/tests/data_validation/TargetRateBarPlots.py b/validmind/tests/data_validation/TargetRateBarPlots.py
index f6b8b3eee..094906fd7 100644
--- a/validmind/tests/data_validation/TargetRateBarPlots.py
+++ b/validmind/tests/data_validation/TargetRateBarPlots.py
@@ -108,4 +108,7 @@ def TargetRateBarPlots(dataset: VMDataset):
 
         figures.append(fig)
 
-    return (*figures, RawData(target_rates_by_category=raw_data))
+    return (
+        *figures,
+        RawData(target_rates_by_category=raw_data, dataset=dataset.input_id),
+    )
diff --git a/validmind/tests/data_validation/TimeSeriesFrequency.py b/validmind/tests/data_validation/TimeSeriesFrequency.py
index aa34d7d75..f89f96e09 100644
--- a/validmind/tests/data_validation/TimeSeriesFrequency.py
+++ b/validmind/tests/data_validation/TimeSeriesFrequency.py
@@ -107,5 +107,5 @@ def TimeSeriesFrequency(dataset: VMDataset):
         frequencies,
         fig,
         len(set(item["Frequency"] for item in frequencies)) == 1,
-        RawData(time_diff_days=time_diff_days),
+        RawData(time_diff_days=time_diff_days, dataset=dataset.input_id),
     )
diff --git a/validmind/tests/data_validation/TimeSeriesOutliers.py b/validmind/tests/data_validation/TimeSeriesOutliers.py
index cc6566b95..6cf926e62 100644
--- a/validmind/tests/data_validation/TimeSeriesOutliers.py
+++ b/validmind/tests/data_validation/TimeSeriesOutliers.py
@@ -5,7 +5,7 @@
 import pandas as pd
 import plotly.graph_objects as go
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.errors import SkipTestError
 from validmind.vm_models import VMDataset
 
@@ -111,4 +111,9 @@ def TimeSeriesOutliers(dataset: VMDataset, zscore_threshold: int = 3):
 
         figures.append(fig)
 
-    return (outlier_df.sort_values(["Column", "Date"]), figures, len(outlier_df) == 0)
+    return (
+        outlier_df.sort_values(["Column", "Date"]),
+        figures,
+        len(outlier_df) == 0,
+        RawData(outliers=outlier_df, dataset=dataset.input_id),
+    )
diff --git a/validmind/tests/data_validation/WOEBinPlots.py b/validmind/tests/data_validation/WOEBinPlots.py
index 3ac139338..3a72a4ddc 100644
--- a/validmind/tests/data_validation/WOEBinPlots.py
+++ b/validmind/tests/data_validation/WOEBinPlots.py
@@ -140,4 +140,4 @@ def WOEBinPlots(
 
         figures.append(fig)
 
-    return (*figures, RawData(woe_iv_data=woe_iv_df))
+    return (*figures, RawData(woe_iv_data=woe_iv_df, dataset=dataset.input_id))
diff --git a/validmind/tests/data_validation/WOEBinTable.py b/validmind/tests/data_validation/WOEBinTable.py
index 94e05db0d..b3ef04098 100644
--- a/validmind/tests/data_validation/WOEBinTable.py
+++ b/validmind/tests/data_validation/WOEBinTable.py
@@ -71,4 +71,4 @@ def WOEBinTable(dataset: VMDataset, breaks_adj: list = None):
 
     return {
         "Weight of Evidence (WoE) and Information Value (IV)": result_table
-    }, RawData(woe_bins=bins)
+    }, RawData(woe_bins=bins, dataset=dataset.input_id)
diff --git a/validmind/tests/data_validation/ZivotAndrewsArch.py b/validmind/tests/data_validation/ZivotAndrewsArch.py
index ec2b560c6..07b689acc 100644
--- a/validmind/tests/data_validation/ZivotAndrewsArch.py
+++ b/validmind/tests/data_validation/ZivotAndrewsArch.py
@@ -6,7 +6,7 @@
 from arch.unitroot import ZivotAndrews
 from numpy.linalg import LinAlgError
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.errors import SkipTestError
 from validmind.logging import get_logger
 from validmind.vm_models import VMDataset
@@ -83,4 +83,7 @@ def ZivotAndrewsArch(dataset: VMDataset):
             }
         )
 
-    return {"Zivot-Andrews Test Results": za_values}
+    return (
+        {"Zivot-Andrews Test Results": za_values},
+        RawData(zivot_andrews=za_values, dataset=dataset.input_id),
+    )

From 0c05a90e11b197f846a49d6cbda2d764c8bf5846 Mon Sep 17 00:00:00 2001
From: Juan <juan@validmind.ai>
Date: Tue, 18 Feb 2025 14:32:33 +0100
Subject: [PATCH 03/10] Fix lint

---
 validmind/tests/data_validation/RollingStatsPlot.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/validmind/tests/data_validation/RollingStatsPlot.py b/validmind/tests/data_validation/RollingStatsPlot.py
index 3ef99dd76..7f05735f6 100644
--- a/validmind/tests/data_validation/RollingStatsPlot.py
+++ b/validmind/tests/data_validation/RollingStatsPlot.py
@@ -117,5 +117,3 @@ def RollingStatsPlot(dataset: VMDataset, window_size: int = 12):
             dataset=dataset.input_id,
         ),
     )
-
-    return tuple(returns)

From 151184cd25a524eee512d49498567210b965ca93 Mon Sep 17 00:00:00 2001
From: Juan <juan@validmind.ai>
Date: Tue, 18 Feb 2025 15:12:05 +0100
Subject: [PATCH 04/10] Update RawData for model validation tests

---
 validmind/tests/model_validation/BertScore.py |  8 +++-
 validmind/tests/model_validation/BleuScore.py |  8 +++-
 .../ClusterSizeDistribution.py                |  4 +-
 .../model_validation/ContextualRecall.py      | 10 ++++-
 .../tests/model_validation/FeaturesAUC.py     |  2 +-
 .../tests/model_validation/MeteorScore.py     |  8 +++-
 .../ModelPredictionResiduals.py               |  6 ++-
 .../tests/model_validation/RegardScore.py     |  7 +++-
 .../RegressionResidualsPlot.py                | 11 +++++-
 .../tests/model_validation/RougeScore.py      |  4 +-
 .../TimeSeriesPredictionWithCI.py             |  2 +
 .../TimeSeriesPredictionsPlot.py              | 12 +++++-
 .../TimeSeriesR2SquareBySegments.py           |  8 +++-
 .../tests/model_validation/TokenDisparity.py  |  6 ++-
 .../tests/model_validation/ToxicityScore.py   |  2 +
 .../embeddings/ClusterDistribution.py         |  2 +-
 .../embeddings/CosineSimilarityComparison.py  |  6 ++-
 .../CosineSimilarityDistribution.py           |  6 ++-
 .../embeddings/CosineSimilarityHeatmap.py     |  6 ++-
 .../embeddings/DescriptiveAnalytics.py        |  2 +
 .../embeddings/EmbeddingsVisualization2D.py   |  6 ++-
 .../embeddings/EuclideanDistanceComparison.py |  8 +++-
 .../embeddings/EuclideanDistanceHeatmap.py    |  4 +-
 .../embeddings/PCAComponentsPairwisePlots.py  |  5 ++-
 .../embeddings/StabilityAnalysisKeyword.py    |  6 ++-
 .../StabilityAnalysisRandomNoise.py           |  6 ++-
 .../embeddings/StabilityAnalysisSynonyms.py   |  6 ++-
 .../StabilityAnalysisTranslation.py           |  6 ++-
 .../embeddings/TSNEComponentsPairwisePlots.py |  7 +++-
 .../ragas/AnswerCorrectness.py                |  2 +-
 .../model_validation/ragas/AspectCritic.py    |  5 ++-
 .../ragas/ContextEntityRecall.py              |  2 +-
 .../ragas/ContextPrecision.py                 |  2 +-
 .../ragas/ContextPrecisionWithoutReference.py |  2 +-
 .../model_validation/ragas/ContextRecall.py   |  2 +-
 .../model_validation/ragas/Faithfulness.py    |  2 +-
 .../ragas/NoiseSensitivity.py                 |  2 +-
 .../ragas/ResponseRelevancy.py                |  2 +-
 .../ragas/SemanticSimilarity.py               |  2 +-
 .../sklearn/AdjustedMutualInformation.py      | 18 ++++-----
 .../sklearn/AdjustedRandIndex.py              | 18 ++++-----
 .../sklearn/CalibrationCurve.py               |  7 +++-
 .../ClassifierThresholdOptimization.py        | 17 ++++++++-
 .../sklearn/ClusterCosineSimilarity.py        |  6 ++-
 .../sklearn/ClusterPerformanceMetrics.py      | 38 ++++++++++++-------
 .../sklearn/CompletenessScore.py              | 17 ++++-----
 .../sklearn/FeatureImportance.py              |  8 +++-
 .../sklearn/FowlkesMallowsScore.py            | 21 +++++-----
 .../sklearn/HomogeneityScore.py               | 23 ++++++-----
 .../sklearn/HyperParametersTuning.py          |  6 ++-
 .../sklearn/KMeansClustersOptimization.py     |  7 +++-
 .../sklearn/MinimumAccuracy.py                | 19 ++++++----
 .../sklearn/MinimumF1Score.py                 | 19 ++++++----
 .../sklearn/MinimumROCAUCScore.py             | 27 ++++++++++---
 .../sklearn/OverfitDiagnosis.py               | 10 ++++-
 .../sklearn/PermutationFeatureImportance.py   |  6 ++-
 .../sklearn/PopulationStabilityIndex.py       |  6 ++-
 .../sklearn/PrecisionRecallCurve.py           |  7 +++-
 .../model_validation/sklearn/ROCCurve.py      |  4 +-
 .../sklearn/RegressionErrors.py               |  8 +++-
 .../sklearn/RegressionPerformance.py          | 21 ++++++----
 .../sklearn/RegressionR2Square.py             | 13 ++++---
 .../sklearn/RobustnessDiagnosis.py            | 11 +++++-
 .../sklearn/SHAPGlobalImportance.py           |  7 +++-
 .../sklearn/ScoreProbabilityAlignment.py      | 12 +++++-
 .../sklearn/SilhouettePlot.py                 |  6 ++-
 .../model_validation/sklearn/VMeasure.py      | 21 +++++-----
 .../CumulativePredictionProbabilities.py      |  6 ++-
 .../statsmodels/DurbinWatsonTest.py           |  7 +++-
 .../model_validation/statsmodels/GINITable.py |  9 ++++-
 .../statsmodels/KolmogorovSmirnov.py          |  4 +-
 .../PredictionProbabilitiesHistogram.py       |  8 +++-
 .../statsmodels/RegressionCoeffs.py           | 10 ++++-
 .../RegressionFeatureSignificance.py          |  4 +-
 .../RegressionModelForecastPlot.py            |  9 ++++-
 .../RegressionModelForecastPlotLevels.py      |  2 +
 .../RegressionModelSensitivityPlot.py         |  2 +
 .../statsmodels/RegressionModelSummary.py     |  6 ++-
 .../RegressionPermutationFeatureImportance.py |  4 +-
 79 files changed, 472 insertions(+), 179 deletions(-)

diff --git a/validmind/tests/model_validation/BertScore.py b/validmind/tests/model_validation/BertScore.py
index bd8257b19..72f04243d 100644
--- a/validmind/tests/model_validation/BertScore.py
+++ b/validmind/tests/model_validation/BertScore.py
@@ -131,4 +131,10 @@ def BertScore(
     # Create a DataFrame from all collected statistics
     result_df = pd.DataFrame(stats_df).reset_index().rename(columns={"index": "Metric"})
 
-    return (result_df, *figures, RawData(bert_scores_df=metrics_df))
+    return (
+        result_df,
+        *figures,
+        RawData(
+            bert_scores_df=metrics_df, model=model.input_id, dataset=dataset.input_id
+        ),
+    )
diff --git a/validmind/tests/model_validation/BleuScore.py b/validmind/tests/model_validation/BleuScore.py
index b524d3227..8d32fe0c1 100644
--- a/validmind/tests/model_validation/BleuScore.py
+++ b/validmind/tests/model_validation/BleuScore.py
@@ -114,4 +114,10 @@ def BleuScore(dataset, model):
     # Create a DataFrame from all collected statistics
     result_df = pd.DataFrame(stats_df).reset_index().rename(columns={"index": "Metric"})
 
-    return (result_df, *figures, RawData(bleu_scores_df=metrics_df))
+    return (
+        result_df,
+        *figures,
+        RawData(
+            bleu_scores_df=metrics_df, model=model.input_id, dataset=dataset.input_id
+        ),
+    )
diff --git a/validmind/tests/model_validation/ClusterSizeDistribution.py b/validmind/tests/model_validation/ClusterSizeDistribution.py
index 9dbb6b485..4cddbe40f 100644
--- a/validmind/tests/model_validation/ClusterSizeDistribution.py
+++ b/validmind/tests/model_validation/ClusterSizeDistribution.py
@@ -72,4 +72,6 @@ def ClusterSizeDistribution(dataset: VMDataset, model: VMModel):
     fig.update_yaxes(title_text="Counts", showgrid=False)
     fig.update_layout(title_text="Cluster distribution", title_x=0.5, barmode="group")
 
-    return fig, RawData(cluster_counts=df_counts)
+    return fig, RawData(
+        cluster_counts=df_counts, model=model.input_id, dataset=dataset.input_id
+    )
diff --git a/validmind/tests/model_validation/ContextualRecall.py b/validmind/tests/model_validation/ContextualRecall.py
index 6bff63856..5930e9f63 100644
--- a/validmind/tests/model_validation/ContextualRecall.py
+++ b/validmind/tests/model_validation/ContextualRecall.py
@@ -118,4 +118,12 @@ def ContextualRecall(dataset, model):
     # Create a DataFrame from all collected statistics
     result_df = pd.DataFrame(stats_df).reset_index().rename(columns={"index": "Metric"})
 
-    return (result_df, *figures, RawData(contextual_recall_scores=metrics_df))
+    return (
+        result_df,
+        *figures,
+        RawData(
+            contextual_recall_scores=metrics_df,
+            model=model.input_id,
+            dataset=dataset.input_id,
+        ),
+    )
diff --git a/validmind/tests/model_validation/FeaturesAUC.py b/validmind/tests/model_validation/FeaturesAUC.py
index cd4ba5afd..44c218c8f 100644
--- a/validmind/tests/model_validation/FeaturesAUC.py
+++ b/validmind/tests/model_validation/FeaturesAUC.py
@@ -95,4 +95,4 @@ def FeaturesAUC(dataset: VMDataset, fontsize: int = 12, figure_height: int = 500
         height=figure_height,
     )
 
-    return fig, RawData(feature_aucs=aucs)
+    return fig, RawData(feature_aucs=aucs, dataset=dataset.input_id)
diff --git a/validmind/tests/model_validation/MeteorScore.py b/validmind/tests/model_validation/MeteorScore.py
index bf679eea0..10ed5b7fb 100644
--- a/validmind/tests/model_validation/MeteorScore.py
+++ b/validmind/tests/model_validation/MeteorScore.py
@@ -117,4 +117,10 @@ def MeteorScore(dataset, model):
     # Create a DataFrame from all collected statistics
     result_df = pd.DataFrame(stats_df).reset_index().rename(columns={"index": "Metric"})
 
-    return (result_df, *figures, RawData(meteor_scores=metrics_df))
+    return (
+        result_df,
+        *figures,
+        RawData(
+            meteor_scores=metrics_df, model=model.input_id, dataset=dataset.input_id
+        ),
+    )
diff --git a/validmind/tests/model_validation/ModelPredictionResiduals.py b/validmind/tests/model_validation/ModelPredictionResiduals.py
index ceaa5fbf4..26a4f837c 100644
--- a/validmind/tests/model_validation/ModelPredictionResiduals.py
+++ b/validmind/tests/model_validation/ModelPredictionResiduals.py
@@ -102,4 +102,8 @@ def ModelPredictionResiduals(
     # Create a summary DataFrame for the KS normality test results
     summary_df = pd.DataFrame([summary])
 
-    return (summary_df, *figures, RawData(residuals=residuals))
+    return (
+        summary_df,
+        *figures,
+        RawData(residuals=residuals, model=model.input_id, dataset=dataset.input_id),
+    )
diff --git a/validmind/tests/model_validation/RegardScore.py b/validmind/tests/model_validation/RegardScore.py
index bcaa57d95..6a362f09e 100644
--- a/validmind/tests/model_validation/RegardScore.py
+++ b/validmind/tests/model_validation/RegardScore.py
@@ -145,5 +145,10 @@ def calculate_stats(df, metric_name):
     return (
         result_df,
         *figures,
-        RawData(true_regard=true_df, pred_regard=pred_df),
+        RawData(
+            true_regard=true_df,
+            pred_regard=pred_df,
+            model=model.input_id,
+            dataset=dataset.input_id,
+        ),
     )
diff --git a/validmind/tests/model_validation/RegressionResidualsPlot.py b/validmind/tests/model_validation/RegressionResidualsPlot.py
index 723bbbe53..813f6bb62 100644
--- a/validmind/tests/model_validation/RegressionResidualsPlot.py
+++ b/validmind/tests/model_validation/RegressionResidualsPlot.py
@@ -105,4 +105,13 @@ def RegressionResidualsPlot(model: VMModel, dataset: VMDataset, bin_size: float
         )
     )
 
-    return (*figures, RawData(residuals=residuals, y_true=y_true, y_pred=y_pred))
+    return (
+        *figures,
+        RawData(
+            residuals=residuals,
+            y_true=y_true,
+            y_pred=y_pred,
+            model=model.input_id,
+            dataset=dataset.input_id,
+        ),
+    )
diff --git a/validmind/tests/model_validation/RougeScore.py b/validmind/tests/model_validation/RougeScore.py
index 7d8f9a106..82a046143 100644
--- a/validmind/tests/model_validation/RougeScore.py
+++ b/validmind/tests/model_validation/RougeScore.py
@@ -121,5 +121,7 @@ def RougeScore(dataset, model, metric="rouge-1"):
     return (
         pd.DataFrame(stats_df).reset_index().rename(columns={"index": "Metric"}),
         *figures,
-        RawData(rouge_scores_df=df_scores),
+        RawData(
+            rouge_scores_df=df_scores, model=model.input_id, dataset=dataset.input_id
+        ),
     )
diff --git a/validmind/tests/model_validation/TimeSeriesPredictionWithCI.py b/validmind/tests/model_validation/TimeSeriesPredictionWithCI.py
index 329e3b42b..6edaed08f 100644
--- a/validmind/tests/model_validation/TimeSeriesPredictionWithCI.py
+++ b/validmind/tests/model_validation/TimeSeriesPredictionWithCI.py
@@ -152,5 +152,7 @@ def TimeSeriesPredictionWithCI(dataset, model, confidence=0.95):
             z_score=z_score,
             lower_confidence=lower_conf,
             upper_confidence=upper_conf,
+            model=model.input_id,
+            dataset=dataset.input_id,
         ),
     )
diff --git a/validmind/tests/model_validation/TimeSeriesPredictionsPlot.py b/validmind/tests/model_validation/TimeSeriesPredictionsPlot.py
index 6d2c84134..9a2239fe6 100644
--- a/validmind/tests/model_validation/TimeSeriesPredictionsPlot.py
+++ b/validmind/tests/model_validation/TimeSeriesPredictionsPlot.py
@@ -4,7 +4,7 @@
 
 import plotly.graph_objects as go
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 
 
 @tags("model_predictions", "visualization")
@@ -70,4 +70,12 @@ def TimeSeriesPredictionsPlot(dataset, model):
         template="plotly_white",
     )
 
-    return fig
+    raw_data = RawData(
+        time_index=time_index,
+        actual_values=dataset.y,
+        predicted_values=y_pred,
+        model=model.input_id,
+        dataset=dataset.input_id,
+    )
+
+    return fig, raw_data
diff --git a/validmind/tests/model_validation/TimeSeriesR2SquareBySegments.py b/validmind/tests/model_validation/TimeSeriesR2SquareBySegments.py
index 79d644e5f..3de28df23 100644
--- a/validmind/tests/model_validation/TimeSeriesR2SquareBySegments.py
+++ b/validmind/tests/model_validation/TimeSeriesR2SquareBySegments.py
@@ -7,7 +7,7 @@
 import plotly.express as px
 from sklearn import metrics
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 
 
 @tags("model_performance", "sklearn")
@@ -105,4 +105,8 @@ def TimeSeriesR2SquareBySegments(dataset, model, segments=None):
         },
     )
 
-    return fig, results_df
+    return (
+        fig,
+        results_df,
+        RawData(summary=results_df, model=model.input_id, dataset=dataset.input_id),
+    )
diff --git a/validmind/tests/model_validation/TokenDisparity.py b/validmind/tests/model_validation/TokenDisparity.py
index 89e654716..5d6961a2e 100644
--- a/validmind/tests/model_validation/TokenDisparity.py
+++ b/validmind/tests/model_validation/TokenDisparity.py
@@ -108,4 +108,8 @@ def TokenDisparity(dataset, model):
     # Create a DataFrame from all collected statistics
     result_df = pd.DataFrame(stats_df).reset_index().rename(columns={"index": "Metric"})
 
-    return (result_df, *figures, RawData(token_counts_df=df))
+    return (
+        result_df,
+        *figures,
+        RawData(token_counts_df=df, model=model.input_id, dataset=dataset.input_id),
+    )
diff --git a/validmind/tests/model_validation/ToxicityScore.py b/validmind/tests/model_validation/ToxicityScore.py
index 6f6314be5..37a5a5ce8 100644
--- a/validmind/tests/model_validation/ToxicityScore.py
+++ b/validmind/tests/model_validation/ToxicityScore.py
@@ -146,5 +146,7 @@ def calculate_stats(df):
             input_toxicity_df=input_df,
             true_toxicity_df=true_df,
             pred_toxicity_df=pred_df,
+            model=model.input_id,
+            dataset=dataset.input_id,
         ),
     )
diff --git a/validmind/tests/model_validation/embeddings/ClusterDistribution.py b/validmind/tests/model_validation/embeddings/ClusterDistribution.py
index 33140b39f..7c604d527 100644
--- a/validmind/tests/model_validation/embeddings/ClusterDistribution.py
+++ b/validmind/tests/model_validation/embeddings/ClusterDistribution.py
@@ -62,4 +62,4 @@ def ClusterDistribution(model: VMModel, dataset: VMDataset, num_clusters: int =
         title="Embeddings Cluster Distribution",
     )
 
-    return fig, RawData(labels=labels)
+    return fig, RawData(labels=labels, model=model.input_id, dataset=dataset.input_id)
diff --git a/validmind/tests/model_validation/embeddings/CosineSimilarityComparison.py b/validmind/tests/model_validation/embeddings/CosineSimilarityComparison.py
index e38561db5..27e919d61 100644
--- a/validmind/tests/model_validation/embeddings/CosineSimilarityComparison.py
+++ b/validmind/tests/model_validation/embeddings/CosineSimilarityComparison.py
@@ -113,5 +113,9 @@ def CosineSimilarityComparison(dataset, models):
     return (
         *figures,
         stats_df,
-        RawData(similarity_matrices=pd.DataFrame(similarity_matrices)),
+        RawData(
+            similarity_matrices=pd.DataFrame(similarity_matrices),
+            dataset=dataset.input_id,
+            models=[model.input_id for model in models],
+        ),
     )
diff --git a/validmind/tests/model_validation/embeddings/CosineSimilarityDistribution.py b/validmind/tests/model_validation/embeddings/CosineSimilarityDistribution.py
index 484370d3e..723a54028 100644
--- a/validmind/tests/model_validation/embeddings/CosineSimilarityDistribution.py
+++ b/validmind/tests/model_validation/embeddings/CosineSimilarityDistribution.py
@@ -59,4 +59,8 @@ def CosineSimilarityDistribution(dataset: VMDataset, model: VMModel):
         nbins=100,
         title="Cosine Similarity Distribution",
         labels={"x": "Cosine Similarity"},
-    ), RawData(similarity_scores=similarity_scores)
+    ), RawData(
+        similarity_scores=similarity_scores,
+        model=model.input_id,
+        dataset=dataset.input_id,
+    )
diff --git a/validmind/tests/model_validation/embeddings/CosineSimilarityHeatmap.py b/validmind/tests/model_validation/embeddings/CosineSimilarityHeatmap.py
index 9f490187a..ae0e8fe8c 100644
--- a/validmind/tests/model_validation/embeddings/CosineSimilarityHeatmap.py
+++ b/validmind/tests/model_validation/embeddings/CosineSimilarityHeatmap.py
@@ -81,4 +81,8 @@ def CosineSimilarityHeatmap(
         yaxis_title=yaxis_title,
     )
 
-    return fig, RawData(similarity_matrix=similarity_matrix)
+    return fig, RawData(
+        similarity_matrix=similarity_matrix,
+        model=model.input_id,
+        dataset=dataset.input_id,
+    )
diff --git a/validmind/tests/model_validation/embeddings/DescriptiveAnalytics.py b/validmind/tests/model_validation/embeddings/DescriptiveAnalytics.py
index f2fd3aeba..91a5ec851 100644
--- a/validmind/tests/model_validation/embeddings/DescriptiveAnalytics.py
+++ b/validmind/tests/model_validation/embeddings/DescriptiveAnalytics.py
@@ -77,5 +77,7 @@ def DescriptiveAnalytics(dataset: VMDataset, model: VMModel):
             embedding_means=embedding_means,
             embedding_medians=embedding_medians,
             embedding_stds=embedding_stds,
+            model=model.input_id,
+            dataset=dataset.input_id,
         ),
     )
diff --git a/validmind/tests/model_validation/embeddings/EmbeddingsVisualization2D.py b/validmind/tests/model_validation/embeddings/EmbeddingsVisualization2D.py
index 543b96d24..f2a254f45 100644
--- a/validmind/tests/model_validation/embeddings/EmbeddingsVisualization2D.py
+++ b/validmind/tests/model_validation/embeddings/EmbeddingsVisualization2D.py
@@ -89,4 +89,8 @@ def EmbeddingsVisualization2D(
     fig = px.scatter(**scatter_kwargs)
     fig.update_layout(width=500, height=500)
 
-    return fig, RawData(tsne_embeddings=reduced_embeddings)
+    return fig, RawData(
+        tsne_embeddings=reduced_embeddings,
+        model=model.input_id,
+        dataset=dataset.input_id,
+    )
diff --git a/validmind/tests/model_validation/embeddings/EuclideanDistanceComparison.py b/validmind/tests/model_validation/embeddings/EuclideanDistanceComparison.py
index 666f34c3f..eaca371c7 100644
--- a/validmind/tests/model_validation/embeddings/EuclideanDistanceComparison.py
+++ b/validmind/tests/model_validation/embeddings/EuclideanDistanceComparison.py
@@ -57,7 +57,7 @@ def EuclideanDistanceComparison(dataset, models):
     figures = []
     all_stats = []
 
-    distance_matrices = {}
+    distance_matrices = []
 
     # Generate all pairs of models for comparison
     for model_A, model_B in combinations(models, 2):
@@ -105,6 +105,10 @@ def EuclideanDistanceComparison(dataset, models):
     stats_df = pd.DataFrame(all_stats)
 
     # Add raw data to return
-    raw_data = RawData(distance_matrices=pd.DataFrame(distance_matrices))
+    raw_data = RawData(
+        distance_matrices=pd.DataFrame(distance_matrices),
+        dataset=dataset.input_id,
+        models=[model.input_id for model in models],
+    )
 
     return (stats_df, *figures, raw_data)
diff --git a/validmind/tests/model_validation/embeddings/EuclideanDistanceHeatmap.py b/validmind/tests/model_validation/embeddings/EuclideanDistanceHeatmap.py
index 430b96f42..da68b0764 100644
--- a/validmind/tests/model_validation/embeddings/EuclideanDistanceHeatmap.py
+++ b/validmind/tests/model_validation/embeddings/EuclideanDistanceHeatmap.py
@@ -79,4 +79,6 @@ def EuclideanDistanceHeatmap(
         yaxis_title=yaxis_title,
     )
 
-    return fig, RawData(distance_matrix=distance_matrix)
+    return fig, RawData(
+        distance_matrix=distance_matrix, model=model.input_id, dataset=dataset.input_id
+    )
diff --git a/validmind/tests/model_validation/embeddings/PCAComponentsPairwisePlots.py b/validmind/tests/model_validation/embeddings/PCAComponentsPairwisePlots.py
index 10a4308a9..a62df5d72 100644
--- a/validmind/tests/model_validation/embeddings/PCAComponentsPairwisePlots.py
+++ b/validmind/tests/model_validation/embeddings/PCAComponentsPairwisePlots.py
@@ -90,4 +90,7 @@ def PCAComponentsPairwisePlots(dataset, model, n_components=3):
         )
         figures.append(fig)
 
-    return (*figures, RawData(pca_results=pca_df))
+    return (
+        *figures,
+        RawData(pca_results=pca_df, model=model.input_id, dataset=dataset.input_id),
+    )
diff --git a/validmind/tests/model_validation/embeddings/StabilityAnalysisKeyword.py b/validmind/tests/model_validation/embeddings/StabilityAnalysisKeyword.py
index d3d0785ae..c39e92def 100644
--- a/validmind/tests/model_validation/embeddings/StabilityAnalysisKeyword.py
+++ b/validmind/tests/model_validation/embeddings/StabilityAnalysisKeyword.py
@@ -97,4 +97,8 @@ def perturb_data(data: str):
         mean_similarity_threshold,
     )
 
-    return results, RawData(original_perturbed_similarity=raw_data)
+    return results, RawData(
+        original_perturbed_similarity=raw_data,
+        model=model.input_id,
+        dataset=dataset.input_id,
+    )
diff --git a/validmind/tests/model_validation/embeddings/StabilityAnalysisRandomNoise.py b/validmind/tests/model_validation/embeddings/StabilityAnalysisRandomNoise.py
index 74efc4c0d..92e250755 100644
--- a/validmind/tests/model_validation/embeddings/StabilityAnalysisRandomNoise.py
+++ b/validmind/tests/model_validation/embeddings/StabilityAnalysisRandomNoise.py
@@ -151,4 +151,8 @@ def perturb_data(data):
         mean_similarity_threshold,
     )
 
-    return *result, RawData(original_perturbed_similarity=raw_data)
+    return *result, RawData(
+        original_perturbed_similarity=raw_data,
+        model=model.input_id,
+        dataset=dataset.input_id,
+    )
diff --git a/validmind/tests/model_validation/embeddings/StabilityAnalysisSynonyms.py b/validmind/tests/model_validation/embeddings/StabilityAnalysisSynonyms.py
index 6c644b80d..a8f14f081 100644
--- a/validmind/tests/model_validation/embeddings/StabilityAnalysisSynonyms.py
+++ b/validmind/tests/model_validation/embeddings/StabilityAnalysisSynonyms.py
@@ -107,4 +107,8 @@ def perturb_data(data):
         mean_similarity_threshold,
     )
 
-    return *result, RawData(original_perturbed_similarity=raw_data)
+    return *result, RawData(
+        original_perturbed_similarity=raw_data,
+        model=model.input_id,
+        dataset=dataset.input_id,
+    )
diff --git a/validmind/tests/model_validation/embeddings/StabilityAnalysisTranslation.py b/validmind/tests/model_validation/embeddings/StabilityAnalysisTranslation.py
index 77dc50b75..6dce9a58a 100644
--- a/validmind/tests/model_validation/embeddings/StabilityAnalysisTranslation.py
+++ b/validmind/tests/model_validation/embeddings/StabilityAnalysisTranslation.py
@@ -134,4 +134,8 @@ def perturb_data(data):
         mean_similarity_threshold,
     )
 
-    return *result, RawData(original_perturbed_similarity=raw_data)
+    return *result, RawData(
+        original_perturbed_similarity=raw_data,
+        model=model.input_id,
+        dataset=dataset.input_id,
+    )
diff --git a/validmind/tests/model_validation/embeddings/TSNEComponentsPairwisePlots.py b/validmind/tests/model_validation/embeddings/TSNEComponentsPairwisePlots.py
index 374049647..816004a0b 100644
--- a/validmind/tests/model_validation/embeddings/TSNEComponentsPairwisePlots.py
+++ b/validmind/tests/model_validation/embeddings/TSNEComponentsPairwisePlots.py
@@ -110,5 +110,10 @@ def TSNEComponentsPairwisePlots(
 
     return (
         *figures,
-        RawData(embeddings_scaled=embeddings_scaled, tsne_results=tsne_results),
+        RawData(
+            embeddings_scaled=embeddings_scaled,
+            tsne_results=tsne_results,
+            model=model.input_id,
+            dataset=dataset.input_id,
+        ),
     )
diff --git a/validmind/tests/model_validation/ragas/AnswerCorrectness.py b/validmind/tests/model_validation/ragas/AnswerCorrectness.py
index 38bed7282..e7fdc6309 100644
--- a/validmind/tests/model_validation/ragas/AnswerCorrectness.py
+++ b/validmind/tests/model_validation/ragas/AnswerCorrectness.py
@@ -144,5 +144,5 @@ def AnswerCorrectness(
         },
         fig_histogram,
         fig_box,
-        RawData(evaluation_results=result_df),
+        RawData(evaluation_results=result_df, dataset=dataset.input_id),
     )
diff --git a/validmind/tests/model_validation/ragas/AspectCritic.py b/validmind/tests/model_validation/ragas/AspectCritic.py
index 3a3ee097a..3f9858c39 100644
--- a/validmind/tests/model_validation/ragas/AspectCritic.py
+++ b/validmind/tests/model_validation/ragas/AspectCritic.py
@@ -195,5 +195,8 @@ def AspectCritic(
             ]
         },
         fig,
-        RawData(evaluation_results=result_df),
+        RawData(
+            evaluation_results=result_df,
+            dataset=dataset.input_id,
+        ),
     )
diff --git a/validmind/tests/model_validation/ragas/ContextEntityRecall.py b/validmind/tests/model_validation/ragas/ContextEntityRecall.py
index 52459877c..2c3516c70 100644
--- a/validmind/tests/model_validation/ragas/ContextEntityRecall.py
+++ b/validmind/tests/model_validation/ragas/ContextEntityRecall.py
@@ -143,5 +143,5 @@ def ContextEntityRecall(
         },
         fig_histogram,
         fig_box,
-        RawData(evaluation_results=result_df),
+        RawData(evaluation_results=result_df, dataset=dataset.input_id),
     )
diff --git a/validmind/tests/model_validation/ragas/ContextPrecision.py b/validmind/tests/model_validation/ragas/ContextPrecision.py
index 97e0fd74e..6be615425 100644
--- a/validmind/tests/model_validation/ragas/ContextPrecision.py
+++ b/validmind/tests/model_validation/ragas/ContextPrecision.py
@@ -135,5 +135,5 @@ def ContextPrecision(
         },
         fig_histogram,
         fig_box,
-        RawData(evaluation_results=result_df),
+        RawData(evaluation_results=result_df, dataset=dataset.input_id),
     )
diff --git a/validmind/tests/model_validation/ragas/ContextPrecisionWithoutReference.py b/validmind/tests/model_validation/ragas/ContextPrecisionWithoutReference.py
index 307e1c36e..916641589 100644
--- a/validmind/tests/model_validation/ragas/ContextPrecisionWithoutReference.py
+++ b/validmind/tests/model_validation/ragas/ContextPrecisionWithoutReference.py
@@ -130,5 +130,5 @@ def ContextPrecisionWithoutReference(
         },
         fig_histogram,
         fig_box,
-        RawData(evaluation_results=result_df),
+        RawData(evaluation_results=result_df, dataset=dataset.input_id),
     )
diff --git a/validmind/tests/model_validation/ragas/ContextRecall.py b/validmind/tests/model_validation/ragas/ContextRecall.py
index 55f8ad941..7503297b8 100644
--- a/validmind/tests/model_validation/ragas/ContextRecall.py
+++ b/validmind/tests/model_validation/ragas/ContextRecall.py
@@ -135,5 +135,5 @@ def ContextRecall(
         },
         fig_histogram,
         fig_box,
-        RawData(evaluation_results=result_df),
+        RawData(evaluation_results=result_df, dataset=dataset.input_id),
     )
diff --git a/validmind/tests/model_validation/ragas/Faithfulness.py b/validmind/tests/model_validation/ragas/Faithfulness.py
index d20dc5b05..989774bdf 100644
--- a/validmind/tests/model_validation/ragas/Faithfulness.py
+++ b/validmind/tests/model_validation/ragas/Faithfulness.py
@@ -140,5 +140,5 @@ def Faithfulness(
         },
         fig_histogram,
         fig_box,
-        RawData(evaluation_results=result_df),
+        RawData(evaluation_results=result_df, dataset=dataset.input_id),
     )
diff --git a/validmind/tests/model_validation/ragas/NoiseSensitivity.py b/validmind/tests/model_validation/ragas/NoiseSensitivity.py
index f68677a25..17cb838cd 100644
--- a/validmind/tests/model_validation/ragas/NoiseSensitivity.py
+++ b/validmind/tests/model_validation/ragas/NoiseSensitivity.py
@@ -179,5 +179,5 @@ def NoiseSensitivity(
         },
         fig_histogram,
         fig_box,
-        RawData(evaluation_results=result_df),
+        RawData(evaluation_results=result_df, dataset=dataset.input_id),
     )
diff --git a/validmind/tests/model_validation/ragas/ResponseRelevancy.py b/validmind/tests/model_validation/ragas/ResponseRelevancy.py
index 29f3fd041..6d6bb9f3b 100644
--- a/validmind/tests/model_validation/ragas/ResponseRelevancy.py
+++ b/validmind/tests/model_validation/ragas/ResponseRelevancy.py
@@ -154,5 +154,5 @@ def ResponseRelevancy(
         },
         fig_histogram,
         fig_box,
-        RawData(evaluation_results=result_df),
+        RawData(evaluation_results=result_df, dataset=dataset.input_id),
     )
diff --git a/validmind/tests/model_validation/ragas/SemanticSimilarity.py b/validmind/tests/model_validation/ragas/SemanticSimilarity.py
index a79cd0bf9..b4ca3049e 100644
--- a/validmind/tests/model_validation/ragas/SemanticSimilarity.py
+++ b/validmind/tests/model_validation/ragas/SemanticSimilarity.py
@@ -133,5 +133,5 @@ def SemanticSimilarity(
         },
         fig_histogram,
         fig_box,
-        RawData(evaluation_results=result_df),
+        RawData(evaluation_results=result_df, dataset=dataset.input_id),
     )
diff --git a/validmind/tests/model_validation/sklearn/AdjustedMutualInformation.py b/validmind/tests/model_validation/sklearn/AdjustedMutualInformation.py
index 71fa45394..5792ffdef 100644
--- a/validmind/tests/model_validation/sklearn/AdjustedMutualInformation.py
+++ b/validmind/tests/model_validation/sklearn/AdjustedMutualInformation.py
@@ -4,7 +4,7 @@
 
 from sklearn.metrics import adjusted_mutual_info_score
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
 
@@ -52,11 +52,11 @@ def AdjustedMutualInformation(model: VMModel, dataset: VMDataset):
     - The interpretability of the score can be complex as it depends on the understanding of information theory
     concepts.
     """
-    return [
-        {
-            "Adjusted Mutual Information": adjusted_mutual_info_score(
-                labels_true=dataset.y,
-                labels_pred=dataset.y_pred(model),
-            )
-        }
-    ]
+    ami_score = adjusted_mutual_info_score(
+        labels_true=dataset.y,
+        labels_pred=dataset.y_pred(model),
+    )
+
+    return [{"Adjusted Mutual Information": ami_score}], RawData(
+        ami_score=ami_score, model=model.input_id, dataset=dataset.input_id
+    )
diff --git a/validmind/tests/model_validation/sklearn/AdjustedRandIndex.py b/validmind/tests/model_validation/sklearn/AdjustedRandIndex.py
index 6fd3cb181..7d529d8d9 100644
--- a/validmind/tests/model_validation/sklearn/AdjustedRandIndex.py
+++ b/validmind/tests/model_validation/sklearn/AdjustedRandIndex.py
@@ -4,7 +4,7 @@
 
 from sklearn.metrics import adjusted_rand_score
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
 
@@ -49,11 +49,11 @@ def AdjustedRandIndex(model: VMModel, dataset: VMDataset):
     - It may be difficult to interpret the implications of an ARI score without context or a benchmark, as it is
     heavily dependent on the characteristics of the dataset used.
     """
-    return [
-        {
-            "Adjusted Rand Index": adjusted_rand_score(
-                labels_true=dataset.y,
-                labels_pred=dataset.y_pred(model),
-            )
-        }
-    ]
+    ari = adjusted_rand_score(
+        labels_true=dataset.y,
+        labels_pred=dataset.y_pred(model),
+    )
+
+    return [{"Adjusted Rand Index": ari}], RawData(
+        ari_score=ari, model=model.input_id, dataset=dataset.input_id
+    )
diff --git a/validmind/tests/model_validation/sklearn/CalibrationCurve.py b/validmind/tests/model_validation/sklearn/CalibrationCurve.py
index 0792f62a5..90bfd5bf8 100644
--- a/validmind/tests/model_validation/sklearn/CalibrationCurve.py
+++ b/validmind/tests/model_validation/sklearn/CalibrationCurve.py
@@ -72,7 +72,10 @@ def CalibrationCurve(model: VMModel, dataset: VMDataset, n_bins: int = 10):
 
     # Create DataFrame for raw data
     raw_data = RawData(
-        mean_predicted_probability=prob_pred, observed_frequency=prob_true
+        mean_predicted_probability=prob_pred,
+        observed_frequency=prob_true,
+        model=model.input_id,
+        dataset=dataset.input_id,
     )
 
     # Create Plotly figure
@@ -114,4 +117,4 @@ def CalibrationCurve(model: VMModel, dataset: VMDataset, n_bins: int = 10):
         template="plotly_white",
     )
 
-    return raw_data, fig
+    return fig, raw_data
diff --git a/validmind/tests/model_validation/sklearn/ClassifierThresholdOptimization.py b/validmind/tests/model_validation/sklearn/ClassifierThresholdOptimization.py
index 0a4d4f442..a8c96c72f 100644
--- a/validmind/tests/model_validation/sklearn/ClassifierThresholdOptimization.py
+++ b/validmind/tests/model_validation/sklearn/ClassifierThresholdOptimization.py
@@ -8,7 +8,7 @@
 from plotly.subplots import make_subplots
 from sklearn.metrics import confusion_matrix, precision_recall_curve, roc_curve
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
 
@@ -255,4 +255,17 @@ def ClassifierThresholdOptimization(
     # Create results table and sort by threshold descending
     table = pd.DataFrame(results).sort_values("threshold", ascending=False)
 
-    return fig, table
+    return (
+        fig,
+        table,
+        RawData(
+            fpr=fpr,
+            tpr=tpr,
+            precision=precision,
+            recall=recall,
+            thresholds_roc=thresholds_roc,
+            thresholds_pr=thresholds_pr,
+            model=model.input_id,
+            dataset=dataset.input_id,
+        ),
+    )
diff --git a/validmind/tests/model_validation/sklearn/ClusterCosineSimilarity.py b/validmind/tests/model_validation/sklearn/ClusterCosineSimilarity.py
index 999625539..6f393eaef 100644
--- a/validmind/tests/model_validation/sklearn/ClusterCosineSimilarity.py
+++ b/validmind/tests/model_validation/sklearn/ClusterCosineSimilarity.py
@@ -84,4 +84,8 @@ def ClusterCosineSimilarity(model: VMModel, dataset: VMDataset):
     if not table:
         raise SkipTestError("No clusters found")
 
-    return table, RawData(cluster_centroids=cluster_centroids)
+    return table, RawData(
+        cluster_centroids=cluster_centroids,
+        model=model.input_id,
+        dataset=dataset.input_id,
+    )
diff --git a/validmind/tests/model_validation/sklearn/ClusterPerformanceMetrics.py b/validmind/tests/model_validation/sklearn/ClusterPerformanceMetrics.py
index e3d23ffe3..ccb737be0 100644
--- a/validmind/tests/model_validation/sklearn/ClusterPerformanceMetrics.py
+++ b/validmind/tests/model_validation/sklearn/ClusterPerformanceMetrics.py
@@ -11,7 +11,7 @@
     v_measure_score,
 )
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
 HOMOGENEITY = """
@@ -115,53 +115,63 @@ def ClusterPerformanceMetrics(model: VMModel, dataset: VMDataset):
     - Does not consider aspects like computational efficiency of the model or its capability to handle high dimensional
     data.
     """
-    return [
+    y_true = dataset.y
+    y_pred = dataset.y_pred(model)
+
+    metrics = [
         {
             "Metric": "Homogeneity Score",
             "Description": HOMOGENEITY,
             "Value": homogeneity_score(
-                labels_true=dataset.y,
-                labels_pred=dataset.y_pred(model),
+                labels_true=y_true,
+                labels_pred=y_pred,
             ),
         },
         {
             "Metric": "Completeness Score",
             "Description": COMPLETENESS,
             "Value": completeness_score(
-                labels_true=dataset.y,
-                labels_pred=dataset.y_pred(model),
+                labels_true=y_true,
+                labels_pred=y_pred,
             ),
         },
         {
             "Metric": "V Measure",
             "Description": V_MEASURE,
             "Value": v_measure_score(
-                labels_true=dataset.y,
-                labels_pred=dataset.y_pred(model),
+                labels_true=y_true,
+                labels_pred=y_pred,
             ),
         },
         {
             "Metric": "Adjusted Rand Index",
             "Description": ADJUSTED_RAND_INDEX,
             "Value": adjusted_rand_score(
-                labels_true=dataset.y,
-                labels_pred=dataset.y_pred(model),
+                labels_true=y_true,
+                labels_pred=y_pred,
             ),
         },
         {
             "Metric": "Adjusted Mutual Information",
             "Description": ADJUSTED_MUTUAL_INFORMATION,
             "Value": adjusted_mutual_info_score(
-                labels_true=dataset.y,
-                labels_pred=dataset.y_pred(model),
+                labels_true=y_true,
+                labels_pred=y_pred,
             ),
         },
         {
             "Metric": "Fowlkes-Mallows score",
             "Description": FOULKES_MALLOWS_SCORE,
             "Value": fowlkes_mallows_score(
-                labels_true=dataset.y,
-                labels_pred=dataset.y_pred(model),
+                labels_true=y_true,
+                labels_pred=y_pred,
             ),
         },
     ]
+
+    return metrics, RawData(
+        true_labels=y_true,
+        predicted_labels=y_pred,
+        model=model.input_id,
+        dataset=dataset.input_id,
+    )
diff --git a/validmind/tests/model_validation/sklearn/CompletenessScore.py b/validmind/tests/model_validation/sklearn/CompletenessScore.py
index 26272822e..ebf990ee4 100644
--- a/validmind/tests/model_validation/sklearn/CompletenessScore.py
+++ b/validmind/tests/model_validation/sklearn/CompletenessScore.py
@@ -4,7 +4,7 @@
 
 from sklearn.metrics import completeness_score
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
 
@@ -47,11 +47,10 @@ def CompletenessScore(model: VMModel, dataset: VMDataset):
     - The Completeness Score only applies to clustering models; it cannot be used for other types of machine learning
     models.
     """
-    return [
-        {
-            "Completeness Score": completeness_score(
-                labels_true=dataset.y,
-                labels_pred=dataset.y_pred(model),
-            )
-        }
-    ]
+    score = completeness_score(
+        labels_true=dataset.y,
+        labels_pred=dataset.y_pred(model),
+    )
+    return [{"Completeness Score": score}], RawData(
+        score=score, model=model.input_id, dataset=dataset.input_id
+    )
diff --git a/validmind/tests/model_validation/sklearn/FeatureImportance.py b/validmind/tests/model_validation/sklearn/FeatureImportance.py
index 61bcb4a01..8f3ed77ca 100644
--- a/validmind/tests/model_validation/sklearn/FeatureImportance.py
+++ b/validmind/tests/model_validation/sklearn/FeatureImportance.py
@@ -5,7 +5,7 @@
 import pandas as pd
 from sklearn.inspection import permutation_importance
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
 
@@ -78,4 +78,8 @@ def FeatureImportance(dataset: VMDataset, model: VMModel, num_features: int = 3)
         else:
             result[f"Feature {i + 1}"] = None
 
-    return pd.DataFrame([result])
+    return pd.DataFrame([result]), RawData(
+        permutation_importance=pfi_values,
+        model=model.input_id,
+        dataset=dataset.input_id,
+    )
diff --git a/validmind/tests/model_validation/sklearn/FowlkesMallowsScore.py b/validmind/tests/model_validation/sklearn/FowlkesMallowsScore.py
index 04b9b08b3..cf314fe83 100644
--- a/validmind/tests/model_validation/sklearn/FowlkesMallowsScore.py
+++ b/validmind/tests/model_validation/sklearn/FowlkesMallowsScore.py
@@ -4,7 +4,7 @@
 
 from sklearn import metrics
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
 
@@ -52,11 +52,14 @@ def FowlkesMallowsScore(dataset: VMDataset, model: VMModel):
     - It does not handle mismatching numbers of clusters between the true and predicted labels. As such, it may return
     misleading results if the predicted labels suggest a different number of clusters than what is in the true labels.
     """
-    return [
-        {
-            "Fowlkes-Mallows score": metrics.fowlkes_mallows_score(
-                labels_true=dataset.y,
-                labels_pred=dataset.y_pred(model),
-            )
-        }
-    ]
+    fowlkes_mallows_score = metrics.fowlkes_mallows_score(
+        labels_true=dataset.y,
+        labels_pred=dataset.y_pred(model),
+    )
+
+    return [{"Fowlkes-Mallows score": fowlkes_mallows_score}], RawData(
+        labels_true=dataset.y,
+        labels_pred=dataset.y_pred(model),
+        model=model.input_id,
+        dataset=dataset.input_id,
+    )
diff --git a/validmind/tests/model_validation/sklearn/HomogeneityScore.py b/validmind/tests/model_validation/sklearn/HomogeneityScore.py
index 1fedd83bd..739d6b086 100644
--- a/validmind/tests/model_validation/sklearn/HomogeneityScore.py
+++ b/validmind/tests/model_validation/sklearn/HomogeneityScore.py
@@ -4,7 +4,7 @@
 
 from sklearn import metrics
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
 
@@ -50,11 +50,16 @@ class labels of the training and testing sets with the labels predicted by the g
     - The score does not address the actual number of clusters formed, or the evenness of cluster sizes. It only checks
     the homogeneity within the given clusters created by the model.
     """
-    return [
-        {
-            "Homogeneity Score": metrics.homogeneity_score(
-                labels_true=dataset.y,
-                labels_pred=dataset.y_pred(model),
-            )
-        }
-    ]
+    homogeneity_score = metrics.homogeneity_score(
+        labels_true=dataset.y,
+        labels_pred=dataset.y_pred(model),
+    )
+
+    raw_data = RawData(
+        y_true=dataset.y,
+        y_pred=dataset.y_pred(model),
+        model=model.input_id,
+        dataset=dataset.input_id,
+    )
+
+    return ([{"Homogeneity Score": homogeneity_score}], raw_data)
diff --git a/validmind/tests/model_validation/sklearn/HyperParametersTuning.py b/validmind/tests/model_validation/sklearn/HyperParametersTuning.py
index dd90a44af..c12d12b29 100644
--- a/validmind/tests/model_validation/sklearn/HyperParametersTuning.py
+++ b/validmind/tests/model_validation/sklearn/HyperParametersTuning.py
@@ -7,7 +7,7 @@
 from sklearn.metrics import make_scorer, recall_score
 from sklearn.model_selection import GridSearchCV
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
 
@@ -162,4 +162,6 @@ def HyperParametersTuning(
 
             results.append(row_result)
 
-    return results
+    return results, RawData(
+        model=model.input_id, dataset=dataset.input_id, param_grid=param_grid
+    )
diff --git a/validmind/tests/model_validation/sklearn/KMeansClustersOptimization.py b/validmind/tests/model_validation/sklearn/KMeansClustersOptimization.py
index f26bddb6a..485f93b34 100644
--- a/validmind/tests/model_validation/sklearn/KMeansClustersOptimization.py
+++ b/validmind/tests/model_validation/sklearn/KMeansClustersOptimization.py
@@ -124,4 +124,9 @@ def KMeansClustersOptimization(
 
     fig.update_layout(showlegend=False)
 
-    return fig, RawData(distortions=distortions, silhouette_avg=silhouette_avg)
+    return fig, RawData(
+        distortions=distortions,
+        silhouette_avg=silhouette_avg,
+        model=model.input_id,
+        dataset=dataset.input_id,
+    )
diff --git a/validmind/tests/model_validation/sklearn/MinimumAccuracy.py b/validmind/tests/model_validation/sklearn/MinimumAccuracy.py
index 330e7ff53..7e470cfc1 100644
--- a/validmind/tests/model_validation/sklearn/MinimumAccuracy.py
+++ b/validmind/tests/model_validation/sklearn/MinimumAccuracy.py
@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 from sklearn.metrics import accuracy_score
 
+from validmind import RawData
 from validmind.tests import tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
@@ -50,10 +51,14 @@ def MinimumAccuracy(dataset: VMDataset, model: VMModel, min_threshold: float = 0
     """
     accuracy = accuracy_score(dataset.y, dataset.y_pred(model))
 
-    return [
-        {
-            "Score": accuracy,
-            "Threshold": min_threshold,
-            "Pass/Fail": "Pass" if accuracy > min_threshold else "Fail",
-        }
-    ], accuracy > min_threshold
+    return (
+        [
+            {
+                "Score": accuracy,
+                "Threshold": min_threshold,
+                "Pass/Fail": "Pass" if accuracy > min_threshold else "Fail",
+            }
+        ],
+        accuracy > min_threshold,
+        RawData(model=model.input_id, dataset=dataset.input_id),
+    )
diff --git a/validmind/tests/model_validation/sklearn/MinimumF1Score.py b/validmind/tests/model_validation/sklearn/MinimumF1Score.py
index 9e2995453..f3894d70c 100644
--- a/validmind/tests/model_validation/sklearn/MinimumF1Score.py
+++ b/validmind/tests/model_validation/sklearn/MinimumF1Score.py
@@ -5,6 +5,7 @@
 import numpy as np
 from sklearn.metrics import f1_score
 
+from validmind import RawData
 from validmind.tests import tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
@@ -58,10 +59,14 @@ def MinimumF1Score(dataset: VMDataset, model: VMModel, min_threshold: float = 0.
     else:
         score = f1_score(dataset.y, dataset.y_pred(model))
 
-    return [
-        {
-            "Score": score,
-            "Threshold": min_threshold,
-            "Pass/Fail": "Pass" if score > min_threshold else "Fail",
-        }
-    ], score > min_threshold
+    return (
+        [
+            {
+                "Score": score,
+                "Threshold": min_threshold,
+                "Pass/Fail": "Pass" if score > min_threshold else "Fail",
+            }
+        ],
+        score > min_threshold,
+        RawData(score=score, model=model.input_id, dataset=dataset.input_id),
+    )
diff --git a/validmind/tests/model_validation/sklearn/MinimumROCAUCScore.py b/validmind/tests/model_validation/sklearn/MinimumROCAUCScore.py
index b6ed15204..61b945f10 100644
--- a/validmind/tests/model_validation/sklearn/MinimumROCAUCScore.py
+++ b/validmind/tests/model_validation/sklearn/MinimumROCAUCScore.py
@@ -6,7 +6,7 @@
 from sklearn.metrics import roc_auc_score
 from sklearn.preprocessing import LabelBinarizer
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
 
@@ -62,19 +62,34 @@ def MinimumROCAUCScore(dataset: VMDataset, model: VMModel, min_threshold: float
         lb = LabelBinarizer()
         lb.fit(y_true)
 
+        y_true_binary = lb.transform(y_true)
+        y_score_binary = lb.transform(dataset.y_pred(model))
+
         roc_auc = roc_auc_score(
-            y_true=lb.transform(y_true),
-            y_score=lb.transform(dataset.y_pred(model)),
+            y_true=y_true_binary,
+            y_score=y_score_binary,
             average="macro",
         )
 
     else:
-        roc_auc = roc_auc_score(y_true=y_true, y_score=dataset.y_prob(model))
+        y_score_prob = dataset.y_prob(model)
+        roc_auc = roc_auc_score(y_true=y_true, y_score=y_score_prob)
 
-    return [
+    results = [
         {
             "Score": roc_auc,
             "Threshold": min_threshold,
             "Pass/Fail": "Pass" if roc_auc > min_threshold else "Fail",
         }
-    ], roc_auc > min_threshold
+    ]
+
+    return (
+        results,
+        roc_auc > min_threshold,
+        RawData(
+            y_true=y_true,
+            roc_auc=roc_auc,
+            model=model.input_id,
+            dataset=dataset.input_id,
+        ),
+    )
diff --git a/validmind/tests/model_validation/sklearn/OverfitDiagnosis.py b/validmind/tests/model_validation/sklearn/OverfitDiagnosis.py
index 28c802e09..759024718 100644
--- a/validmind/tests/model_validation/sklearn/OverfitDiagnosis.py
+++ b/validmind/tests/model_validation/sklearn/OverfitDiagnosis.py
@@ -10,7 +10,7 @@
 import seaborn as sns
 from sklearn import metrics
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.logging import get_logger
 from validmind.vm_models import VMDataset, VMModel
 
@@ -299,4 +299,10 @@ def OverfitDiagnosis(
                 }
             )
 
-    return ({"Overfit Diagnosis": test_results}, *figures)
+    return (
+        {"Overfit Diagnosis": test_results},
+        *figures,
+        RawData(
+            model=model.input_id, datasets=[dataset.input_id for dataset in datasets]
+        ),
+    )
diff --git a/validmind/tests/model_validation/sklearn/PermutationFeatureImportance.py b/validmind/tests/model_validation/sklearn/PermutationFeatureImportance.py
index f65f8955e..3f1d0d196 100644
--- a/validmind/tests/model_validation/sklearn/PermutationFeatureImportance.py
+++ b/validmind/tests/model_validation/sklearn/PermutationFeatureImportance.py
@@ -111,4 +111,8 @@ def PermutationFeatureImportance(
         height=figure_height,
     )
 
-    return fig, RawData(permutation_importance=pfi_values)
+    return fig, RawData(
+        permutation_importance=pfi_values,
+        model=model.input_id,
+        dataset=dataset.input_id,
+    )
diff --git a/validmind/tests/model_validation/sklearn/PopulationStabilityIndex.py b/validmind/tests/model_validation/sklearn/PopulationStabilityIndex.py
index c6be015f5..3c3ac8ee6 100644
--- a/validmind/tests/model_validation/sklearn/PopulationStabilityIndex.py
+++ b/validmind/tests/model_validation/sklearn/PopulationStabilityIndex.py
@@ -209,5 +209,9 @@ def PopulationStabilityIndex(
             ],
         },
         fig,
-        RawData(psi_raw=psi_results),
+        RawData(
+            psi_raw=psi_results,
+            model=model.input_id,
+            datasets=[datasets[0].input_id, datasets[1].input_id],
+        ),
     )
diff --git a/validmind/tests/model_validation/sklearn/PrecisionRecallCurve.py b/validmind/tests/model_validation/sklearn/PrecisionRecallCurve.py
index 8a722fa8b..b521a333a 100644
--- a/validmind/tests/model_validation/sklearn/PrecisionRecallCurve.py
+++ b/validmind/tests/model_validation/sklearn/PrecisionRecallCurve.py
@@ -83,4 +83,9 @@ def PrecisionRecallCurve(model: VMModel, dataset: VMDataset):
         ),
     )
 
-    return fig, RawData(precision=precision, recall=recall)
+    return fig, RawData(
+        precision=precision,
+        recall=recall,
+        model=model.input_id,
+        dataset=dataset.input_id,
+    )
diff --git a/validmind/tests/model_validation/sklearn/ROCCurve.py b/validmind/tests/model_validation/sklearn/ROCCurve.py
index cb7314953..9f3161931 100644
--- a/validmind/tests/model_validation/sklearn/ROCCurve.py
+++ b/validmind/tests/model_validation/sklearn/ROCCurve.py
@@ -103,5 +103,7 @@ def ROCCurve(model: VMModel, dataset: VMDataset):
                 height=500,
             ),
         ),
-        RawData(fpr=fpr, tpr=tpr, auc=auc),
+        RawData(
+            fpr=fpr, tpr=tpr, auc=auc, model=model.input_id, dataset=dataset.input_id
+        ),
     )
diff --git a/validmind/tests/model_validation/sklearn/RegressionErrors.py b/validmind/tests/model_validation/sklearn/RegressionErrors.py
index f398d6d21..f10ce3a48 100644
--- a/validmind/tests/model_validation/sklearn/RegressionErrors.py
+++ b/validmind/tests/model_validation/sklearn/RegressionErrors.py
@@ -6,7 +6,7 @@
 import pandas as pd
 from sklearn import metrics
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 
 
 @tags("sklearn", "model_performance")
@@ -62,7 +62,11 @@ def RegressionErrors(model, dataset):
     y_pred = dataset.y_pred(model)
     y_true = y_true.astype(y_pred.dtype)
 
-    return _regression_errors(y_true, y_pred)
+    results_df = _regression_errors(y_true, y_pred)
+
+    return results_df, RawData(
+        y_true=y_true, y_pred=y_pred, model=model.input_id, dataset=dataset.input_id
+    )
 
 
 def _regression_errors(y_true, y_pred):
diff --git a/validmind/tests/model_validation/sklearn/RegressionPerformance.py b/validmind/tests/model_validation/sklearn/RegressionPerformance.py
index 2ffdbd6cf..2c1eed0e4 100644
--- a/validmind/tests/model_validation/sklearn/RegressionPerformance.py
+++ b/validmind/tests/model_validation/sklearn/RegressionPerformance.py
@@ -5,7 +5,7 @@
 import numpy as np
 from sklearn.metrics import mean_absolute_error, mean_squared_error
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.logging import get_logger
 from validmind.vm_models import VMDataset, VMModel
 
@@ -74,10 +74,15 @@ def RegressionPerformance(model: VMModel, dataset: VMDataset):
     # MBD calculation
     metrics["Mean Bias Deviation (MBD)"] = np.mean(y_pred - y_true)
 
-    return [
-        {
-            "Metric": metric,
-            "Value": value,
-        }
-        for metric, value in metrics.items()
-    ]
+    return (
+        [
+            {
+                "Metric": metric,
+                "Value": value,
+            }
+            for metric, value in metrics.items()
+        ],
+        RawData(
+            y_true=y_true, y_pred=y_pred, model=model.input_id, dataset=dataset.input_id
+        ),
+    )
diff --git a/validmind/tests/model_validation/sklearn/RegressionR2Square.py b/validmind/tests/model_validation/sklearn/RegressionR2Square.py
index 6f6dc6bfc..302184e6d 100644
--- a/validmind/tests/model_validation/sklearn/RegressionR2Square.py
+++ b/validmind/tests/model_validation/sklearn/RegressionR2Square.py
@@ -5,7 +5,7 @@
 import pandas as pd
 from sklearn import metrics
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.tests.model_validation.statsmodels.statsutils import adj_r2_score
 
 
@@ -55,11 +55,14 @@ def RegressionR2Square(dataset, model):
     y_pred = dataset.y_pred(model)
     y_true = y_true.astype(y_pred.dtype)
 
+    r2 = metrics.r2_score(y_true, y_pred)
+    adj_r2 = adj_r2_score(y_true, y_pred, len(y_true), len(dataset.feature_columns))
+
     return pd.DataFrame(
         {
-            "R-squared (R2) Score": [metrics.r2_score(y_true, y_pred)],
-            "Adjusted R-squared (R2) Score": [
-                adj_r2_score(y_true, y_pred, len(y_true), len(dataset.feature_columns))
-            ],
+            "R-squared (R2) Score": [r2],
+            "Adjusted R-squared (R2) Score": [adj_r2],
         }
+    ), RawData(
+        y_true=y_true, y_pred=y_pred, model=model.input_id, dataset=dataset.input_id
     )
diff --git a/validmind/tests/model_validation/sklearn/RobustnessDiagnosis.py b/validmind/tests/model_validation/sklearn/RobustnessDiagnosis.py
index d42007ed7..4b0710211 100644
--- a/validmind/tests/model_validation/sklearn/RobustnessDiagnosis.py
+++ b/validmind/tests/model_validation/sklearn/RobustnessDiagnosis.py
@@ -14,7 +14,7 @@
 
 from validmind.errors import MissingOrInvalidModelPredictFnError
 from validmind.logging import get_logger
-from validmind.tests import tags, tasks
+from validmind.tests import RawData, tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
 logger = get_logger(__name__)
@@ -327,4 +327,11 @@ def RobustnessDiagnosis(
         results_df["Perturbation Size"] == 0.0, "Perturbation Size"
     ] = "Baseline (0.0)"
 
-    return results_df, fig, all(results_df["Passed"])
+    return (
+        results_df,
+        fig,
+        all(results_df["Passed"]),
+        RawData(
+            model=model.input_id, datasets=[dataset.input_id for dataset in datasets]
+        ),
+    )
diff --git a/validmind/tests/model_validation/sklearn/SHAPGlobalImportance.py b/validmind/tests/model_validation/sklearn/SHAPGlobalImportance.py
index a7e162b83..56165fdf6 100644
--- a/validmind/tests/model_validation/sklearn/SHAPGlobalImportance.py
+++ b/validmind/tests/model_validation/sklearn/SHAPGlobalImportance.py
@@ -229,5 +229,10 @@ def SHAPGlobalImportance(
     return (
         generate_shap_plot("mean", shap_values, shap_sample),
         generate_shap_plot("summary", shap_values, shap_sample),
-        RawData(shap_values=shap_values, shap_sample=shap_sample),
+        RawData(
+            shap_values=shap_values,
+            shap_sample=shap_sample,
+            model=model.input_id,
+            dataset=dataset.input_id,
+        ),
     )
diff --git a/validmind/tests/model_validation/sklearn/ScoreProbabilityAlignment.py b/validmind/tests/model_validation/sklearn/ScoreProbabilityAlignment.py
index 7246ca573..6c2d6c08b 100644
--- a/validmind/tests/model_validation/sklearn/ScoreProbabilityAlignment.py
+++ b/validmind/tests/model_validation/sklearn/ScoreProbabilityAlignment.py
@@ -5,7 +5,7 @@
 import pandas as pd
 import plotly.graph_objects as go
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
 
@@ -128,4 +128,12 @@ def ScoreProbabilityAlignment(
         height=600,
     )
 
-    return results_df, fig
+    # Include raw data for post-processing
+    raw_data = RawData(
+        score_bins=df[["score_bin", score_column]],
+        predicted_probabilities=df["probability"],
+        model=model.input_id,
+        dataset=dataset.input_id,
+    )
+
+    return results_df, fig, raw_data
diff --git a/validmind/tests/model_validation/sklearn/SilhouettePlot.py b/validmind/tests/model_validation/sklearn/SilhouettePlot.py
index 836c8791b..2247a33ac 100644
--- a/validmind/tests/model_validation/sklearn/SilhouettePlot.py
+++ b/validmind/tests/model_validation/sklearn/SilhouettePlot.py
@@ -110,5 +110,9 @@ def SilhouettePlot(model: VMModel, dataset: VMDataset):
             "Silhouette Score": silhouette_avg,
         },
         fig,
-        RawData(sample_silhouette_values=sample_silhouette_values),
+        RawData(
+            sample_silhouette_values=sample_silhouette_values,
+            model=model.input_id,
+            dataset=dataset.input_id,
+        ),
     )
diff --git a/validmind/tests/model_validation/sklearn/VMeasure.py b/validmind/tests/model_validation/sklearn/VMeasure.py
index 2a86c7390..9153ed392 100644
--- a/validmind/tests/model_validation/sklearn/VMeasure.py
+++ b/validmind/tests/model_validation/sklearn/VMeasure.py
@@ -4,7 +4,7 @@
 
 from sklearn import metrics
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
 
@@ -48,11 +48,14 @@ def VMeasure(dataset: VMDataset, model: VMModel):
     the other. The V Measure Score does not provide flexibility in assigning different weights to homogeneity and
     completeness.
     """
-    return [
-        {
-            "V Measure": metrics.v_measure_score(
-                labels_true=dataset.y,
-                labels_pred=dataset.y_pred(model),
-            )
-        }
-    ]
+    v_measure = metrics.v_measure_score(
+        labels_true=dataset.y,
+        labels_pred=dataset.y_pred(model),
+    )
+
+    return (
+        [{"V Measure": v_measure}],
+        RawData(
+            v_measure_score=v_measure, model=model.input_id, dataset=dataset.input_id
+        ),
+    )
diff --git a/validmind/tests/model_validation/statsmodels/CumulativePredictionProbabilities.py b/validmind/tests/model_validation/statsmodels/CumulativePredictionProbabilities.py
index 4f069acc1..733787529 100644
--- a/validmind/tests/model_validation/statsmodels/CumulativePredictionProbabilities.py
+++ b/validmind/tests/model_validation/statsmodels/CumulativePredictionProbabilities.py
@@ -64,7 +64,11 @@ def CumulativePredictionProbabilities(dataset, model, title="Cumulative Probabil
 
     fig, fig_data = _plot_cumulative_prob(df, dataset.target_column, title)
 
-    return fig, RawData(cumulative_probabilities=fig_data)
+    return fig, RawData(
+        cumulative_probabilities=fig_data,
+        model=model.input_id,
+        dataset=dataset.input_id,
+    )
 
 
 def _plot_cumulative_prob(df, target_col, title):
diff --git a/validmind/tests/model_validation/statsmodels/DurbinWatsonTest.py b/validmind/tests/model_validation/statsmodels/DurbinWatsonTest.py
index 759e15553..02147456b 100644
--- a/validmind/tests/model_validation/statsmodels/DurbinWatsonTest.py
+++ b/validmind/tests/model_validation/statsmodels/DurbinWatsonTest.py
@@ -5,7 +5,7 @@
 import pandas as pd
 from statsmodels.stats.stattools import durbin_watson
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 
 
 @tasks("regression")
@@ -81,4 +81,9 @@ def get_autocorrelation(dw_value, threshold):
             "threshold": [str(threshold)],
             "autocorrelation": [get_autocorrelation(dw_statistic, threshold)],
         }
+    ), RawData(
+        residuals=residuals,
+        dw_statistic=dw_statistic,
+        model=model.input_id,
+        dataset=dataset.input_id,
     )
diff --git a/validmind/tests/model_validation/statsmodels/GINITable.py b/validmind/tests/model_validation/statsmodels/GINITable.py
index 5e90dd8d3..39a8815c3 100644
--- a/validmind/tests/model_validation/statsmodels/GINITable.py
+++ b/validmind/tests/model_validation/statsmodels/GINITable.py
@@ -6,7 +6,7 @@
 import pandas as pd
 from sklearn.metrics import roc_auc_score, roc_curve
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 
 
 @tags("model_performance")
@@ -76,4 +76,11 @@ def GINITable(dataset, model):
             "GINI": [gini],
             "KS": [max(tpr - fpr)],
         }
+    ), RawData(
+        fpr=fpr,
+        tpr=tpr,
+        y_true=y_true,
+        y_prob=y_prob,
+        model=model.input_id,
+        dataset=dataset.input_id,
     )
diff --git a/validmind/tests/model_validation/statsmodels/KolmogorovSmirnov.py b/validmind/tests/model_validation/statsmodels/KolmogorovSmirnov.py
index f4042d2c7..61f76390b 100644
--- a/validmind/tests/model_validation/statsmodels/KolmogorovSmirnov.py
+++ b/validmind/tests/model_validation/statsmodels/KolmogorovSmirnov.py
@@ -4,7 +4,7 @@
 
 from statsmodels.stats.diagnostic import kstest_normal
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.errors import InvalidTestParametersError
 from validmind.vm_models import VMDataset, VMModel
 
@@ -66,4 +66,4 @@ def KolmogorovSmirnov(model: VMModel, dataset: VMDataset, dist: str = "norm"):
             "P-Value": result["pvalue"],
         }
         for k, result in ks_values.items()
-    ]
+    ], RawData(ks_values=ks_values, dataset=dataset.input_id)
diff --git a/validmind/tests/model_validation/statsmodels/PredictionProbabilitiesHistogram.py b/validmind/tests/model_validation/statsmodels/PredictionProbabilitiesHistogram.py
index 964ab5aca..05887165d 100644
--- a/validmind/tests/model_validation/statsmodels/PredictionProbabilitiesHistogram.py
+++ b/validmind/tests/model_validation/statsmodels/PredictionProbabilitiesHistogram.py
@@ -6,7 +6,7 @@
 import plotly.graph_objects as go
 from matplotlib import cm
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 
 
 @tags("visualization", "credit_risk")
@@ -60,7 +60,11 @@ def PredictionProbabilitiesHistogram(
 
     fig = _plot_prob_histogram(df, dataset.target_column, title)
 
-    return fig
+    return fig, RawData(
+        probabilities=df["probabilities"],
+        model=model.input_id,
+        dataset=dataset.input_id,
+    )
 
 
 def _plot_prob_histogram(df, target_col, title):
diff --git a/validmind/tests/model_validation/statsmodels/RegressionCoeffs.py b/validmind/tests/model_validation/statsmodels/RegressionCoeffs.py
index 70f92b3ac..aee3cbe90 100644
--- a/validmind/tests/model_validation/statsmodels/RegressionCoeffs.py
+++ b/validmind/tests/model_validation/statsmodels/RegressionCoeffs.py
@@ -7,7 +7,7 @@
 import plotly.graph_objects as go
 from scipy import stats
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.errors import SkipTestError
 
 
@@ -97,4 +97,10 @@ def RegressionCoeffs(model):
         yaxis_title="Coefficients",
     )
 
-    return fig, coefficients
+    return (
+        fig,
+        coefficients,
+        RawData(
+            model=model.input_id, std_err=std_err, lower_ci=lower_ci, upper_ci=upper_ci
+        ),
+    )
diff --git a/validmind/tests/model_validation/statsmodels/RegressionFeatureSignificance.py b/validmind/tests/model_validation/statsmodels/RegressionFeatureSignificance.py
index 3094d51c7..3d017d1c6 100644
--- a/validmind/tests/model_validation/statsmodels/RegressionFeatureSignificance.py
+++ b/validmind/tests/model_validation/statsmodels/RegressionFeatureSignificance.py
@@ -90,4 +90,6 @@ def RegressionFeatureSignificance(
 
     plt.close()
 
-    return fig, RawData(coefficients=coefficients, pvalues=pvalues)
+    return fig, RawData(
+        coefficients=coefficients, pvalues=pvalues, model=model.input_id
+    )
diff --git a/validmind/tests/model_validation/statsmodels/RegressionModelForecastPlot.py b/validmind/tests/model_validation/statsmodels/RegressionModelForecastPlot.py
index 820202f16..95bd0fdd4 100644
--- a/validmind/tests/model_validation/statsmodels/RegressionModelForecastPlot.py
+++ b/validmind/tests/model_validation/statsmodels/RegressionModelForecastPlot.py
@@ -7,7 +7,7 @@
 import matplotlib.pyplot as plt
 import pandas as pd
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.logging import get_logger
 from validmind.vm_models import VMDataset, VMModel
 
@@ -87,4 +87,9 @@ def RegressionModelForecastPlot(
 
     plt.close()
 
-    return fig
+    return fig, RawData(
+        observed_values=dataset.y.tolist(),
+        forecast_values=dataset.y_pred(model).tolist(),
+        model=model.input_id,
+        dataset=dataset.input_id,
+    )
diff --git a/validmind/tests/model_validation/statsmodels/RegressionModelForecastPlotLevels.py b/validmind/tests/model_validation/statsmodels/RegressionModelForecastPlotLevels.py
index c634a300d..467ea1440 100644
--- a/validmind/tests/model_validation/statsmodels/RegressionModelForecastPlotLevels.py
+++ b/validmind/tests/model_validation/statsmodels/RegressionModelForecastPlotLevels.py
@@ -94,4 +94,6 @@ def RegressionModelForecastPlotLevels(
     return fig, RawData(
         y_transformed=dataset_y_transformed,
         y_pred_transformed=y_pred_transformed,
+        model=model.input_id,
+        dataset=dataset.input_id,
     )
diff --git a/validmind/tests/model_validation/statsmodels/RegressionModelSensitivityPlot.py b/validmind/tests/model_validation/statsmodels/RegressionModelSensitivityPlot.py
index d3635fd0c..d8f5c120d 100644
--- a/validmind/tests/model_validation/statsmodels/RegressionModelSensitivityPlot.py
+++ b/validmind/tests/model_validation/statsmodels/RegressionModelSensitivityPlot.py
@@ -114,4 +114,6 @@ def RegressionModelSensitivityPlot(
     return fig, RawData(
         transformed_target=transformed_target,
         transformed_predictions=transformed_predictions,
+        model=model.input_id,
+        dataset=dataset.input_id,
     )
diff --git a/validmind/tests/model_validation/statsmodels/RegressionModelSummary.py b/validmind/tests/model_validation/statsmodels/RegressionModelSummary.py
index 2814810fd..88d0b9162 100644
--- a/validmind/tests/model_validation/statsmodels/RegressionModelSummary.py
+++ b/validmind/tests/model_validation/statsmodels/RegressionModelSummary.py
@@ -4,7 +4,7 @@
 
 from sklearn.metrics import mean_squared_error, r2_score
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
 from .statsutils import adj_r2_score
@@ -58,4 +58,6 @@ def RegressionModelSummary(dataset: VMDataset, model: VMModel):
             "MSE": mean_squared_error(y_true, y_pred, squared=True),
             "RMSE": mean_squared_error(y_true, y_pred, squared=False),
         }
-    ]
+    ], RawData(
+        y_true=y_true, y_pred=y_pred, model=model.input_id, dataset=dataset.input_id
+    )
diff --git a/validmind/tests/model_validation/statsmodels/RegressionPermutationFeatureImportance.py b/validmind/tests/model_validation/statsmodels/RegressionPermutationFeatureImportance.py
index cb314ab0b..108cd45cf 100644
--- a/validmind/tests/model_validation/statsmodels/RegressionPermutationFeatureImportance.py
+++ b/validmind/tests/model_validation/statsmodels/RegressionPermutationFeatureImportance.py
@@ -97,4 +97,6 @@ def RegressionPermutationFeatureImportance(
         height=figure_height,
     )
 
-    return fig, RawData(importances=importances)
+    return fig, RawData(
+        importances=importances, model=model.input_id, dataset=dataset.input_id
+    )

From eb84b9c0127f46f6e1ae0b7d15bdd42cf5777300 Mon Sep 17 00:00:00 2001
From: Juan <juan@validmind.ai>
Date: Tue, 18 Feb 2025 15:30:47 +0100
Subject: [PATCH 05/10] Update RawData for ongoing monitoring tests

---
 .../CalibrationCurveDrift.py                  | 12 ++++++++++-
 .../ClassificationAccuracyDrift.py            | 12 +++++++++--
 .../ConfusionMatrixDrift.py                   |  9 ++++++++-
 .../CumulativePredictionProbabilitiesDrift.py | 20 +++++++++++++++++--
 .../tests/ongoing_monitoring/FeatureDrift.py  | 11 ++++++++--
 .../PredictionAcrossEachFeature.py            | 10 ++++++++--
 .../PredictionCorrelation.py                  | 15 ++++++++++++--
 .../PredictionProbabilitiesHistogramDrift.py  | 15 ++++++++++++--
 .../tests/ongoing_monitoring/ROCCurveDrift.py | 18 +++++++++++++++--
 .../ongoing_monitoring/ScoreBandsDrift.py     | 13 ++++++++++--
 .../TargetPredictionDistributionPlot.py       | 15 ++++++++++++--
 11 files changed, 130 insertions(+), 20 deletions(-)

diff --git a/validmind/tests/ongoing_monitoring/CalibrationCurveDrift.py b/validmind/tests/ongoing_monitoring/CalibrationCurveDrift.py
index 872d92238..f29848137 100644
--- a/validmind/tests/ongoing_monitoring/CalibrationCurveDrift.py
+++ b/validmind/tests/ongoing_monitoring/CalibrationCurveDrift.py
@@ -9,7 +9,7 @@
 import plotly.graph_objects as go
 from sklearn.calibration import calibration_curve
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.errors import SkipTestError
 from validmind.vm_models import VMDataset, VMModel
 
@@ -217,4 +217,14 @@ def CalibrationCurveDrift(
         fig,
         {"Mean Predicted Probabilities": pred_df, "Fraction of Positives": true_df},
         pass_fail_bool,
+        RawData(
+            prob_true_ref=prob_true_ref,
+            prob_pred_ref=prob_pred_ref,
+            prob_true_mon=prob_true_mon,
+            prob_pred_mon=prob_pred_mon,
+            bin_labels=bin_labels,
+            model=model.input_id,
+            dataset_ref=datasets[0].input_id,
+            dataset_mon=datasets[1].input_id,
+        ),
     )
diff --git a/validmind/tests/ongoing_monitoring/ClassificationAccuracyDrift.py b/validmind/tests/ongoing_monitoring/ClassificationAccuracyDrift.py
index 37bc97119..85b43a576 100644
--- a/validmind/tests/ongoing_monitoring/ClassificationAccuracyDrift.py
+++ b/validmind/tests/ongoing_monitoring/ClassificationAccuracyDrift.py
@@ -8,7 +8,7 @@
 import pandas as pd
 from sklearn.metrics import classification_report
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
 
@@ -145,4 +145,12 @@ def ClassificationAccuracyDrift(
     # Calculate overall pass/fail
     pass_fail_bool = (df["Pass/Fail"] == "Pass").all()
 
-    return ({"Classification Accuracy Metrics": df}, pass_fail_bool)
+    raw_data = RawData(
+        report_reference=report_ref,
+        report_monitoring=report_mon,
+        model=model.input_id,
+        dataset_reference=datasets[0].input_id,
+        dataset_monitoring=datasets[1].input_id,
+    )
+
+    return ({"Classification Accuracy Metrics": df}, pass_fail_bool, raw_data)
diff --git a/validmind/tests/ongoing_monitoring/ConfusionMatrixDrift.py b/validmind/tests/ongoing_monitoring/ConfusionMatrixDrift.py
index 7579f81dc..37a1a2ac0 100644
--- a/validmind/tests/ongoing_monitoring/ConfusionMatrixDrift.py
+++ b/validmind/tests/ongoing_monitoring/ConfusionMatrixDrift.py
@@ -8,7 +8,7 @@
 import pandas as pd
 from sklearn.metrics import confusion_matrix
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
 
@@ -190,4 +190,11 @@ def ConfusionMatrixDrift(
     return (
         {"Confusion Matrix Metrics": metrics_df, "Sample Counts": counts_df},
         pass_fail_bool,
+        RawData(
+            confusion_matrix_reference=cm_ref,
+            confusion_matrix_monitoring=cm_mon,
+            model=model.input_id,
+            dataset_reference=datasets[0].input_id,
+            dataset_monitoring=datasets[1].input_id,
+        ),
     )
diff --git a/validmind/tests/ongoing_monitoring/CumulativePredictionProbabilitiesDrift.py b/validmind/tests/ongoing_monitoring/CumulativePredictionProbabilitiesDrift.py
index 32b03db43..d4143c5d7 100644
--- a/validmind/tests/ongoing_monitoring/CumulativePredictionProbabilitiesDrift.py
+++ b/validmind/tests/ongoing_monitoring/CumulativePredictionProbabilitiesDrift.py
@@ -8,7 +8,7 @@
 import plotly.graph_objects as go
 from plotly.subplots import make_subplots
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
 
@@ -83,6 +83,7 @@ def CumulativePredictionProbabilitiesDrift(
     diff_color = "rgba(148, 103, 189, 0.8)"  # Purple with 0.8 opacity
 
     figures = []
+    raw_data = {}
     for class_value in classes:
         # Create figure with secondary y-axis
         fig = make_subplots(
@@ -175,4 +176,19 @@ def CumulativePredictionProbabilitiesDrift(
 
         figures.append(fig)
 
-    return tuple(figures)
+        # Store raw data for current class
+        raw_data[f"class_{class_value}_ref_probs"] = ref_probs
+        raw_data[f"class_{class_value}_mon_probs"] = mon_probs
+        raw_data[f"class_{class_value}_ref_sorted"] = ref_sorted
+        raw_data[f"class_{class_value}_ref_cumsum"] = ref_cumsum
+        raw_data[f"class_{class_value}_mon_sorted"] = mon_sorted
+        raw_data[f"class_{class_value}_mon_cumsum"] = mon_cumsum
+
+    return tuple(figures) + (
+        RawData(
+            model=model.input_id,
+            dataset_reference=datasets[0].input_id,
+            dataset_monitoring=datasets[1].input_id,
+            **raw_data,
+        ),
+    )
diff --git a/validmind/tests/ongoing_monitoring/FeatureDrift.py b/validmind/tests/ongoing_monitoring/FeatureDrift.py
index 9b3ee571a..52547e2f7 100644
--- a/validmind/tests/ongoing_monitoring/FeatureDrift.py
+++ b/validmind/tests/ongoing_monitoring/FeatureDrift.py
@@ -6,7 +6,7 @@
 import pandas as pd
 import plotly.graph_objects as go
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 
 
 def calculate_psi_score(actual, expected):
@@ -183,4 +183,11 @@ def FeatureDrift(
     # Calculate overall pass/fail
     pass_fail_bool = (psi_df["Pass/Fail"] == "Pass").all()
 
-    return ({"PSI Scores": psi_df}, *figures, pass_fail_bool)
+    # Prepare raw data
+    raw_data = RawData(
+        distributions=distributions,
+        dataset_reference=datasets[0].input_id,
+        dataset_monitoring=datasets[1].input_id,
+    )
+
+    return ({"PSI Scores": psi_df}, *figures, pass_fail_bool, raw_data)
diff --git a/validmind/tests/ongoing_monitoring/PredictionAcrossEachFeature.py b/validmind/tests/ongoing_monitoring/PredictionAcrossEachFeature.py
index 29c68b194..10c075c78 100644
--- a/validmind/tests/ongoing_monitoring/PredictionAcrossEachFeature.py
+++ b/validmind/tests/ongoing_monitoring/PredictionAcrossEachFeature.py
@@ -5,7 +5,7 @@
 
 import matplotlib.pyplot as plt
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 
 
 @tags("visualization")
@@ -74,4 +74,10 @@ def PredictionAcrossEachFeature(datasets, model):
         figures_to_save.append(fig)
         plt.close()
 
-    return tuple(figures_to_save)
+    return tuple(figures_to_save), RawData(
+        y_prob_reference=y_prob_reference,
+        y_prob_monitoring=y_prob_monitoring,
+        model=model.input_id,
+        dataset_reference=datasets[0].input_id,
+        dataset_monitoring=datasets[1].input_id,
+    )
diff --git a/validmind/tests/ongoing_monitoring/PredictionCorrelation.py b/validmind/tests/ongoing_monitoring/PredictionCorrelation.py
index 4c5f53c12..ddbc14503 100644
--- a/validmind/tests/ongoing_monitoring/PredictionCorrelation.py
+++ b/validmind/tests/ongoing_monitoring/PredictionCorrelation.py
@@ -5,7 +5,7 @@
 import pandas as pd
 import plotly.graph_objects as go
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 
 
 @tags("visualization")
@@ -140,4 +140,15 @@ def PredictionCorrelation(datasets, model, drift_pct_threshold=20):
     # Calculate overall pass/fail
     pass_fail_bool = (corr_final["Pass/Fail"] == "Pass").all()
 
-    return ({"Correlation Pair Table": corr_final}, fig, pass_fail_bool)
+    return (
+        {"Correlation Pair Table": corr_final},
+        fig,
+        pass_fail_bool,
+        RawData(
+            reference_correlations=corr_ref.to_dict(),
+            monitoring_correlations=corr_mon.to_dict(),
+            model=model.input_id,
+            dataset_reference=datasets[0].input_id,
+            dataset_monitoring=datasets[1].input_id,
+        ),
+    )
diff --git a/validmind/tests/ongoing_monitoring/PredictionProbabilitiesHistogramDrift.py b/validmind/tests/ongoing_monitoring/PredictionProbabilitiesHistogramDrift.py
index 60a435b99..150e544be 100644
--- a/validmind/tests/ongoing_monitoring/PredictionProbabilitiesHistogramDrift.py
+++ b/validmind/tests/ongoing_monitoring/PredictionProbabilitiesHistogramDrift.py
@@ -10,7 +10,7 @@
 from plotly.subplots import make_subplots
 from scipy import stats
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
 
@@ -201,4 +201,15 @@ def PredictionProbabilitiesHistogramDrift(
             }
         )
 
-    return fig, tables, all_passed
+    return (
+        fig,
+        tables,
+        all_passed,
+        RawData(
+            reference_probabilities=y_prob_ref,
+            monitoring_probabilities=y_prob_mon,
+            model=model.input_id,
+            dataset_reference=datasets[0].input_id,
+            dataset_monitoring=datasets[1].input_id,
+        ),
+    )
diff --git a/validmind/tests/ongoing_monitoring/ROCCurveDrift.py b/validmind/tests/ongoing_monitoring/ROCCurveDrift.py
index a01b82bfc..b2cc2e335 100644
--- a/validmind/tests/ongoing_monitoring/ROCCurveDrift.py
+++ b/validmind/tests/ongoing_monitoring/ROCCurveDrift.py
@@ -8,7 +8,7 @@
 import plotly.graph_objects as go
 from sklearn.metrics import roc_auc_score, roc_curve
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.errors import SkipTestError
 from validmind.vm_models import VMDataset, VMModel
 
@@ -147,4 +147,18 @@ def ROCCurveDrift(datasets: List[VMDataset], model: VMModel):
         height=500,
     )
 
-    return fig1, fig2
+    return (
+        fig1,
+        fig2,
+        RawData(
+            fpr_ref=fpr_ref,
+            tpr_ref=tpr_ref,
+            auc_ref=auc_ref,
+            fpr_mon=fpr_mon,
+            tpr_mon=tpr_mon,
+            auc_mon=auc_mon,
+            model=model.input_id,
+            dataset_reference=datasets[0].input_id,
+            dataset_monitoring=datasets[1].input_id,
+        ),
+    )
diff --git a/validmind/tests/ongoing_monitoring/ScoreBandsDrift.py b/validmind/tests/ongoing_monitoring/ScoreBandsDrift.py
index 2e9b8f76a..b122e26bd 100644
--- a/validmind/tests/ongoing_monitoring/ScoreBandsDrift.py
+++ b/validmind/tests/ongoing_monitoring/ScoreBandsDrift.py
@@ -7,7 +7,7 @@
 import numpy as np
 import pandas as pd
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
 
@@ -209,4 +209,13 @@ def process_dataset(dataset, model):
         tables[table_name] = pd.DataFrame(rows)
         all_passed &= metric_passed
 
-    return tables, all_passed
+    # Collect raw data
+    raw_data = RawData(
+        ref_results=ref_results,
+        mon_results=mon_results,
+        model=model.input_id,
+        dataset_reference=datasets[0].input_id,
+        dataset_monitoring=datasets[1].input_id,
+    )
+
+    return tables, all_passed, raw_data
diff --git a/validmind/tests/ongoing_monitoring/TargetPredictionDistributionPlot.py b/validmind/tests/ongoing_monitoring/TargetPredictionDistributionPlot.py
index eab0e8a7c..d3eba152b 100644
--- a/validmind/tests/ongoing_monitoring/TargetPredictionDistributionPlot.py
+++ b/validmind/tests/ongoing_monitoring/TargetPredictionDistributionPlot.py
@@ -7,7 +7,7 @@
 import plotly.graph_objects as go
 from scipy.stats import kurtosis, skew
 
-from validmind import tags, tasks
+from validmind import RawData, tags, tasks
 
 
 @tags("visualization")
@@ -142,4 +142,15 @@ def TargetPredictionDistributionPlot(datasets, model, drift_pct_threshold=20):
 
     pass_fail_bool = (moments["Pass/Fail"] == "Pass").all()
 
-    return ({"Distribution Moments": moments}, fig, pass_fail_bool)
+    return (
+        {"Distribution Moments": moments},
+        fig,
+        pass_fail_bool,
+        RawData(
+            pred_ref=pred_ref,
+            pred_monitor=pred_monitor,
+            model=model.input_id,
+            dataset_reference=datasets[0].input_id,
+            dataset_monitoring=datasets[1].input_id,
+        ),
+    )

From f4922ee89b30bbc5b94ba755ef708eac4cfe8325 Mon Sep 17 00:00:00 2001
From: Juan <juan@validmind.ai>
Date: Tue, 18 Feb 2025 15:36:45 +0100
Subject: [PATCH 06/10] Update RawData for prompt validation tests

---
 validmind/tests/prompt_validation/Clarity.py             | 2 +-
 validmind/tests/prompt_validation/NegativeInstruction.py | 2 +-
 validmind/tests/prompt_validation/Robustness.py          | 7 ++++++-
 validmind/tests/prompt_validation/Specificity.py         | 2 +-
 4 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/validmind/tests/prompt_validation/Clarity.py b/validmind/tests/prompt_validation/Clarity.py
index 14721e382..c6c5c6cc0 100644
--- a/validmind/tests/prompt_validation/Clarity.py
+++ b/validmind/tests/prompt_validation/Clarity.py
@@ -106,5 +106,5 @@ def Clarity(model, min_threshold=7):
             }
         ],
         passed,
-        RawData(response=response),
+        RawData(response=response, model=model.input_id),
     )
diff --git a/validmind/tests/prompt_validation/NegativeInstruction.py b/validmind/tests/prompt_validation/NegativeInstruction.py
index ad0039d54..b2d62e4c9 100644
--- a/validmind/tests/prompt_validation/NegativeInstruction.py
+++ b/validmind/tests/prompt_validation/NegativeInstruction.py
@@ -117,5 +117,5 @@ def NegativeInstruction(model, min_threshold=7):
             }
         ],
         passed,
-        RawData(model_response=response),
+        RawData(model_response=response, model=model.input_id),
     )
diff --git a/validmind/tests/prompt_validation/Robustness.py b/validmind/tests/prompt_validation/Robustness.py
index 7216b1ecd..627554e28 100644
--- a/validmind/tests/prompt_validation/Robustness.py
+++ b/validmind/tests/prompt_validation/Robustness.py
@@ -130,5 +130,10 @@ def Robustness(model, dataset, num_tests=10):
     return (
         results,
         all(result["Pass/Fail"] == "Pass" for result in results),
-        RawData(generated_inputs=generated_inputs, responses=responses),
+        RawData(
+            generated_inputs=generated_inputs,
+            responses=responses,
+            model=model.input_id,
+            dataset=dataset.input_id,
+        ),
     )
diff --git a/validmind/tests/prompt_validation/Specificity.py b/validmind/tests/prompt_validation/Specificity.py
index 7c1999a32..0a2dfb700 100644
--- a/validmind/tests/prompt_validation/Specificity.py
+++ b/validmind/tests/prompt_validation/Specificity.py
@@ -113,5 +113,5 @@ def Specificity(model, min_threshold=7):
             }
         ],
         passed,
-        RawData(response=response),
+        RawData(response=response, model=model.input_id),
     )

From c4db885aacbb85a8119448c9a15d1254286e6922 Mon Sep 17 00:00:00 2001
From: Juan <juan@validmind.ai>
Date: Tue, 18 Feb 2025 15:53:02 +0100
Subject: [PATCH 07/10] Add section for comparison tests in rawdata notebook

---
 .../how_to/understand_utilize_rawdata.ipynb   | 44 +++++++++++--------
 1 file changed, 25 insertions(+), 19 deletions(-)

diff --git a/notebooks/how_to/understand_utilize_rawdata.ipynb b/notebooks/how_to/understand_utilize_rawdata.ipynb
index b5c1dd6f6..93d664a7c 100644
--- a/notebooks/how_to/understand_utilize_rawdata.ipynb
+++ b/notebooks/how_to/understand_utilize_rawdata.ipynb
@@ -564,12 +564,14 @@
     "\n",
     "### Using `RawData` in comparison tests\n",
     "\n",
-    "TBC."
+    "When running comparison tests, the `RawData` object will contain the raw data for each individual test result as well as the comparison results between the test results. To support this, the RawData object contains the model and dataset input_ids for each of the datasets and models in the test, so that the post-processing function can use them to customize the output. The example below shows how to use the `RawData` object to customize the output of a comparison test and add a table to the test result that shows the confusion matrix for each individual test result as well as the comparison results between the test results.\n",
+    "\n",
+    "When designing post-processing functions that need to handle both individual and comparison test results, you can check the structure of the raw data to determine which case you're dealing with. In the example below, we check if `confusion_matrix` is a list (comparison test with multiple matrices) or a single matrix (individual test). For comparison tests, the function creates two tables: one showing the confusion matrices for each test case, and another showing the percentage drift between them. For individual tests, it creates a single table with the confusion matrix values. This pattern of checking the raw data structure can be applied to other tests to create versatile post-processing functions that work in both scenarios.\n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 13,
    "id": "bcbbe9f4",
    "metadata": {},
    "outputs": [],
@@ -626,6 +628,14 @@
     "    return result"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "41edd959",
+   "metadata": {},
+   "source": [
+    "Let's first run the confusion matrix test on a single dataset-model pair to see how our post-processing function handles individual results:"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -646,6 +656,14 @@
     ")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "a2482c54",
+   "metadata": {},
+   "source": [
+    "Now let's run a comparison test between test and train datasets to see how the function handles multiple results:"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -665,33 +683,21 @@
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "012ec495",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "result_cm.raw_data.confusion_matrix"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "4626bb96",
+   "cell_type": "markdown",
+   "id": "9f7d361a",
    "metadata": {},
-   "outputs": [],
    "source": [
-    "result_cm.raw_data.model"
+    "Let's inspect the raw data to see how comparison tests structure their data - notice how the `RawData` object contains not just the confusion matrices for both datasets, but also tracks which dataset and model each result came from:"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "07825ec9",
+   "id": "012ec495",
    "metadata": {},
    "outputs": [],
    "source": [
-    "result_cm.raw_data.dataset"
+    "result_cm.raw_data.inspect()"
    ]
   }
  ],

From c84c6790d5be7ecd97ee03b162011974d83e7bcb Mon Sep 17 00:00:00 2001
From: Juan <juan@validmind.ai>
Date: Tue, 18 Feb 2025 16:44:17 +0100
Subject: [PATCH 08/10] Fix unit tests

---
 .../data_validation/test_IQROutliersTable.py  |  7 +--
 .../test_IsolationForestOutliers.py           | 30 +++++-------
 .../data_validation/test_JarqueBera.py        |  5 +-
 .../data_validation/test_LJungBox.py          |  5 +-
 .../data_validation/test_MissingValues.py     |  8 +--
 .../data_validation/test_ShapiroWilk.py       |  7 ++-
 .../test_TimeSeriesOutliers.py                |  3 +-
 .../data_validation/test_ZivotAndrewsArch.py  |  7 ++-
 .../sklearn/test_FeatureImportance.py         | 25 ++++++----
 .../sklearn/test_RegressionErrors.py          | 11 +++--
 .../sklearn/test_RegressionR2Square.py        | 35 +++++++------
 .../statsmodels/test_DurbinWatsonTest.py      | 10 +++-
 .../statsmodels/test_GINITable.py             | 49 ++++++++++---------
 .../test_PredictionProbabilitiesHistogram.py  | 17 +++++--
 .../statsmodels/test_RegressionCoeffs.py      |  7 ++-
 .../test_TimeSeriesPredictionsPlot.py         | 15 ++++--
 .../test_TimeSeriesR2SquareBySegments.py      | 11 +++--
 17 files changed, 152 insertions(+), 100 deletions(-)

diff --git a/tests/unit_tests/data_validation/test_IQROutliersTable.py b/tests/unit_tests/data_validation/test_IQROutliersTable.py
index eed153eed..0174d724f 100644
--- a/tests/unit_tests/data_validation/test_IQROutliersTable.py
+++ b/tests/unit_tests/data_validation/test_IQROutliersTable.py
@@ -38,10 +38,11 @@ def setUp(self):
         )
 
     def test_outliers_structure(self):
-        result = IQROutliersTable(self.vm_dataset)
+        result, raw_data = IQROutliersTable(self.vm_dataset)
 
         # Check basic structure
         self.assertIsInstance(result, dict)
+        self.assertIsInstance(raw_data, vm.RawData)
         self.assertIn("Summary of Outliers Detected by IQR Method", result)
 
         # Check result structure
@@ -59,7 +60,7 @@ def test_outliers_structure(self):
             self.assertIn("Maximum Outlier Value", summary)
 
     def test_outliers_detection(self):
-        result = IQROutliersTable(self.vm_dataset)
+        result, raw_data = IQROutliersTable(self.vm_dataset)
         outliers_summary = result["Summary of Outliers Detected by IQR Method"]
 
         # Check that outliers are detected in the 'with_outliers' column
@@ -76,7 +77,7 @@ def test_outliers_detection(self):
         self.assertIsNone(normal_summary)
 
     def test_binary_exclusion(self):
-        result = IQROutliersTable(self.vm_dataset)
+        result, raw_data = IQROutliersTable(self.vm_dataset)
         outliers_summary = result["Summary of Outliers Detected by IQR Method"]
 
         # Verify binary column is not in results
diff --git a/tests/unit_tests/data_validation/test_IsolationForestOutliers.py b/tests/unit_tests/data_validation/test_IsolationForestOutliers.py
index 4a49371b9..19a7d494f 100644
--- a/tests/unit_tests/data_validation/test_IsolationForestOutliers.py
+++ b/tests/unit_tests/data_validation/test_IsolationForestOutliers.py
@@ -28,25 +28,19 @@ def setUp(self):
         )
 
     def test_outliers_detection(self):
-        result = IsolationForestOutliers(self.vm_dataset, contamination=0.1)
+        figure, raw_data = IsolationForestOutliers(self.vm_dataset, contamination=0.1)
 
-        # Check return type
-        self.assertIsInstance(result, tuple)
+        # Check return types
+        self.assertIsInstance(figure, plt.Figure)
+        self.assertIsInstance(raw_data, vm.RawData)
 
-        # Separate figures and raw data
-        figures = result
-
-        # Check that at least one figure is returned
-        self.assertGreater(len(figures), 0)
-
-        # Check each figure
-        for fig in figures:
-            self.assertIsInstance(fig, plt.Figure)
+        # Check that the figure has at least one axes
+        self.assertGreater(len(figure.axes), 0)
 
     def test_feature_columns_validation(self):
         # Test with valid feature columns
         try:
-            IsolationForestOutliers(
+            figure, raw_data = IsolationForestOutliers(
                 self.vm_dataset, feature_columns=["feature1", "feature2"]
             )
         except ValueError:
@@ -60,13 +54,13 @@ def test_feature_columns_validation(self):
 
     def test_contamination_parameter(self):
         # Test with different contamination levels
-        figures_low_contamination = IsolationForestOutliers(
+        figure_low, raw_data_low = IsolationForestOutliers(
             self.vm_dataset, contamination=0.05
         )
-        figures_high_contamination = IsolationForestOutliers(
+        figure_high, raw_data_high = IsolationForestOutliers(
             self.vm_dataset, contamination=0.2
         )
 
-        # Check that figures are returned for both contamination levels
-        self.assertGreater(len(figures_low_contamination), 0)
-        self.assertGreater(len(figures_high_contamination), 0)
+        # Check that figures have at least one axes
+        self.assertGreater(len(figure_low.axes), 0)
+        self.assertGreater(len(figure_high.axes), 0)
diff --git a/tests/unit_tests/data_validation/test_JarqueBera.py b/tests/unit_tests/data_validation/test_JarqueBera.py
index bd7bed930..5c17bc0ad 100644
--- a/tests/unit_tests/data_validation/test_JarqueBera.py
+++ b/tests/unit_tests/data_validation/test_JarqueBera.py
@@ -29,11 +29,14 @@ def test_returns_dataframe_and_rawdata(self):
         )
 
         # Run the function
-        result = JarqueBera(vm_dataset)
+        result, raw_data = JarqueBera(vm_dataset)
 
         # Check if result is a DataFrame
         self.assertIsInstance(result, pd.DataFrame)
 
+        # Check if raw_data is a RawData object
+        self.assertIsInstance(raw_data, vm.RawData)
+
         # Check if the DataFrame has the expected columns
         expected_columns = ["column", "stat", "pvalue", "skew", "kurtosis"]
         self.assertListEqual(list(result.columns), expected_columns)
diff --git a/tests/unit_tests/data_validation/test_LJungBox.py b/tests/unit_tests/data_validation/test_LJungBox.py
index 3880dcf3d..d50e501b9 100644
--- a/tests/unit_tests/data_validation/test_LJungBox.py
+++ b/tests/unit_tests/data_validation/test_LJungBox.py
@@ -22,11 +22,14 @@ def test_returns_dataframe_with_expected_shape(self):
         )
 
         # Run the function
-        result = LJungBox(vm_dataset)
+        result, raw_data = LJungBox(vm_dataset)
 
         # Check if result is a DataFrame
         self.assertIsInstance(result, pd.DataFrame)
 
+        # Check if raw_data is a RawData object
+        self.assertIsInstance(raw_data, vm.RawData)
+
         # Check if the DataFrame has the expected columns
         expected_columns = ["column", "stat", "pvalue"]
         self.assertListEqual(list(result.columns), expected_columns)
diff --git a/tests/unit_tests/data_validation/test_MissingValues.py b/tests/unit_tests/data_validation/test_MissingValues.py
index bbccff0b3..f8d7aa2b0 100644
--- a/tests/unit_tests/data_validation/test_MissingValues.py
+++ b/tests/unit_tests/data_validation/test_MissingValues.py
@@ -28,11 +28,13 @@ def setUp(self):
         )
 
     def test_missing_values_structure(self):
-        summary, passed = MissingValues(self.vm_dataset)
+        # Run the function
+        summary, passed, raw_data = MissingValues(self.vm_dataset)
 
         # Check return types
         self.assertIsInstance(summary, list)
         self.assertIsInstance(passed, bool)
+        self.assertIsInstance(raw_data, vm.RawData)
 
         # Check summary structure
         for column_summary in summary:
@@ -42,7 +44,7 @@ def test_missing_values_structure(self):
             self.assertIn("Pass/Fail", column_summary)
 
     def test_missing_values_counts(self):
-        summary, passed = MissingValues(self.vm_dataset)
+        summary, passed, raw_data = MissingValues(self.vm_dataset)
 
         # Get results for each column
         no_missing = next(s for s in summary if s["Column"] == "no_missing")
@@ -69,7 +71,7 @@ def test_missing_values_counts(self):
 
     def test_threshold_parameter(self):
         # Test with higher threshold that allows some missing values
-        summary, passed = MissingValues(self.vm_dataset, min_threshold=25)
+        summary, passed, raw_data = MissingValues(self.vm_dataset, min_threshold=25)
 
         # Get results
         some_missing = next(s for s in summary if s["Column"] == "some_missing")
diff --git a/tests/unit_tests/data_validation/test_ShapiroWilk.py b/tests/unit_tests/data_validation/test_ShapiroWilk.py
index 59989998e..dca9cd8b5 100644
--- a/tests/unit_tests/data_validation/test_ShapiroWilk.py
+++ b/tests/unit_tests/data_validation/test_ShapiroWilk.py
@@ -30,11 +30,14 @@ def setUp(self):
 
     def test_returns_dataframe_and_rawdata(self):
         # Run the function
-        result_df = ShapiroWilk(self.vm_dataset)
+        result_df, raw_data = ShapiroWilk(self.vm_dataset)
 
         # Check if result_df is a DataFrame
         self.assertIsInstance(result_df, pd.DataFrame)
 
+        # Check if raw_data is a RawData object
+        self.assertIsInstance(raw_data, vm.RawData)
+
         # Check if the DataFrame has the expected columns
         expected_columns = ["column", "stat", "pvalue"]
         self.assertListEqual(list(result_df.columns), expected_columns)
@@ -44,7 +47,7 @@ def test_returns_dataframe_and_rawdata(self):
 
     def test_handles_different_distributions(self):
         # Run the function
-        result_df = ShapiroWilk(self.vm_dataset)
+        result_df, raw_data = ShapiroWilk(self.vm_dataset)
 
         # The normal distribution should have a higher p-value than the exponential distribution
         normal_pvalue = result_df[result_df["column"] == "normal_dist"]["pvalue"].iloc[
diff --git a/tests/unit_tests/data_validation/test_TimeSeriesOutliers.py b/tests/unit_tests/data_validation/test_TimeSeriesOutliers.py
index 193653a91..62beb3b93 100644
--- a/tests/unit_tests/data_validation/test_TimeSeriesOutliers.py
+++ b/tests/unit_tests/data_validation/test_TimeSeriesOutliers.py
@@ -36,12 +36,13 @@ def setUp(self):
         )
 
     def test_time_series_outliers(self):
-        outlier_df, figures, passed = TimeSeriesOutliers(self.vm_dataset)
+        outlier_df, figures, passed, raw_data = TimeSeriesOutliers(self.vm_dataset)
 
         # Check return types
         self.assertIsInstance(outlier_df, pd.DataFrame)
         self.assertIsInstance(figures, list)
         self.assertIsInstance(passed, bool)
+        self.assertIsInstance(raw_data, vm.RawData)
 
         # Check that we have the expected number of figures (one per feature)
         self.assertEqual(len(figures), 2)
diff --git a/tests/unit_tests/data_validation/test_ZivotAndrewsArch.py b/tests/unit_tests/data_validation/test_ZivotAndrewsArch.py
index 97e7716c7..9e9b038cc 100644
--- a/tests/unit_tests/data_validation/test_ZivotAndrewsArch.py
+++ b/tests/unit_tests/data_validation/test_ZivotAndrewsArch.py
@@ -39,10 +39,13 @@ def setUp(self):
         )
 
     def test_zivot_andrews(self):
-        result = ZivotAndrewsArch(self.vm_dataset)
+        result, raw_data = ZivotAndrewsArch(self.vm_dataset)
 
-        # Check return type and structure
+        # Check return types
         self.assertIsInstance(result, dict)
+        self.assertIsInstance(raw_data, vm.RawData)
+
+        # Check results structure
         self.assertIn("Zivot-Andrews Test Results", result)
 
         # Check results structure
diff --git a/tests/unit_tests/model_validation/sklearn/test_FeatureImportance.py b/tests/unit_tests/model_validation/sklearn/test_FeatureImportance.py
index 7cd346336..c7c75491b 100644
--- a/tests/unit_tests/model_validation/sklearn/test_FeatureImportance.py
+++ b/tests/unit_tests/model_validation/sklearn/test_FeatureImportance.py
@@ -58,23 +58,28 @@ def setUp(self):
 
     def test_returns_dataframe_and_rawdata(self):
         # Run the function
-        result_df = FeatureImportance(self.vm_dataset, self.vm_model)
+        results_df, raw_data = FeatureImportance(self.vm_dataset, self.vm_model)
 
-        # Check if result_df is a DataFrame
-        self.assertIsInstance(result_df, pd.DataFrame)
+        # Check if results_df is a DataFrame
+        self.assertIsInstance(results_df, pd.DataFrame)
+
+        # Check if raw_data is a RawData object
+        self.assertIsInstance(raw_data, vm.RawData)
 
         # Check if DataFrame has expected columns
         expected_columns = ["Feature 1", "Feature 2", "Feature 3"]
-        self.assertTrue(all(col in result_df.columns for col in expected_columns))
+        self.assertTrue(all(col in results_df.columns for col in expected_columns))
 
     def test_feature_importance_ranking(self):
         # Run with all features
-        result_df = FeatureImportance(self.vm_dataset, self.vm_model, num_features=4)
+        results_df, raw_data = FeatureImportance(
+            self.vm_dataset, self.vm_model, num_features=4
+        )
 
         # Get feature names and scores
         features = []
         for i in range(1, 5):
-            feature_info = result_df[f"Feature {i}"].iloc[0]
+            feature_info = results_df[f"Feature {i}"].iloc[0]
             feature_name = feature_info.split(";")[0].strip("[]")
             features.append(feature_name)
 
@@ -87,21 +92,21 @@ def test_feature_importance_ranking(self):
     def test_num_features_parameter(self):
         # Test with different num_features values
         for num_features in [2, 3, 4]:
-            result_df = FeatureImportance(
+            results_df, raw_data = FeatureImportance(
                 self.vm_dataset, self.vm_model, num_features=num_features
             )
 
             # Check number of columns matches num_features
             feature_columns = [
-                col for col in result_df.columns if col.startswith("Feature")
+                col for col in results_df.columns if col.startswith("Feature")
             ]
             self.assertEqual(len(feature_columns), num_features)
 
     def test_feature_importance_scores(self):
-        result_df = FeatureImportance(self.vm_dataset, self.vm_model)
+        results_df, raw_data = FeatureImportance(self.vm_dataset, self.vm_model)
 
         # Get first feature score
-        first_feature = result_df["Feature 1"].iloc[0]
+        first_feature = results_df["Feature 1"].iloc[0]
         score = float(first_feature.split(";")[1].strip("[] "))
 
         # Check if score is positive (since we're using absolute importance)
diff --git a/tests/unit_tests/model_validation/sklearn/test_RegressionErrors.py b/tests/unit_tests/model_validation/sklearn/test_RegressionErrors.py
index aced18ed8..71042035d 100644
--- a/tests/unit_tests/model_validation/sklearn/test_RegressionErrors.py
+++ b/tests/unit_tests/model_validation/sklearn/test_RegressionErrors.py
@@ -48,11 +48,14 @@ def setUp(self):
 
     def test_returns_dataframe_and_raw_data(self):
         # Run the function
-        results = RegressionErrors(self.vm_model, self.vm_dataset)
+        results, raw_data = RegressionErrors(self.vm_model, self.vm_dataset)
 
         # Check if results is a DataFrame
         self.assertIsInstance(results, pd.DataFrame)
 
+        # Check if raw_data is a RawData object
+        self.assertIsInstance(raw_data, vm.RawData)
+
         # Check if DataFrame has expected columns
         expected_columns = [
             "Mean Absolute Error (MAE)",
@@ -67,7 +70,7 @@ def test_returns_dataframe_and_raw_data(self):
         self.assertEqual(len(results), 1)
 
     def test_error_metrics_range(self):
-        results = RegressionErrors(self.vm_model, self.vm_dataset)
+        results, raw_data = RegressionErrors(self.vm_model, self.vm_dataset)
 
         # All error metrics should be non-negative (except MBD)
         self.assertGreaterEqual(results["Mean Absolute Error (MAE)"].iloc[0], 0)
@@ -109,7 +112,7 @@ def test_perfect_prediction(self):
         vm_perfect_dataset.assign_predictions(vm_perfect_model)
 
         # Calculate errors
-        results = RegressionErrors(vm_perfect_model, vm_perfect_dataset)
+        results, raw_data = RegressionErrors(vm_perfect_model, vm_perfect_dataset)
 
         # All error metrics should be very close to 0
         self.assertAlmostEqual(
@@ -124,7 +127,7 @@ def test_perfect_prediction(self):
         )
 
     def test_error_metrics_consistency(self):
-        results = RegressionErrors(self.vm_model, self.vm_dataset)
+        results, raw_data = RegressionErrors(self.vm_model, self.vm_dataset)
 
         # MSE should be greater than or equal to MAE squared
         mae = results["Mean Absolute Error (MAE)"].iloc[0]
diff --git a/tests/unit_tests/model_validation/sklearn/test_RegressionR2Square.py b/tests/unit_tests/model_validation/sklearn/test_RegressionR2Square.py
index 71447956a..eba7643bb 100644
--- a/tests/unit_tests/model_validation/sklearn/test_RegressionR2Square.py
+++ b/tests/unit_tests/model_validation/sklearn/test_RegressionR2Square.py
@@ -50,17 +50,20 @@ def setUp(self):
 
     def test_returns_dataframe(self):
         # Run the function
-        result = RegressionR2Square(self.vm_dataset, self.vm_model)
+        results, raw_data = RegressionR2Square(self.vm_dataset, self.vm_model)
 
-        # Check if result is a DataFrame
-        self.assertIsInstance(result, pd.DataFrame)
+        # Check if results is a DataFrame
+        self.assertIsInstance(results, pd.DataFrame)
+
+        # Check if raw_data is a RawData object
+        self.assertIsInstance(raw_data, vm.RawData)
 
         # Check if DataFrame has expected columns
         expected_columns = ["R-squared (R2) Score", "Adjusted R-squared (R2) Score"]
-        self.assertTrue(all(col in result.columns for col in expected_columns))
+        self.assertTrue(all(col in results.columns for col in expected_columns))
 
         # Check if DataFrame has exactly one row
-        self.assertEqual(len(result), 1)
+        self.assertEqual(len(results), 1)
 
     def test_perfect_prediction(self):
         # Create a perfect prediction scenario
@@ -93,21 +96,21 @@ def test_perfect_prediction(self):
         # Assign predictions to the dataset
         vm_perfect_dataset.assign_predictions(vm_perfect_model)
 
-        # Calculate R2 scores
-        result = RegressionR2Square(vm_perfect_dataset, vm_perfect_model)
+        # Calculate metrics
+        results, raw_data = RegressionR2Square(vm_perfect_dataset, vm_perfect_model)
 
-        # Both R2 and Adjusted R2 should be very close to 1 for perfect predictions
-        self.assertAlmostEqual(result["R-squared (R2) Score"].iloc[0], 1.0, places=5)
+        # For perfect separation:
+        self.assertAlmostEqual(results["R-squared (R2) Score"].iloc[0], 1.0, places=5)
         self.assertAlmostEqual(
-            result["Adjusted R-squared (R2) Score"].iloc[0], 1.0, places=5
+            results["Adjusted R-squared (R2) Score"].iloc[0], 1.0, places=5
         )
 
     def test_r2_score_range(self):
-        result = RegressionR2Square(self.vm_dataset, self.vm_model)
+        results, raw_data = RegressionR2Square(self.vm_dataset, self.vm_model)
 
         # R2 score should be between 0 and 1 for a reasonable model
-        r2_score = result["R-squared (R2) Score"].iloc[0]
-        adj_r2_score = result["Adjusted R-squared (R2) Score"].iloc[0]
+        r2_score = results["R-squared (R2) Score"].iloc[0]
+        adj_r2_score = results["Adjusted R-squared (R2) Score"].iloc[0]
 
         self.assertGreaterEqual(r2_score, 0)
         self.assertLessEqual(r2_score, 1)
@@ -146,8 +149,8 @@ def test_poor_prediction(self):
         vm_poor_dataset.assign_predictions(vm_poor_model)
 
         # Calculate R2 scores
-        result = RegressionR2Square(vm_poor_dataset, vm_poor_model)
+        results, raw_data = RegressionR2Square(vm_poor_dataset, vm_poor_model)
 
         # R2 score should be close to 0 for poor predictions
-        self.assertLess(result["R-squared (R2) Score"].iloc[0], 0.1)
-        self.assertLess(result["Adjusted R-squared (R2) Score"].iloc[0], 0.1)
+        self.assertLess(results["R-squared (R2) Score"].iloc[0], 0.1)
+        self.assertLess(results["Adjusted R-squared (R2) Score"].iloc[0], 0.1)
diff --git a/tests/unit_tests/model_validation/statsmodels/test_DurbinWatsonTest.py b/tests/unit_tests/model_validation/statsmodels/test_DurbinWatsonTest.py
index 6388ece40..44a460a1e 100644
--- a/tests/unit_tests/model_validation/statsmodels/test_DurbinWatsonTest.py
+++ b/tests/unit_tests/model_validation/statsmodels/test_DurbinWatsonTest.py
@@ -63,11 +63,14 @@ def setUp(self):
 
     def test_returns_dataframe_and_rawdata(self):
         # Run the function
-        results = DurbinWatsonTest(self.vm_dataset, self.vm_model)
+        results, raw_data = DurbinWatsonTest(self.vm_dataset, self.vm_model)
 
         # Check if results is a DataFrame
         self.assertIsInstance(results, pd.DataFrame)
 
+        # Check if raw_data is a RawData object
+        self.assertIsInstance(raw_data, vm.RawData)
+
         # Check if DataFrame has expected columns
         expected_columns = ["dw_statistic", "threshold", "autocorrelation"]
         self.assertTrue(all(col in results.columns for col in expected_columns))
@@ -107,11 +110,14 @@ def test_no_autocorrelation(self):
         vm_no_auto_dataset.assign_predictions(vm_no_auto_model)
 
         # Run the function
-        results = DurbinWatsonTest(vm_no_auto_dataset, vm_no_auto_model)
+        results, raw_data = DurbinWatsonTest(vm_no_auto_dataset, vm_no_auto_model)
 
         # Check if results is a DataFrame
         self.assertIsInstance(results, pd.DataFrame)
 
+        # Check if raw_data is a RawData object
+        self.assertIsInstance(raw_data, vm.RawData)
+
         # Check if DataFrame has expected columns
         expected_columns = ["dw_statistic", "threshold", "autocorrelation"]
         self.assertTrue(all(col in results.columns for col in expected_columns))
diff --git a/tests/unit_tests/model_validation/statsmodels/test_GINITable.py b/tests/unit_tests/model_validation/statsmodels/test_GINITable.py
index bb44381d0..c91814899 100644
--- a/tests/unit_tests/model_validation/statsmodels/test_GINITable.py
+++ b/tests/unit_tests/model_validation/statsmodels/test_GINITable.py
@@ -56,17 +56,20 @@ def setUp(self):
 
     def test_returns_dataframe_and_rawdata(self):
         # Run the function
-        result = GINITable(self.vm_dataset, self.vm_model)
+        results, raw_data = GINITable(self.vm_dataset, self.vm_model)
 
-        # Check if result is a DataFrame
-        self.assertIsInstance(result, pd.DataFrame)
+        # Check if results is a DataFrame
+        self.assertIsInstance(results, pd.DataFrame)
+
+        # Check if raw_data is a RawData object
+        self.assertIsInstance(raw_data, RawData)
 
         # Check if DataFrame has expected columns
         expected_columns = ["AUC", "GINI", "KS"]
-        self.assertTrue(all(col in result.columns for col in expected_columns))
+        self.assertTrue(all(col in results.columns for col in expected_columns))
 
         # Check if DataFrame has exactly one row
-        self.assertEqual(len(result), 1)
+        self.assertEqual(len(results), 1)
 
     def test_perfect_separation(self):
         # Create a dataset with perfect class separation
@@ -100,15 +103,15 @@ def test_perfect_separation(self):
         vm_perfect_dataset.assign_predictions(vm_perfect_model)
 
         # Calculate metrics
-        result = GINITable(vm_perfect_dataset, vm_perfect_model)
+        results, raw_data = GINITable(vm_perfect_dataset, vm_perfect_model)
 
         # For perfect separation:
         # - AUC should be 1.0
         # - GINI should be 1.0 (2*AUC - 1)
         # - KS should be 1.0
-        self.assertAlmostEqual(result["AUC"].iloc[0], 1.0, places=5)
-        self.assertAlmostEqual(result["GINI"].iloc[0], 1.0, places=5)
-        self.assertAlmostEqual(result["KS"].iloc[0], 1.0, places=5)
+        self.assertAlmostEqual(results["AUC"].iloc[0], 1.0, places=5)
+        self.assertAlmostEqual(results["GINI"].iloc[0], 1.0, places=5)
+        self.assertAlmostEqual(results["KS"].iloc[0], 1.0, places=5)
 
     def test_random_prediction(self):
         # Create a dataset with random predictions
@@ -144,36 +147,36 @@ def test_random_prediction(self):
         vm_random_dataset.assign_predictions(vm_random_model)
 
         # Calculate metrics
-        result = GINITable(vm_random_dataset, vm_random_model)
+        results, raw_data = GINITable(vm_random_dataset, vm_random_model)
 
         # For random predictions:
         # - AUC should be close to 0.5
         # - GINI should be close to 0.0 (2*0.5 - 1)
         # - KS should be close to 0.0
-        self.assertGreater(result["AUC"].iloc[0], 0.4)
-        self.assertLess(result["AUC"].iloc[0], 0.6)
-        self.assertGreater(result["GINI"].iloc[0], -0.2)
-        self.assertLess(result["GINI"].iloc[0], 0.2)
-        self.assertLess(result["KS"].iloc[0], 0.2)
+        self.assertGreater(results["AUC"].iloc[0], 0.4)
+        self.assertLess(results["AUC"].iloc[0], 0.6)
+        self.assertGreater(results["GINI"].iloc[0], -0.2)
+        self.assertLess(results["GINI"].iloc[0], 0.2)
+        self.assertLess(results["KS"].iloc[0], 0.2)
 
     def test_metric_ranges(self):
         # Test regular case
-        result = GINITable(self.vm_dataset, self.vm_model)
+        results, raw_data = GINITable(self.vm_dataset, self.vm_model)
 
         # Check metric ranges
         # AUC should be between 0 and 1
-        self.assertGreaterEqual(result["AUC"].iloc[0], 0.0)
-        self.assertLessEqual(result["AUC"].iloc[0], 1.0)
+        self.assertGreaterEqual(results["AUC"].iloc[0], 0.0)
+        self.assertLessEqual(results["AUC"].iloc[0], 1.0)
 
         # GINI should be between -1 and 1
-        self.assertGreaterEqual(result["GINI"].iloc[0], -1.0)
-        self.assertLessEqual(result["GINI"].iloc[0], 1.0)
+        self.assertGreaterEqual(results["GINI"].iloc[0], -1.0)
+        self.assertLessEqual(results["GINI"].iloc[0], 1.0)
 
         # KS should be between 0 and 1
-        self.assertGreaterEqual(result["KS"].iloc[0], 0.0)
-        self.assertLessEqual(result["KS"].iloc[0], 1.0)
+        self.assertGreaterEqual(results["KS"].iloc[0], 0.0)
+        self.assertLessEqual(results["KS"].iloc[0], 1.0)
 
         # GINI should be 2*AUC - 1
         self.assertAlmostEqual(
-            result["GINI"].iloc[0], 2 * result["AUC"].iloc[0] - 1, places=5
+            results["GINI"].iloc[0], 2 * results["AUC"].iloc[0] - 1, places=5
         )
diff --git a/tests/unit_tests/model_validation/statsmodels/test_PredictionProbabilitiesHistogram.py b/tests/unit_tests/model_validation/statsmodels/test_PredictionProbabilitiesHistogram.py
index cd55a22a4..12e098642 100644
--- a/tests/unit_tests/model_validation/statsmodels/test_PredictionProbabilitiesHistogram.py
+++ b/tests/unit_tests/model_validation/statsmodels/test_PredictionProbabilitiesHistogram.py
@@ -58,11 +58,16 @@ def setUp(self):
 
     def test_returns_figure_and_raw_data(self):
         # Run the function
-        figure = PredictionProbabilitiesHistogram(self.vm_dataset, self.vm_model)
+        figure, raw_data = PredictionProbabilitiesHistogram(
+            self.vm_dataset, self.vm_model
+        )
 
-        # Check if result contains a Plotly Figure and RawData
+        # Check if figure is a Plotly Figure
         self.assertIsInstance(figure, go.Figure)
 
+        # Check if raw_data is a RawData object
+        self.assertIsInstance(raw_data, vm.RawData)
+
         # Check if figure has traces
         self.assertGreater(len(figure.data), 0)
 
@@ -98,13 +103,17 @@ def test_perfect_separation(self):
         vm_perfect_dataset.assign_predictions(vm_perfect_model)
 
         # Generate histogram
-        figure = PredictionProbabilitiesHistogram(vm_perfect_dataset, vm_perfect_model)
+        figure, raw_data = PredictionProbabilitiesHistogram(
+            vm_perfect_dataset, vm_perfect_model
+        )
 
         # Check if there are exactly two traces (one for each class)
         self.assertEqual(len(figure.data), 2)
 
     def test_probability_ranges(self):
-        figure = PredictionProbabilitiesHistogram(self.vm_dataset, self.vm_model)
+        figure, raw_data = PredictionProbabilitiesHistogram(
+            self.vm_dataset, self.vm_model
+        )
 
         # Check if probabilities are within [0, 1] range
         for trace in figure.data:
diff --git a/tests/unit_tests/model_validation/statsmodels/test_RegressionCoeffs.py b/tests/unit_tests/model_validation/statsmodels/test_RegressionCoeffs.py
index 88f0448f7..da9581e78 100644
--- a/tests/unit_tests/model_validation/statsmodels/test_RegressionCoeffs.py
+++ b/tests/unit_tests/model_validation/statsmodels/test_RegressionCoeffs.py
@@ -47,7 +47,7 @@ def test_returns_tuple(self):
 
         # Check if result is a tuple
         self.assertIsInstance(result, tuple)
-        self.assertEqual(len(result), 2)
+        self.assertEqual(len(result), 3)  # Now expecting 3 values
 
         # Check if first element is a Plotly Figure
         self.assertIsInstance(result[0], go.Figure)
@@ -55,9 +55,12 @@ def test_returns_tuple(self):
         # Check if second element is a table (DataFrame)
         self.assertIsInstance(result[1], pd.DataFrame)
 
+        # Check if third element is a RawData object
+        self.assertIsInstance(result[2], vm.RawData)
+
     def test_plot_properties(self):
         # Run the function
-        fig, _ = RegressionCoeffs(self.vm_model)
+        fig, df, raw_data = RegressionCoeffs(self.vm_model)  # Unpack all three values
 
         # Check if there is exactly one trace (bar plot)
         self.assertEqual(len(fig.data), 1)
diff --git a/tests/unit_tests/model_validation/test_TimeSeriesPredictionsPlot.py b/tests/unit_tests/model_validation/test_TimeSeriesPredictionsPlot.py
index 5437f365d..024a9df36 100644
--- a/tests/unit_tests/model_validation/test_TimeSeriesPredictionsPlot.py
+++ b/tests/unit_tests/model_validation/test_TimeSeriesPredictionsPlot.py
@@ -50,19 +50,24 @@ def setUp(self):
 
     def test_return_types(self):
         """Test if function returns a tuple with a Plotly figure and raw data."""
-        result = TimeSeriesPredictionsPlot(self.vm_dataset, self.vm_model)
-        self.assertIsInstance(result, go.Figure)
+        fig, raw_data = TimeSeriesPredictionsPlot(self.vm_dataset, self.vm_model)
+
+        # Check if first element is a Plotly figure
+        self.assertIsInstance(fig, go.Figure)
+
+        # Check if second element is a RawData object
+        self.assertIsInstance(raw_data, vm.RawData)
 
     def test_figure_properties(self):
         """Test if figure has expected properties."""
-        fig = TimeSeriesPredictionsPlot(self.vm_dataset, self.vm_model)
+        fig, raw_data = TimeSeriesPredictionsPlot(self.vm_dataset, self.vm_model)
 
         # Check if figure has exactly two traces (Actual and Predicted)
         self.assertEqual(len(fig.data), 2)
 
     def test_data_length(self):
         """Test if the plotted data has correct length."""
-        fig = TimeSeriesPredictionsPlot(self.vm_dataset, self.vm_model)
+        fig, raw_data = TimeSeriesPredictionsPlot(self.vm_dataset, self.vm_model)
 
         # Both traces should have same length as input data
         self.assertEqual(len(fig.data[0].x), len(self.df))
@@ -72,7 +77,7 @@ def test_data_length(self):
 
     def test_datetime_index(self):
         """Test if x-axis uses datetime values."""
-        fig = TimeSeriesPredictionsPlot(self.vm_dataset, self.vm_model)
+        fig, raw_data = TimeSeriesPredictionsPlot(self.vm_dataset, self.vm_model)
 
         # Check if x values are datetime objects
         self.assertTrue(
diff --git a/tests/unit_tests/model_validation/test_TimeSeriesR2SquareBySegments.py b/tests/unit_tests/model_validation/test_TimeSeriesR2SquareBySegments.py
index 9d339e829..7fe4fec33 100644
--- a/tests/unit_tests/model_validation/test_TimeSeriesR2SquareBySegments.py
+++ b/tests/unit_tests/model_validation/test_TimeSeriesR2SquareBySegments.py
@@ -46,15 +46,20 @@ def setUp(self):
 
     def test_return_types(self):
         """Test if function returns expected types."""
-        fig, results_df = TimeSeriesR2SquareBySegments(self.vm_dataset, self.vm_model)
+        fig, results_df, raw_data = TimeSeriesR2SquareBySegments(
+            self.vm_dataset, self.vm_model
+        )
 
         # Check return types
         self.assertIsInstance(fig, go.Figure)
         self.assertIsInstance(results_df, pd.DataFrame)
+        self.assertIsInstance(raw_data, vm.RawData)
 
     def test_results_dataframe(self):
         """Test if results DataFrame has expected structure."""
-        _, results_df = TimeSeriesR2SquareBySegments(self.vm_dataset, self.vm_model)
+        _, results_df, raw_data = TimeSeriesR2SquareBySegments(
+            self.vm_dataset, self.vm_model
+        )
 
         # Check columns
         expected_columns = ["Segments", "Start Date", "End Date", "R-Squared"]
@@ -78,7 +83,7 @@ def test_custom_segments(self):
             "end_date": [dates[32], dates[65], dates[-1]],
         }
 
-        _, results_df = TimeSeriesR2SquareBySegments(
+        _, results_df, raw_data = TimeSeriesR2SquareBySegments(
             self.vm_dataset, self.vm_model, segments=custom_segments
         )
 

From 18c9a0db3f35c743cfddb682f6fa1d80834ae3db Mon Sep 17 00:00:00 2001
From: Juan <juan@validmind.ai>
Date: Tue, 18 Feb 2025 16:59:12 +0100
Subject: [PATCH 09/10] Remove RawData from RobustnessDiagnosis

---
 .../tests/model_validation/sklearn/RobustnessDiagnosis.py    | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/validmind/tests/model_validation/sklearn/RobustnessDiagnosis.py b/validmind/tests/model_validation/sklearn/RobustnessDiagnosis.py
index 4b0710211..604e74815 100644
--- a/validmind/tests/model_validation/sklearn/RobustnessDiagnosis.py
+++ b/validmind/tests/model_validation/sklearn/RobustnessDiagnosis.py
@@ -14,7 +14,7 @@
 
 from validmind.errors import MissingOrInvalidModelPredictFnError
 from validmind.logging import get_logger
-from validmind.tests import RawData, tags, tasks
+from validmind.tests import tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
 logger = get_logger(__name__)
@@ -331,7 +331,4 @@ def RobustnessDiagnosis(
         results_df,
         fig,
         all(results_df["Passed"]),
-        RawData(
-            model=model.input_id, datasets=[dataset.input_id for dataset in datasets]
-        ),
     )

From 6233b22344e37d7a80a8c093911fa421d8f29df5 Mon Sep 17 00:00:00 2001
From: Juan <juan@validmind.ai>
Date: Fri, 28 Feb 2025 13:58:27 +0100
Subject: [PATCH 10/10] 2.8.12

---
 pyproject.toml           | 2 +-
 validmind/__version__.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index e4362b32c..97133b421 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -10,7 +10,7 @@ description = "ValidMind Library"
 license = "Commercial License"
 name = "validmind"
 readme = "README.pypi.md"
-version = "2.8.11"
+version = "2.8.12"
 
 [tool.poetry.dependencies]
 aiohttp = {extras = ["speedups"], version = "*"}
diff --git a/validmind/__version__.py b/validmind/__version__.py
index bb309c0ec..7953a978c 100644
--- a/validmind/__version__.py
+++ b/validmind/__version__.py
@@ -1 +1 @@
-__version__ = "2.8.11"
+__version__ = "2.8.12"