validmind 2.8.27__py3-none-any.whl → 2.8.29__py3-none-any.whl

validmind/__version__.py

@@ -1 +1 @@
-__version__ = "2.8.27"
+__version__ = "2.8.28"

validmind/ai/utils.py

@@ -130,7 +130,7 @@ def get_judge_config(judge_llm=None, judge_embeddings=None):
                     "ensure that you are connected to the ValidMind API and confirm ValidMind AI is enabled for your account."
                 )
         if isinstance(judge_embeddings, FunctionModel) and judge_embeddings is not None:
-            if isinstance(judge_llm.model, BaseChatModel):
+            if isinstance(judge_embeddings.model, Embeddings):
                 judge_embeddings = judge_embeddings.model
             else:
                 raise ValueError(

validmind/models/function.py

@@ -35,7 +35,8 @@ class FunctionModel(VMModel):
 
     Attributes:
         predict_fn (callable): The predict function that should take a dictionary of
-            input features and return a prediction.
+            input features and return a prediction. Can return simple values or
+            dictionary objects.
         input_id (str, optional): The input ID for the model. Defaults to None.
         name (str, optional): The name of the model. Defaults to the name of the predict_fn.
         prompt (Prompt, optional): If using a prompt, the prompt object that defines the template
@@ -55,6 +56,13 @@ class FunctionModel(VMModel):
             X (pandas.DataFrame): The input features to predict on
 
         Returns:
-            List[Any]: The predictions
+            List[Any]: The predictions. Can contain simple values or dictionary objects
+                       depending on what the predict_fn returns.
         """
-        return [self.predict_fn(x) for x in X.to_dict(orient="records")]
+        predictions = []
+        for x in X.to_dict(orient="records"):
+            result = self.predict_fn(x)
+            # Handle both simple values and complex dictionary returns
+            predictions.append(result)
+
+        return predictions

validmind/tests/data_validation/ACFandPACFPlot.py

@@ -2,6 +2,8 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
+from typing import Tuple
+
 import pandas as pd
 import plotly.graph_objects as go
 from statsmodels.tsa.stattools import acf, pacf
@@ -12,7 +14,7 @@ from validmind.vm_models import VMDataset
 
 @tags("time_series_data", "forecasting", "statistical_test", "visualization")
 @tasks("regression")
-def ACFandPACFPlot(dataset: VMDataset):
+def ACFandPACFPlot(dataset: VMDataset) -> Tuple[go.Figure, RawData]:
     """
     Analyzes time series data using Autocorrelation Function (ACF) and Partial Autocorrelation Function (PACF) plots to
     reveal trends and correlations.

validmind/tests/data_validation/ADF.py

@@ -2,6 +2,8 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
+from typing import Dict
+
 import pandas as pd
 from statsmodels.tsa.stattools import adfuller
 
@@ -16,7 +18,7 @@ logger = get_logger(__name__)
     "time_series_data", "statsmodels", "forecasting", "statistical_test", "stationarity"
 )
 @tasks("regression")
-def ADF(dataset: VMDataset):
+def ADF(dataset: VMDataset) -> Dict[str, pd.DataFrame]:
     """
     Assesses the stationarity of a time series dataset using the Augmented Dickey-Fuller (ADF) test.
 

validmind/tests/data_validation/AutoAR.py

@@ -2,6 +2,8 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
+from typing import Dict
+
 import pandas as pd
 from statsmodels.tsa.ar_model import AutoReg
 from statsmodels.tsa.stattools import adfuller
@@ -15,7 +17,7 @@ logger = get_logger(__name__)
 
 @tags("time_series_data", "statsmodels", "forecasting", "statistical_test")
 @tasks("regression")
-def AutoAR(dataset: VMDataset, max_ar_order: int = 3):
+def AutoAR(dataset: VMDataset, max_ar_order: int = 3) -> Dict[str, pd.DataFrame]:
     """
     Automatically identifies the optimal Autoregressive (AR) order for a time series using BIC and AIC criteria.
 

validmind/tests/data_validation/AutoMA.py

@@ -2,6 +2,8 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
+from typing import Dict, Tuple
+
 import pandas as pd
 from statsmodels.tsa.arima.model import ARIMA
 from statsmodels.tsa.stattools import adfuller
@@ -15,7 +17,9 @@ logger = get_logger(__name__)
 
 @tags("time_series_data", "statsmodels", "forecasting", "statistical_test")
 @tasks("regression")
-def AutoMA(dataset: VMDataset, max_ma_order: int = 3):
+def AutoMA(
+    dataset: VMDataset, max_ma_order: int = 3
+) -> Tuple[Dict[str, pd.DataFrame], RawData]:
     """
     Automatically selects the optimal Moving Average (MA) order for each variable in a time series dataset based on
     minimal BIC and AIC values.

validmind/tests/data_validation/AutoStationarity.py

@@ -2,6 +2,8 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
+from typing import Dict
+
 import numpy as np
 import pandas as pd
 from statsmodels.tsa.stattools import adfuller
@@ -12,7 +14,9 @@ from validmind.vm_models import VMDataset
 
 @tags("time_series_data", "statsmodels", "forecasting", "statistical_test")
 @tasks("regression")
-def AutoStationarity(dataset: VMDataset, max_order: int = 5, threshold: float = 0.05):
+def AutoStationarity(
+    dataset: VMDataset, max_order: int = 5, threshold: float = 0.05
+) -> Dict[str, pd.DataFrame]:
     """
     Automates Augmented Dickey-Fuller test to assess stationarity across multiple time series in a DataFrame.
 

validmind/tests/data_validation/BivariateScatterPlots.py

@@ -3,15 +3,17 @@
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
 import itertools
+from typing import Tuple
 
 import plotly.express as px
+import plotly.graph_objects as go
 
 from validmind import RawData, tags, tasks
 
 
 @tags("tabular_data", "numerical_data", "visualization")
 @tasks("classification")
-def BivariateScatterPlots(dataset):
+def BivariateScatterPlots(dataset) -> Tuple[go.Figure, RawData]:
     """
     Generates bivariate scatterplots to visually inspect relationships between pairs of numerical predictor variables
     in machine learning classification tasks.

validmind/tests/data_validation/BoxPierce.py

@@ -2,6 +2,9 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
+
+from typing import Tuple
+
 import pandas as pd
 from statsmodels.stats.diagnostic import acorr_ljungbox
 
@@ -10,7 +13,7 @@ from validmind import RawData, tags, tasks
 
 @tasks("regression")
 @tags("time_series_data", "forecasting", "statistical_test", "statsmodels")
-def BoxPierce(dataset):
+def BoxPierce(dataset) -> Tuple[pd.DataFrame, RawData]:
     """
     Detects autocorrelation in time-series data through the Box-Pierce test to validate model performance.
 

validmind/tests/data_validation/ChiSquaredFeaturesTable.py

@@ -12,7 +12,7 @@ from validmind.errors import SkipTestError
 
 @tags("tabular_data", "categorical_data", "statistical_test")
 @tasks("classification")
-def ChiSquaredFeaturesTable(dataset, p_threshold=0.05):
+def ChiSquaredFeaturesTable(dataset, p_threshold=0.05) -> pd.DataFrame:
     """
     Assesses the statistical association between categorical features and a target variable using the Chi-Squared test.
 

validmind/tests/data_validation/ClassImbalance.py

@@ -20,7 +20,7 @@ from validmind.vm_models import VMDataset
 @tasks("classification")
 def ClassImbalance(
     dataset: VMDataset, min_percent_threshold: int = 10
-) -> Tuple[Dict[str, Any], go.Figure, bool]:
+) -> Tuple[Dict[str, Any], go.Figure, bool, RawData]:
     """
     Evaluates and quantifies class distribution imbalance in a dataset used by a machine learning model.
 

validmind/tests/data_validation/DatasetDescription.py

@@ -4,6 +4,7 @@
 
 import re
 from collections import Counter
+from typing import Any, Dict, List, Tuple
 
 import numpy as np
 
@@ -142,7 +143,9 @@ def describe_column(df, column):
 
 @tags("tabular_data", "time_series_data", "text_data")
 @tasks("classification", "regression", "text_classification", "text_summarization")
-def DatasetDescription(dataset: VMDataset):
+def DatasetDescription(
+    dataset: VMDataset,
+) -> Tuple[Dict[str, List[Dict[str, Any]]], RawData]:
     """
     Provides comprehensive analysis and statistical summaries of each column in a machine learning model's dataset.
 

validmind/tests/data_validation/DatasetSplit.py

@@ -2,7 +2,8 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
-from typing import List
+
+from typing import Any, Dict, List, Tuple
 
 from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset
@@ -17,7 +18,7 @@ DATASET_LABELS = {
 
 @tags("tabular_data", "time_series_data", "text_data")
 @tasks("classification", "regression", "text_classification", "text_summarization")
-def DatasetSplit(datasets: List[VMDataset]):
+def DatasetSplit(datasets: List[VMDataset]) -> Tuple[List[Dict[str, Any]], RawData]:
     """
     Evaluates and visualizes the distribution proportions among training, testing, and validation datasets of an ML
     model.

validmind/tests/data_validation/DescriptiveStatistics.py

@@ -2,6 +2,8 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
+from typing import Any, Dict
+
 import pandas as pd
 
 from validmind import tags, tasks
@@ -46,7 +48,7 @@ def get_summary_statistics_categorical(df, categorical_fields):
 
 @tags("tabular_data", "time_series_data", "data_quality")
 @tasks("classification", "regression")
-def DescriptiveStatistics(dataset: VMDataset):
+def DescriptiveStatistics(dataset: VMDataset) -> Dict[str, Any]:
     """
     Performs a detailed descriptive statistical analysis of both numerical and categorical data within a model's
     dataset.

validmind/tests/data_validation/DickeyFullerGLS.py

@@ -2,6 +2,8 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
+from typing import Any, Dict, Tuple
+
 import pandas as pd
 from arch.unitroot import DFGLS
 from numpy.linalg import LinAlgError
@@ -16,7 +18,7 @@ logger = get_logger(__name__)
 
 @tags("time_series_data", "forecasting", "unit_root_test")
 @tasks("regression")
-def DickeyFullerGLS(dataset: VMDataset):
+def DickeyFullerGLS(dataset: VMDataset) -> Tuple[Dict[str, Any], RawData]:
     """
     Assesses stationarity in time series data using the Dickey-Fuller GLS test to determine the order of integration.
 

validmind/tests/data_validation/Duplicates.py

@@ -2,6 +2,8 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
+from typing import Any, Dict, Tuple
+
 import pandas as pd
 
 from validmind import tags, tasks
@@ -9,7 +11,7 @@ from validmind import tags, tasks
 
 @tags("tabular_data", "data_quality", "text_data")
 @tasks("classification", "regression")
-def Duplicates(dataset, min_threshold=1):
+def Duplicates(dataset, min_threshold=1) -> Tuple[Dict[str, Any], bool]:
     """
     Tests dataset for duplicate entries, ensuring model reliability via data quality verification.
 

validmind/tests/data_validation/EngleGrangerCoint.py

@@ -2,6 +2,9 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
+
+from typing import Dict
+
 import pandas as pd
 from statsmodels.tsa.stattools import coint
 
@@ -12,7 +15,9 @@ from validmind.vm_models import VMDataset
 
 @tags("time_series_data", "statistical_test", "forecasting")
 @tasks("regression")
-def EngleGrangerCoint(dataset: VMDataset, threshold: float = 0.05):
+def EngleGrangerCoint(
+    dataset: VMDataset, threshold: float = 0.05
+) -> Dict[str, pd.DataFrame]:
     """
     Assesses the degree of co-movement between pairs of time series data using the Engle-Granger cointegration test.
 

validmind/tests/data_validation/FeatureTargetCorrelationPlot.py

@@ -3,6 +3,8 @@
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
 
+from typing import Tuple
+
 import numpy as np
 import plotly.graph_objects as go
 
@@ -11,7 +13,7 @@ from validmind import RawData, tags, tasks
 
 @tags("tabular_data", "visualization", "correlation")
 @tasks("classification", "regression")
-def FeatureTargetCorrelationPlot(dataset, fig_height=600):
+def FeatureTargetCorrelationPlot(dataset, fig_height=600) -> Tuple[go.Figure, RawData]:
     """
     Visualizes the correlation between input features and the model's target output in a color-coded horizontal bar
     plot.

validmind/tests/data_validation/HighCardinality.py

@@ -2,6 +2,8 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
+from typing import Any, Dict, List, Tuple
+
 from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset
 
@@ -13,7 +15,7 @@ def HighCardinality(
     num_threshold: int = 100,
     percent_threshold: float = 0.1,
     threshold_type: str = "percent",
-):
+) -> Tuple[List[Dict[str, Any]], bool, RawData]:
     """
     Assesses the number of unique values in categorical columns to detect high cardinality and potential overfitting.
 

validmind/tests/data_validation/HighPearsonCorrelation.py

@@ -2,6 +2,9 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
+
+from typing import Any, Dict, List, Tuple
+
 from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset
 
@@ -13,7 +16,7 @@ def HighPearsonCorrelation(
     max_threshold: float = 0.3,
     top_n_correlations: int = 10,
     feature_columns: list = None,
-):
+) -> Tuple[List[Dict[str, Any]], bool, RawData]:
     """
     Identifies highly correlated feature pairs in a dataset suggesting feature redundancy or multicollinearity.
 

validmind/tests/data_validation/IQROutliersBarPlot.py

@@ -2,6 +2,9 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
+
+from typing import Tuple
+
 import plotly.graph_objects as go
 
 from validmind import RawData, tags, tasks
@@ -22,7 +25,7 @@ def compute_outliers(series, threshold):
 @tasks("classification", "regression")
 def IQROutliersBarPlot(
     dataset: VMDataset, threshold: float = 1.5, fig_width: int = 800
-):
+) -> Tuple[go.Figure, RawData]:
     """
     Visualizes outlier distribution across percentiles in numerical data using the Interquartile Range (IQR) method.
 

validmind/tests/data_validation/IQROutliersTable.py

@@ -2,6 +2,9 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
+
+from typing import Any, Dict, Tuple
+
 from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset
 
@@ -18,7 +21,9 @@ def compute_outliers(series, threshold=1.5):
 
 @tags("tabular_data", "numerical_data")
 @tasks("classification", "regression")
-def IQROutliersTable(dataset: VMDataset, threshold: float = 1.5):
+def IQROutliersTable(
+    dataset: VMDataset, threshold: float = 1.5
+) -> Tuple[Dict[str, Any], RawData]:
     """
     Determines and summarizes outliers in numerical features using the Interquartile Range method.
 

validmind/tests/data_validation/IsolationForestOutliers.py

@@ -3,7 +3,9 @@
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
 import itertools
+from typing import Tuple
 
+import matplotlib.figure
 import matplotlib.pyplot as plt
 import seaborn as sns
 from sklearn.ensemble import IsolationForest
@@ -19,7 +21,7 @@ def IsolationForestOutliers(
     random_state: int = 0,
     contamination: float = 0.1,
     feature_columns: list = None,
-):
+) -> Tuple[matplotlib.figure.Figure, RawData]:
     """
     Detects outliers in a dataset using the Isolation Forest algorithm and visualizes results through scatter plots.
 

validmind/tests/data_validation/JarqueBera.py

@@ -2,6 +2,8 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
+from typing import Tuple
+
 import pandas as pd
 from statsmodels.stats.stattools import jarque_bera
 
@@ -10,7 +12,7 @@ from validmind import RawData, tags, tasks
 
 @tasks("classification", "regression")
 @tags("tabular_data", "data_distribution", "statistical_test", "statsmodels")
-def JarqueBera(dataset):
+def JarqueBera(dataset) -> Tuple[pd.DataFrame, RawData]:
     """
     Assesses normality of dataset features in an ML model using the Jarque-Bera test.
 

validmind/tests/data_validation/KPSS.py

@@ -2,6 +2,8 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
+from typing import Any, Dict
+
 import pandas as pd
 from statsmodels.tsa.stattools import kpss
 
@@ -15,7 +17,7 @@ logger = get_logger(__name__)
 
 @tags("time_series_data", "stationarity", "unit_root_test", "statsmodels")
 @tasks("data_validation")
-def KPSS(dataset: VMDataset):
+def KPSS(dataset: VMDataset) -> Dict[str, Any]:
     """
     Assesses the stationarity of time-series data in a machine learning model using the KPSS unit root test.
 

validmind/tests/data_validation/LJungBox.py

@@ -2,6 +2,8 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
+from typing import Tuple
+
 import pandas as pd
 from statsmodels.stats.diagnostic import acorr_ljungbox
 
@@ -10,7 +12,7 @@ from validmind import RawData, tags, tasks
 
 @tasks("regression")
 @tags("time_series_data", "forecasting", "statistical_test", "statsmodels")
-def LJungBox(dataset):
+def LJungBox(dataset) -> Tuple[pd.DataFrame, RawData]:
     """
     Assesses autocorrelations in dataset features by performing a Ljung-Box test on each feature.
 

validmind/tests/data_validation/LaggedCorrelationHeatmap.py

@@ -2,9 +2,12 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
+from typing import Tuple
+
 import numpy as np
 import pandas as pd
 import plotly.figure_factory as ff
+import plotly.graph_objects as go
 
 from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset
@@ -15,7 +18,9 @@ COOLWARM = [[0, "rgb(95,5,255)"], [0.5, "rgb(255,255,255)"], [1, "rgb(255,5,0)"]
 
 @tags("time_series_data", "visualization")
 @tasks("regression")
-def LaggedCorrelationHeatmap(dataset: VMDataset, num_lags: int = 10):
+def LaggedCorrelationHeatmap(
+    dataset: VMDataset, num_lags: int = 10
+) -> Tuple[go.Figure, RawData]:
     """
     Assesses and visualizes correlation between target variable and lagged independent variables in a time-series
     dataset.

validmind/tests/data_validation/MissingValues.py

@@ -2,13 +2,17 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
+from typing import Any, Dict, List, Tuple
+
 from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset
 
 
 @tags("tabular_data", "data_quality")
 @tasks("classification", "regression")
-def MissingValues(dataset: VMDataset, min_threshold: int = 1):
+def MissingValues(
+    dataset: VMDataset, min_threshold: int = 1
+) -> Tuple[List[Dict[str, Any]], bool, RawData]:
     """
     Evaluates dataset quality by ensuring missing value ratio across all features does not exceed a set threshold.
 

validmind/tests/data_validation/MissingValuesBarPlot.py

@@ -2,6 +2,8 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
+from typing import Tuple
+
 import plotly.graph_objects as go
 
 from validmind import RawData, tags, tasks
@@ -12,7 +14,7 @@ from validmind.vm_models import VMDataset
 @tasks("classification", "regression")
 def MissingValuesBarPlot(
     dataset: VMDataset, threshold: int = 80, fig_height: int = 600
-):
+) -> Tuple[go.Figure, RawData]:
     """
     Assesses the percentage and distribution of missing values in the dataset via a bar plot, with emphasis on
     identifying high-risk columns based on a user-defined threshold.

validmind/tests/data_validation/MutualInformation.py

@@ -2,6 +2,9 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
+
+from typing import Tuple
+
 import plotly.graph_objects as go
 from sklearn.feature_selection import mutual_info_classif, mutual_info_regression
 
@@ -14,7 +17,7 @@ from validmind.vm_models.result import RawData
 @tasks("classification", "regression")
 def MutualInformation(
     dataset: VMDataset, min_threshold: float = 0.01, task: str = "classification"
-):
+) -> Tuple[go.Figure, RawData]:
     """
     Calculates mutual information scores between features and target variable to evaluate feature relevance.
 

validmind/tests/data_validation/PearsonCorrelationMatrix.py

@@ -3,6 +3,8 @@
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
 
+from typing import Tuple
+
 import plotly.graph_objects as go
 
 from validmind import RawData, tags, tasks
@@ -10,7 +12,7 @@ from validmind import RawData, tags, tasks
 
 @tags("tabular_data", "numerical_data", "correlation")
 @tasks("classification", "regression")
-def PearsonCorrelationMatrix(dataset):
+def PearsonCorrelationMatrix(dataset) -> Tuple[go.Figure, RawData]:
     """
     Evaluates linear dependency between numerical variables in a dataset via a Pearson Correlation coefficient heat map.
 

validmind/tests/data_validation/PhillipsPerronArch.py

@@ -2,6 +2,8 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
+from typing import Any, Dict
+
 import numpy as np
 import pandas as pd
 from arch.unitroot import PhillipsPerron
@@ -17,7 +19,7 @@ logger = get_logger(__name__)
 
 @tags("time_series_data", "forecasting", "statistical_test", "unit_root_test")
 @tasks("regression")
-def PhillipsPerronArch(dataset: VMDataset):
+def PhillipsPerronArch(dataset: VMDataset) -> Dict[str, Any]:
     """
     Assesses the stationarity of time series data in each feature of the ML model using the Phillips-Perron test.
 

validmind/tests/data_validation/ProtectedClassesCombination.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
 import sys
+from typing import Dict, Tuple
 
 import pandas as pd
 import plotly.graph_objects as go
@@ -11,6 +12,7 @@ import plotly.subplots as sp
 from validmind import RawData, tags, tasks
 from validmind.errors import MissingDependencyError
 from validmind.logging import get_logger
+from validmind.vm_models import VMDataset, VMModel
 
 try:
     from fairlearn.metrics import (
@@ -33,7 +35,9 @@ logger = get_logger(__name__)
 
 @tags("bias_and_fairness")
 @tasks("classification", "regression")
-def ProtectedClassesCombination(dataset, model, protected_classes=None):
+def ProtectedClassesCombination(
+    dataset: VMDataset, model: VMModel, protected_classes=None
+) -> Tuple[Dict[str, pd.DataFrame], Dict[str, pd.DataFrame], go.Figure, RawData]:
     """
     Visualizes combinations of protected classes and their corresponding error metric differences.
 

validmind/tests/data_validation/ProtectedClassesDescription.py

@@ -3,6 +3,8 @@
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
 
+from typing import List, Tuple, Union
+
 import pandas as pd
 import plotly.graph_objects as go
 
@@ -14,7 +16,9 @@ logger = get_logger(__name__)
 
 @tags("bias_and_fairness", "descriptive_statistics")
 @tasks("classification", "regression")
-def ProtectedClassesDescription(dataset, protected_classes=None):
+def ProtectedClassesDescription(
+    dataset, protected_classes=None
+) -> Tuple[pd.DataFrame, Union[go.Figure, List[go.Figure]], RawData]:
     """
     Visualizes the distribution of protected classes in the dataset relative to the target variable
     and provides descriptive statistics.