validmind 2.8.10__py3-none-any.whl → 2.8.12__py3-none-any.whl

validmind/__version__.py

@@ -1 +1 @@
-__version__ = "2.8.10"
+__version__ = "2.8.12"

validmind/ai/test_descriptions.py

@@ -37,8 +37,10 @@ def _get_llm_global_context():
     return context if context_enabled and context else None
 
 
-def _truncate_summary(summary: str, test_id: str, max_tokens: int = 100_000):
-    if len(summary) < max_tokens:
+def _truncate_summary(
+    summary: Union[str, None], test_id: str, max_tokens: int = 100_000
+):
+    if summary is None or len(summary) < max_tokens:
         # since string itself is less than max_tokens, definitely small enough
         return summary
 

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),
+    )

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)

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,
+        ),
     )

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
+    )

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),
     )

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)

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)

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):

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),
+    )

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),
     )

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
         ),
     )

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)

validmind/tests/data_validation/IsolationForestOutliers.py

@@ -8,7 +8,7 @@ import matplotlib.pyplot as plt
 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),
+    )

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)

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)

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)

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),
+    )

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
+        ),
     )

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,
     )

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)

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,
         ),
     )

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))

validmind/tests/data_validation/ProtectedClassesDisparity.py

@@ -7,7 +7,7 @@ import sys
 
 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),
+    )

validmind/tests/data_validation/ProtectedClassesThresholdOptimizer.py

@@ -8,7 +8,7 @@ import sys
 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):

validmind/tests/data_validation/RollingStatsPlot.py

@@ -113,6 +113,7 @@ 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,
         ),
     )

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
+    )

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))

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)

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))

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))

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)

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),
+    )

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),
     )

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),
+    )

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))

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)

validmind/tests/data_validation/ZivotAndrewsArch.py

@@ -6,7 +6,7 @@ import pandas as pd
 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),
+    )

validmind/tests/data_validation/nlp/CommonWords.py

@@ -94,4 +94,4 @@ def CommonWords(dataset: VMDataset):
         xaxis_tickangle=-45,
     )
 
-    return fig, RawData(words=x, frequencies=y)
+    return fig, RawData(words=x, frequencies=y, dataset=dataset.input_id)

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)

validmind/tests/data_validation/nlp/LanguageDetection.py

@@ -71,5 +71,5 @@ def LanguageDetection(dataset):
             title="Language Distribution",
             labels={"x": "Language Codes"},
         ),
-        RawData(detected_languages=languages),
+        RawData(detected_languages=languages, dataset=dataset.input_id),
     )

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),
     )

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),
+    )

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):

validmind/tests/data_validation/nlp/Sentiment.py

@@ -77,4 +77,6 @@ def Sentiment(dataset):
 
     plt.close("all")
 
-    return fig, RawData(sentiment_scores=vader_sentiment.tolist())
+    return fig, RawData(
+        sentiment_scores=vader_sentiment.tolist(), dataset=dataset.input_id
+    )

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))

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)

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
+        ),
+    )

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
+        ),
+    )

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
+    )

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,
+        ),
+    )