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

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, *tuple(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,
+        ),
+    )

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)

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, *tuple(figures), RawData(meteor_scores=metrics_df))
+    return (
+        result_df,
+        *figures,
+        RawData(
+            meteor_scores=metrics_df, model=model.input_id, dataset=dataset.input_id
+        ),
+    )

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

validmind/tests/model_validation/RegardScore.py

@@ -145,5 +145,10 @@ def RegardScore(dataset, model):
     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,
+        ),
     )

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

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

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

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

validmind/tests/model_validation/TimeSeriesR2SquareBySegments.py

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

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

validmind/tests/model_validation/ToxicityScore.py

@@ -141,10 +141,12 @@ def ToxicityScore(dataset, model):
 
     return (
         result_df,
-        *tuple(figures),
+        *figures,
         RawData(
             input_toxicity_df=input_df,
             true_toxicity_df=true_df,
             pred_toxicity_df=pred_df,
+            model=model.input_id,
+            dataset=dataset.input_id,
         ),
     )

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)

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

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

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

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

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

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)

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

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

validmind/tests/model_validation/embeddings/StabilityAnalysisKeyword.py

@@ -97,4 +97,8 @@ def StabilityAnalysisKeyword(
         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,
+    )

validmind/tests/model_validation/embeddings/StabilityAnalysisRandomNoise.py

@@ -151,4 +151,8 @@ def StabilityAnalysisRandomNoise(
         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,
+    )

validmind/tests/model_validation/embeddings/StabilityAnalysisSynonyms.py

@@ -107,4 +107,8 @@ def StabilityAnalysisSynonyms(
         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,
+    )

validmind/tests/model_validation/embeddings/StabilityAnalysisTranslation.py

@@ -134,4 +134,8 @@ def StabilityAnalysisTranslation(
         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,
+    )

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