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

validmind/tests/model_validation/ragas/SemanticSimilarity.py

@@ -3,12 +3,15 @@
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
 import warnings
+from typing import Dict, Tuple
 
 import plotly.express as px
+import plotly.graph_objects as go
 from datasets import Dataset
 
 from validmind import RawData, tags, tasks
 from validmind.errors import MissingDependencyError
+from validmind.vm_models import VMDataset
 
 from .utils import get_ragas_config, get_renamed_columns
 
@@ -30,12 +33,12 @@ except ImportError as e:
 @tags("ragas", "llm")
 @tasks("text_qa", "text_generation", "text_summarization")
 def SemanticSimilarity(
-    dataset,
-    response_column="response",
-    reference_column="reference",
+    dataset: VMDataset,
+    response_column: str = "response",
+    reference_column: str = "reference",
     judge_llm=None,
     judge_embeddings=None,
-):
+) -> Tuple[Dict[str, list], go.Figure, go.Figure, RawData]:
     """
     Calculates the semantic similarity between generated responses and ground truths
 

validmind/tests/model_validation/sklearn/AdjustedMutualInformation.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, List, Tuple
+
 from sklearn.metrics import adjusted_mutual_info_score
 
 from validmind import RawData, tags, tasks
@@ -10,7 +12,9 @@ from validmind.vm_models import VMDataset, VMModel
 
 @tags("sklearn", "model_performance", "clustering")
 @tasks("clustering")
-def AdjustedMutualInformation(model: VMModel, dataset: VMDataset):
+def AdjustedMutualInformation(
+    model: VMModel, dataset: VMDataset
+) -> Tuple[List[Dict[str, float]], RawData]:
     """
     Evaluates clustering model performance by measuring mutual information between true and predicted labels, adjusting
     for chance.

validmind/tests/model_validation/sklearn/AdjustedRandIndex.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, List, Tuple
+
 from sklearn.metrics import adjusted_rand_score
 
 from validmind import RawData, tags, tasks
@@ -10,7 +12,9 @@ from validmind.vm_models import VMDataset, VMModel
 
 @tags("sklearn", "model_performance", "clustering")
 @tasks("clustering")
-def AdjustedRandIndex(model: VMModel, dataset: VMDataset):
+def AdjustedRandIndex(
+    model: VMModel, dataset: VMDataset
+) -> Tuple[List[Dict[str, float]], RawData]:
     """
     Measures the similarity between two data clusters using the Adjusted Rand Index (ARI) metric in clustering machine
     learning models.

validmind/tests/model_validation/sklearn/CalibrationCurve.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 sklearn.calibration import calibration_curve
 
@@ -12,7 +14,9 @@ from validmind.vm_models.result import RawData
 
 @tags("sklearn", "model_performance", "classification")
 @tasks("classification")
-def CalibrationCurve(model: VMModel, dataset: VMDataset, n_bins: int = 10):
+def CalibrationCurve(
+    model: VMModel, dataset: VMDataset, n_bins: int = 10
+) -> Tuple[go.Figure, RawData]:
     """
     Evaluates the calibration of probability estimates by comparing predicted probabilities against observed
     frequencies.

validmind/tests/model_validation/sklearn/ClassifierPerformance.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, List
+
 import numpy as np
 from sklearn.metrics import classification_report, roc_auc_score
 from sklearn.preprocessing import LabelBinarizer
@@ -20,7 +22,9 @@ def multiclass_roc_auc_score(y_test, y_pred, average="macro"):
     "sklearn", "binary_classification", "multiclass_classification", "model_performance"
 )
 @tasks("classification", "text_classification")
-def ClassifierPerformance(dataset: VMDataset, model: VMModel, average: str = "macro"):
+def ClassifierPerformance(
+    dataset: VMDataset, model: VMModel, average: str = "macro"
+) -> Dict[str, List[Dict[str, float]]]:
     """
     Evaluates performance of binary or multiclass classification models using precision, recall, F1-Score, accuracy,
     and ROC AUC scores.

validmind/tests/model_validation/sklearn/ClusterCosineSimilarity.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, List, Tuple
+
 import numpy as np
 from sklearn.metrics.pairwise import cosine_similarity
 
@@ -12,7 +14,9 @@ from validmind.vm_models import VMDataset, VMModel
 
 @tags("sklearn", "model_performance", "clustering")
 @tasks("clustering")
-def ClusterCosineSimilarity(model: VMModel, dataset: VMDataset):
+def ClusterCosineSimilarity(
+    model: VMModel, dataset: VMDataset
+) -> Tuple[List[Dict[str, float]], RawData]:
     """
     Measures the intra-cluster similarity of a clustering model using cosine similarity.
 

validmind/tests/model_validation/sklearn/ClusterPerformanceMetrics.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, List, Tuple
+
 from sklearn.metrics import (
     adjusted_mutual_info_score,
     adjusted_rand_score,
@@ -69,7 +71,9 @@ identify members of the same class (precision) and the ability to capture all me
 
 @tags("sklearn", "model_performance", "clustering")
 @tasks("clustering")
-def ClusterPerformanceMetrics(model: VMModel, dataset: VMDataset):
+def ClusterPerformanceMetrics(
+    model: VMModel, dataset: VMDataset
+) -> Tuple[List[Dict[str, float]], RawData]:
     """
     Evaluates the performance of clustering machine learning models using multiple established metrics.
 

validmind/tests/model_validation/sklearn/CompletenessScore.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, List, Tuple
+
 from sklearn.metrics import completeness_score
 
 from validmind import RawData, tags, tasks
@@ -10,7 +12,9 @@ from validmind.vm_models import VMDataset, VMModel
 
 @tags("sklearn", "model_performance", "clustering")
 @tasks("clustering")
-def CompletenessScore(model: VMModel, dataset: VMDataset):
+def CompletenessScore(
+    model: VMModel, dataset: VMDataset
+) -> Tuple[List[Dict[str, float]], RawData]:
     """
     Evaluates a clustering model's capacity to categorize instances from a single class into the same cluster.
 

validmind/tests/model_validation/sklearn/ConfusionMatrix.py

@@ -3,8 +3,11 @@
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
 
+from typing import Tuple
+
 import numpy as np
 import plotly.figure_factory as ff
+import plotly.graph_objects as go
 from sklearn.metrics import confusion_matrix
 
 from validmind import RawData, tags, tasks
@@ -23,7 +26,7 @@ def ConfusionMatrix(
     dataset: VMDataset,
     model: VMModel,
     threshold: float = 0.5,
-):
+) -> Tuple[go.Figure, RawData]:
     """
     Evaluates and visually represents the classification ML model's predictive performance using a Confusion Matrix
     heatmap.

validmind/tests/model_validation/sklearn/FeatureImportance.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 sklearn.inspection import permutation_importance
 
@@ -11,7 +13,9 @@ from validmind.vm_models import VMDataset, VMModel
 
 @tags("model_explainability", "sklearn")
 @tasks("regression", "time_series_forecasting")
-def FeatureImportance(dataset: VMDataset, model: VMModel, num_features: int = 3):
+def FeatureImportance(
+    dataset: VMDataset, model: VMModel, num_features: int = 3
+) -> Tuple[pd.DataFrame, RawData]:
     """
     Compute feature importance scores for a given model and generate a summary table
     with the top important features.

validmind/tests/model_validation/sklearn/FowlkesMallowsScore.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, List, Tuple
+
 from sklearn import metrics
 
 from validmind import RawData, tags, tasks
@@ -10,7 +12,9 @@ from validmind.vm_models import VMDataset, VMModel
 
 @tags("sklearn", "model_performance")
 @tasks("clustering")
-def FowlkesMallowsScore(dataset: VMDataset, model: VMModel):
+def FowlkesMallowsScore(
+    dataset: VMDataset, model: VMModel
+) -> Tuple[List[Dict[str, float]], RawData]:
     """
     Evaluates the similarity between predicted and actual cluster assignments in a model using the Fowlkes-Mallows
     score.

validmind/tests/model_validation/sklearn/HomogeneityScore.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, List, Tuple
+
 from sklearn import metrics
 
 from validmind import RawData, tags, tasks
@@ -10,7 +12,9 @@ from validmind.vm_models import VMDataset, VMModel
 
 @tags("sklearn", "model_performance")
 @tasks("clustering")
-def HomogeneityScore(dataset: VMDataset, model: VMModel):
+def HomogeneityScore(
+    dataset: VMDataset, model: VMModel
+) -> Tuple[List[Dict[str, float]], RawData]:
     """
     Assesses clustering homogeneity by comparing true and predicted labels, scoring from 0 (heterogeneous) to 1
     (homogeneous).

validmind/tests/model_validation/sklearn/HyperParametersTuning.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 typing import Dict, List, Union
+from typing import Dict, List, Tuple, Union
 
 from sklearn.metrics import make_scorer, recall_score
 from sklearn.model_selection import GridSearchCV
@@ -11,8 +11,6 @@ from validmind import RawData, tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
 
-@tags("sklearn", "model_performance")
-@tasks("classification", "clustering")
 def custom_recall(y_true, y_pred_proba, threshold=0.5):
     y_pred = (y_pred_proba >= threshold).astype(int)
     return recall_score(y_true, y_pred)
@@ -65,7 +63,7 @@ def HyperParametersTuning(
     scoring: Union[str, List, Dict] = None,
     thresholds: Union[float, List[float]] = None,
     fit_params: dict = None,
-):
+) -> Tuple[List[Dict[str, float]], RawData]:
     """
     Performs exhaustive grid search over specified parameter ranges to find optimal model configurations
     across different metrics and decision thresholds.

validmind/tests/model_validation/sklearn/KMeansClustersOptimization.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 typing import List, Union
+from typing import List, Optional, Tuple
 
 import numpy as np
 import plotly.graph_objects as go
@@ -19,8 +19,8 @@ from validmind.vm_models import VMDataset, VMModel
 @tags("sklearn", "model_performance", "kmeans")
 @tasks("clustering")
 def KMeansClustersOptimization(
-    model: VMModel, dataset: VMDataset, n_clusters: Union[List[int], None] = None
-):
+    model: VMModel, dataset: VMDataset, n_clusters: Optional[List[int]] = None
+) -> Tuple[go.Figure, RawData]:
     """
     Optimizes the number of clusters in K-means models using Elbow and Silhouette methods.
 

validmind/tests/model_validation/sklearn/MinimumAccuracy.py

@@ -1,6 +1,8 @@
 # Copyright © 2023-2024 ValidMind Inc. All rights reserved.
 # 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, List, Tuple
+
 from sklearn.metrics import accuracy_score
 
 from validmind import RawData
@@ -12,7 +14,9 @@ from validmind.vm_models import VMDataset, VMModel
     "sklearn", "binary_classification", "multiclass_classification", "model_performance"
 )
 @tasks("classification", "text_classification")
-def MinimumAccuracy(dataset: VMDataset, model: VMModel, min_threshold: float = 0.7):
+def MinimumAccuracy(
+    dataset: VMDataset, model: VMModel, min_threshold: float = 0.7
+) -> Tuple[List[Dict[str, float]], bool, RawData]:
     """
     Checks if the model's prediction accuracy meets or surpasses a specified threshold.
 

validmind/tests/model_validation/sklearn/MinimumF1Score.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, List, Tuple
+
 import numpy as np
 from sklearn.metrics import f1_score
 
@@ -14,7 +16,9 @@ from validmind.vm_models import VMDataset, VMModel
     "sklearn", "binary_classification", "multiclass_classification", "model_performance"
 )
 @tasks("classification", "text_classification")
-def MinimumF1Score(dataset: VMDataset, model: VMModel, min_threshold: float = 0.5):
+def MinimumF1Score(
+    dataset: VMDataset, model: VMModel, min_threshold: float = 0.5
+) -> Tuple[List[Dict[str, float]], bool, RawData]:
     """
     Assesses if the model's F1 score on the validation set meets a predefined minimum threshold, ensuring balanced
     performance between precision and recall.

validmind/tests/model_validation/sklearn/MinimumROCAUCScore.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, List, Tuple
+
 import numpy as np
 from sklearn.metrics import roc_auc_score
 from sklearn.preprocessing import LabelBinarizer
@@ -14,7 +16,9 @@ from validmind.vm_models import VMDataset, VMModel
     "sklearn", "binary_classification", "multiclass_classification", "model_performance"
 )
 @tasks("classification", "text_classification")
-def MinimumROCAUCScore(dataset: VMDataset, model: VMModel, min_threshold: float = 0.5):
+def MinimumROCAUCScore(
+    dataset: VMDataset, model: VMModel, min_threshold: float = 0.5
+) -> Tuple[List[Dict[str, float]], bool, RawData]:
     """
     Validates model by checking if the ROC AUC score meets or surpasses a specified threshold.
 

validmind/tests/model_validation/sklearn/ModelParameters.py

@@ -2,14 +2,19 @@
 # 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, Optional
+
 import pandas as pd
 
 from validmind import tags, tasks
+from validmind.vm_models import VMModel
 
 
 @tags("model_training", "metadata")
 @tasks("classification", "regression")
-def ModelParameters(model, model_params=None):
+def ModelParameters(
+    model: VMModel, model_params: Optional[List[str]] = None
+) -> pd.DataFrame:
     """
     Extracts and displays model parameters in a structured format for transparency and reproducibility.
 

validmind/tests/model_validation/sklearn/ModelsPerformanceComparison.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, List
+
 import numpy as np
 from sklearn.metrics import classification_report
 
@@ -19,7 +21,9 @@ from .ClassifierPerformance import multiclass_roc_auc_score
     "model_comparison",
 )
 @tasks("classification", "text_classification")
-def ModelsPerformanceComparison(dataset: VMDataset, models: list[VMModel]):
+def ModelsPerformanceComparison(
+    dataset: VMDataset, models: list[VMModel]
+) -> Dict[str, List[Dict[str, float]]]:
     """
     Evaluates and compares the performance of multiple Machine Learning models using various metrics like accuracy,
     precision, recall, and F1 score.

validmind/tests/model_validation/sklearn/OverfitDiagnosis.py

@@ -2,11 +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 List
+from typing import Dict, List, Tuple
 
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
+import plotly.graph_objects as go
 import seaborn as sns
 from sklearn import metrics
 
@@ -178,7 +179,7 @@ def OverfitDiagnosis(
     datasets: List[VMDataset],
     metric: str = None,
     cut_off_threshold: float = DEFAULT_THRESHOLD,
-):
+) -> Tuple[Dict[str, List[Dict[str, float]]], go.Figure, RawData]:
     """
     Assesses potential overfitting in a model's predictions, identifying regions where performance between training and
     testing sets deviates significantly.

validmind/tests/model_validation/sklearn/PermutationFeatureImportance.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 typing import Union
+from typing import Optional, Tuple
 
 import plotly.graph_objects as go
 from sklearn.inspection import permutation_importance
@@ -26,9 +26,9 @@ logger = get_logger(__name__)
 def PermutationFeatureImportance(
     model: VMModel,
     dataset: VMDataset,
-    fontsize: Union[int, None] = None,
-    figure_height: Union[int, None] = None,
-):
+    fontsize: Optional[int] = None,
+    figure_height: Optional[int] = None,
+) -> Tuple[go.Figure, RawData]:
     """
     Assesses the significance of each feature in a model by evaluating the impact on model performance when feature
     values are randomly rearranged.

validmind/tests/model_validation/sklearn/PopulationStabilityIndex.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 typing import List
+from typing import Dict, List, Tuple
 
 import numpy as np
 import pandas as pd
@@ -82,7 +82,7 @@ def calculate_psi(score_initial, score_new, num_bins=10, mode="fixed"):
 @tasks("classification", "text_classification")
 def PopulationStabilityIndex(
     datasets: List[VMDataset], model: VMModel, num_bins: int = 10, mode: str = "fixed"
-):
+) -> Tuple[Dict[str, List[Dict[str, float]]], go.Figure, RawData]:
     """
     Assesses the Population Stability Index (PSI) to quantify the stability of an ML model's predictions across
     different datasets.

validmind/tests/model_validation/sklearn/PrecisionRecallCurve.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 numpy as np
 import plotly.graph_objects as go
 from sklearn.metrics import precision_recall_curve
@@ -14,7 +16,9 @@ from validmind.vm_models import VMDataset, VMModel
 
 @tags("sklearn", "binary_classification", "model_performance", "visualization")
 @tasks("classification", "text_classification")
-def PrecisionRecallCurve(model: VMModel, dataset: VMDataset):
+def PrecisionRecallCurve(
+    model: VMModel, dataset: VMDataset
+) -> Tuple[go.Figure, RawData]:
     """
     Evaluates the precision-recall trade-off for binary classification models and visualizes the Precision-Recall curve.
 

validmind/tests/model_validation/sklearn/ROCCurve.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 numpy as np
 import plotly.graph_objects as go
 from sklearn.metrics import roc_auc_score, roc_curve
@@ -19,7 +21,7 @@ from validmind.vm_models import VMDataset, VMModel
     "visualization",
 )
 @tasks("classification", "text_classification")
-def ROCCurve(model: VMModel, dataset: VMDataset):
+def ROCCurve(model: VMModel, dataset: VMDataset) -> Tuple[go.Figure, RawData]:
     """
     Evaluates binary classification model performance by generating and plotting the Receiver Operating Characteristic
     (ROC) curve and calculating the Area Under Curve (AUC) score.

validmind/tests/model_validation/sklearn/RegressionErrors.py

@@ -2,16 +2,21 @@
 # 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
 from sklearn import metrics
 
 from validmind import RawData, tags, tasks
+from validmind.vm_models import VMDataset, VMModel
 
 
 @tags("sklearn", "model_performance")
 @tasks("regression", "classification")
-def RegressionErrors(model, dataset):
+def RegressionErrors(
+    model: VMModel, dataset: VMDataset
+) -> Tuple[pd.DataFrame, RawData]:
     """
     Assesses the performance and error distribution of a regression model using various error metrics.
 

validmind/tests/model_validation/sklearn/RegressionErrorsComparison.py

@@ -2,19 +2,24 @@
 # 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
+
 import numpy as np
 import pandas as pd
 from sklearn import metrics
 
 from validmind import tags, tasks
 from validmind.logging import get_logger
+from validmind.vm_models import VMDataset, VMModel
 
 logger = get_logger(__name__)
 
 
 @tags("model_performance", "sklearn")
 @tasks("regression", "time_series_forecasting")
-def RegressionErrorsComparison(datasets, models):
+def RegressionErrorsComparison(
+    datasets: List[VMDataset], models: List[VMModel]
+) -> pd.DataFrame:
     """
     Assesses multiple regression error metrics to compare model performance across different datasets, emphasizing
     systematic overestimation or underestimation and large percentage errors.

validmind/tests/model_validation/sklearn/RegressionPerformance.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, List, Tuple
+
 import numpy as np
 from sklearn.metrics import mean_absolute_error, mean_squared_error
 
@@ -14,7 +16,9 @@ logger = get_logger(__name__)
 
 @tags("sklearn", "model_performance")
 @tasks("regression")
-def RegressionPerformance(model: VMModel, dataset: VMDataset):
+def RegressionPerformance(
+    model: VMModel, dataset: VMDataset
+) -> Tuple[List[Dict[str, float]], RawData]:
     """
     Evaluates the performance of a regression model using five different metrics: MAE, MSE, RMSE, MAPE, and MBD.
 

validmind/tests/model_validation/sklearn/RegressionR2Square.py

@@ -2,16 +2,21 @@
 # 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 sklearn import metrics
 
 from validmind import RawData, tags, tasks
 from validmind.tests.model_validation.statsmodels.statsutils import adj_r2_score
+from validmind.vm_models import VMDataset, VMModel
 
 
 @tags("sklearn", "model_performance")
 @tasks("regression")
-def RegressionR2Square(dataset, model):
+def RegressionR2Square(
+    dataset: VMDataset, model: VMModel
+) -> Tuple[pd.DataFrame, RawData]:
     """
     Assesses the overall goodness-of-fit of a regression model by evaluating R-squared (R2) and Adjusted R-squared (Adj
     R2) scores to determine the model's explanatory power over the dependent variable.

validmind/tests/model_validation/sklearn/RegressionR2SquareComparison.py

@@ -2,16 +2,21 @@
 # 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
+
 import pandas as pd
 from sklearn import metrics
 
 from validmind import tags, tasks
 from validmind.tests.model_validation.statsmodels.statsutils import adj_r2_score
+from validmind.vm_models import VMDataset, VMModel
 
 
 @tags("model_performance", "sklearn")
 @tasks("regression", "time_series_forecasting")
-def RegressionR2SquareComparison(datasets, models):
+def RegressionR2SquareComparison(
+    datasets: List[VMDataset], models: List[VMModel]
+) -> pd.DataFrame:
     """
     Compares R-Squared and Adjusted R-Squared values for different regression models across multiple datasets to assess
     model performance and relevance of features.

validmind/tests/model_validation/sklearn/RobustnessDiagnosis.py

@@ -12,9 +12,9 @@ import plotly.graph_objects as go
 import seaborn as sns
 from sklearn import metrics
 
+from validmind import tags, tasks
 from validmind.errors import MissingOrInvalidModelPredictFnError
 from validmind.logging import get_logger
-from validmind.tests import tags, tasks
 from validmind.vm_models import VMDataset, VMModel
 
 logger = get_logger(__name__)
@@ -221,7 +221,7 @@ def RobustnessDiagnosis(
     metric: str = None,
     scaling_factor_std_dev_list: List[float] = DEFAULT_STD_DEV_LIST,
     performance_decay_threshold: float = DEFAULT_DECAY_THRESHOLD,
-):
+) -> Tuple[pd.DataFrame, go.Figure, bool]:
     """
     Assesses the robustness of a machine learning model by evaluating performance decay under noisy conditions.