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

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

@@ -14,7 +14,9 @@ from validmind.errors import SkipTestError
 from validmind.vm_models import VMDataset
 
 
-@tags("tabular_data", "binary_classification", "multiclass_classification")
+@tags(
+    "tabular_data", "binary_classification", "multiclass_classification", "data_quality"
+)
 @tasks("classification")
 def ClassImbalance(
     dataset: VMDataset, min_percent_threshold: int = 10
@@ -104,5 +106,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

@@ -6,12 +6,10 @@ import re
 from collections import Counter
 
 import numpy as np
-from ydata_profiling.config import Settings
-from ydata_profiling.model.typeset import ProfilingTypeSet
 
 from validmind import RawData, tags, tasks
-from validmind.errors import UnsupportedColumnTypeError
 from validmind.logging import get_logger
+from validmind.utils import infer_datatypes
 from validmind.vm_models import VMDataset
 
 DEFAULT_HISTOGRAM_BINS = 10
@@ -20,25 +18,6 @@ DEFAULT_HISTOGRAM_BIN_SIZES = [5, 10, 20, 50]
 logger = get_logger(__name__)
 
 
-def infer_datatypes(df):
-    column_type_mappings = {}
-    typeset = ProfilingTypeSet(Settings())
-    variable_types = typeset.infer_type(df)
-
-    for column, type in variable_types.items():
-        if str(type) == "Unsupported":
-            if df[column].isnull().all():
-                column_type_mappings[column] = {"id": column, "type": "Null"}
-            else:
-                raise UnsupportedColumnTypeError(
-                    f"Unsupported type for column {column}. Please review all values in this dataset column."
-                )
-        else:
-            column_type_mappings[column] = {"id": column, "type": str(type)}
-
-    return list(column_type_mappings.values())
-
-
 def get_numerical_histograms(df, column):
     """
     Returns a collection of histograms for a numerical column, each one
@@ -50,7 +29,7 @@ def get_numerical_histograms(df, column):
     # bins='sturges'. Cannot use 'auto' until we review and fix its performance
     #  on datasets with too many unique values
     #
-    # 'sturges': R’s default method, only accounts for data size. Only optimal
+    # 'sturges': R's default method, only accounts for data size. Only optimal
     # for gaussian data and underestimates number of bins for large non-gaussian datasets.
     default_hist = np.histogram(values_cleaned, bins="sturges")
 
@@ -242,4 +221,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/DescriptiveStatistics.py

@@ -44,7 +44,7 @@ def get_summary_statistics_categorical(df, categorical_fields):
     return summary_stats
 
 
-@tags("tabular_data", "time_series_data")
+@tags("tabular_data", "time_series_data", "data_quality")
 @tasks("classification", "regression")
 def DescriptiveStatistics(dataset: VMDataset):
     """

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/Skewness.py

@@ -2,10 +2,8 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
-from ydata_profiling.config import Settings
-from ydata_profiling.model.typeset import ProfilingTypeSet
-
 from validmind import tags, tasks
+from validmind.utils import infer_datatypes
 
 
 @tags("data_quality", "tabular_data")
@@ -49,8 +47,11 @@ def Skewness(dataset, max_threshold=1):
     - Subjective threshold for risk grading, requiring expert input and recurrent iterations for refinement.
     """
 
-    typeset = ProfilingTypeSet(Settings())
-    dataset_types = typeset.infer_type(dataset.df)
+    # Use the imported infer_datatypes function
+    dataset_types = infer_datatypes(dataset.df)
+
+    # Convert the list of dictionaries to a dictionary for easy access
+    dataset_types_dict = {item["id"]: item["type"] for item in dataset_types}
 
     skewness = dataset.df.skew(numeric_only=True)
 
@@ -58,7 +59,7 @@ def Skewness(dataset, max_threshold=1):
     passed = True
 
     for col in skewness.index:
-        if str(dataset_types[col]) != "Numeric":
+        if dataset_types_dict.get(col) != "Numeric":
             continue
 
         col_skewness = skewness[col]

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