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

validmind/tests/model_validation/ragas/ContextRecall.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,13 +33,13 @@ except ImportError as e:
 @tags("ragas", "llm", "retrieval_performance")
 @tasks("text_qa", "text_generation", "text_summarization", "text_classification")
 def ContextRecall(
-    dataset,
+    dataset: VMDataset,
     user_input_column: str = "user_input",
     retrieved_contexts_column: str = "retrieved_contexts",
     reference_column: str = "reference",
     judge_llm=None,
     judge_embeddings=None,
-):
+) -> Tuple[Dict[str, list], go.Figure, go.Figure, RawData]:
     """
     Context recall measures the extent to which the retrieved context aligns with the
     annotated answer, treated as the ground truth. It is computed based on the `ground
@@ -109,6 +112,7 @@ def ContextRecall(
     }
 
     df = get_renamed_columns(dataset._df, required_columns)
+    df = df[required_columns.keys()]
 
     result_df = evaluate(
         Dataset.from_pandas(df),

validmind/tests/model_validation/ragas/Faithfulness.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,13 +33,13 @@ except ImportError as e:
 @tags("ragas", "llm", "rag_performance")
 @tasks("text_qa", "text_generation", "text_summarization")
 def Faithfulness(
-    dataset,
-    user_input_column="user_input",
-    response_column="response",
-    retrieved_contexts_column="retrieved_contexts",
+    dataset: VMDataset,
+    user_input_column: str = "user_input",
+    response_column: str = "response",
+    retrieved_contexts_column: str = "retrieved_contexts",
     judge_llm=None,
     judge_embeddings=None,
-):  # noqa
+) -> Tuple[Dict[str, list], go.Figure, go.Figure, RawData]:
     """
     Evaluates the faithfulness of the generated answers with respect to retrieved contexts.
 
@@ -115,6 +118,7 @@ def Faithfulness(
 
     df = get_renamed_columns(dataset._df, required_columns)
 
+    df = df[required_columns.keys()]
     result_df = evaluate(
         Dataset.from_pandas(df),
         metrics=[faithfulness()],

validmind/tests/model_validation/ragas/NoiseSensitivity.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
 
@@ -32,15 +35,15 @@ VALID_FOCUS_VALUES = ["relevant", "irrelevant"]
 @tags("ragas", "llm", "rag_performance")
 @tasks("text_qa", "text_generation", "text_summarization")
 def NoiseSensitivity(
-    dataset,
-    response_column="response",
-    retrieved_contexts_column="retrieved_contexts",
-    reference_column="reference",
-    focus="relevant",
-    user_input_column="user_input",
+    dataset: VMDataset,
+    response_column: str = "response",
+    retrieved_contexts_column: str = "retrieved_contexts",
+    reference_column: str = "reference",
+    focus: str = "relevant",
+    user_input_column: str = "user_input",
     judge_llm=None,
     judge_embeddings=None,
-):
+) -> Tuple[Dict[str, list], go.Figure, go.Figure, RawData]:
     """
     Assesses the sensitivity of a Large Language Model (LLM) to noise in retrieved context by measuring how often it
     generates incorrect responses.

validmind/tests/model_validation/ragas/ResponseRelevancy.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,13 +33,13 @@ except ImportError as e:
 @tags("ragas", "llm", "rag_performance")
 @tasks("text_qa", "text_generation", "text_summarization")
 def ResponseRelevancy(
-    dataset,
-    user_input_column="user_input",
-    retrieved_contexts_column=None,
-    response_column="response",
+    dataset: VMDataset,
+    user_input_column: str = "user_input",
+    retrieved_contexts_column: str = None,
+    response_column: str = "response",
     judge_llm=None,
     judge_embeddings=None,
-):
+) -> Tuple[Dict[str, list], go.Figure, go.Figure, RawData]:
     """
     Assesses how pertinent the generated answer is to the given prompt.
 
@@ -124,6 +127,7 @@ def ResponseRelevancy(
         required_columns["retrieved_contexts"] = retrieved_contexts_column
 
     df = get_renamed_columns(dataset._df, required_columns)
+    df = df[required_columns.keys()]
 
     metrics = [response_relevancy()]
 
@@ -134,7 +138,6 @@ def ResponseRelevancy(
     ).to_pandas()
 
     score_column = "answer_relevancy"
-
     fig_histogram = px.histogram(
         x=result_df[score_column].to_list(), nbins=10, title="Response Relevancy"
     )

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.