validmind 2.1.1__py3-none-any.whl → 2.2.4__py3-none-any.whl

validmind/tests/model_validation/TokenDisparity.py

@@ -2,139 +2,102 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
-import itertools
-from dataclasses import dataclass
-
 import pandas as pd
 import plotly.graph_objects as go
-from plotly.subplots import make_subplots
-from transformers import BertTokenizer
 
-from validmind.vm_models import Figure, Metric
+from validmind import tags, tasks
 
 
-@dataclass
-class TokenDisparity(Metric):
+@tags("nlp", "text_data", "visualization")
+@tasks("text_classification", "text_summarization")
+def TokenDisparity(dataset, model):
     """
-    Assess and visualize token count disparity between model's predicted and actual dataset.
-
-    **Purpose**:
-    The Token Disparity metric is designed to assess the distributional congruence between the model's predicted
-    outputs and the actual data. This is achieved by constructing histograms that illustrate the disparity in token
-    count between the two columns. Additionally, this metric is used to measure the model's verbosity in comparison to
-    the genuine dataset.
-
-    **Test Mechanism**:
-    The mechanism of running this test involves tokenizing both columns: one containing the actual data and the other
-    containing the model's predictions. The BERT tokenizer is used for tokenizing the contents of each column. After
-    tokenization, tokens in each column are counted and represented in two distinct histograms to facilitate the
-    visualization of token count distribution in the actual and predicted data. To quantify the difference in
-    distribution, the histogram of the actual tokens is compared with the histogram of the predicted tokens.
-
-    **Signs of High Risk**:
-    High risk or potential failure in model performance may be suggested by:
-
-    - Significant incongruities in distribution patterns between the two histograms.
-    - Marked divergence of the predicted histogram from the reference histogram, indicating that the model may be
-    generating output with unexpected verbosity.
-    - This might result in an output that has a significantly higher or lower number of tokens than expected.
-
-    **Strengths**:
-    Strengths of the Token Disparity metric include:
-
-    - It provides a clear and visual comparison of predicted versus actual token distributions, enhancing understanding
-    of the model's output consistency and verbosity.
-    - It is able to detect potential issues with the model's output generation capability, such as over-production or
-    under-production of tokens compared to the actual data set.
-
-    **Limitations**:
-    Limitations of the Token Disparity metric include:
-
-    - The metric focuses solely on token count, disregarding the semantics behind those tokens. Consequently, it may
-    miss out on issues related to relevance or meaningfulness of produced tokens.
-    - The assumption that similar token count between predicted and actual data suggests accurate output, which is not
-    always the case.
-    - Dependence on the BERT tokenizer, which may not always be the optimum choice for all types of text data.
+    Evaluates the token disparity between reference and generated texts, visualizing the results through histograms
+    and bar charts, alongside compiling a comprehensive table of descriptive statistics for token counts.
+
+    **Purpose:**
+    This function is designed to assess the token disparity between reference and generated texts. Token disparity is
+    important for understanding how closely the length and token usage of generated texts match the reference texts.
+
+    **Test Mechanism:**
+    The function starts by extracting the true and predicted values from the provided dataset and model. It then calculates
+    the number of tokens in each reference and generated text. Histograms and bar charts are generated for the token counts
+    of both reference and generated texts to visualize their distribution. Additionally, a table of descriptive statistics
+    (mean, median, standard deviation, minimum, and maximum) is compiled for the token counts, providing a comprehensive
+    summary of the model's performance.
+
+    **Signs of High Risk:**
+    - Significant disparity in token counts between reference and generated texts could indicate issues with text generation
+      quality, such as verbosity or lack of detail.
+    - Consistently low token counts in generated texts compared to references might suggest that the model is producing
+      incomplete or overly concise outputs.
+
+    **Strengths:**
+    - Provides a simple yet effective evaluation of text length and token usage.
+    - Visual representations (histograms and bar charts) make it easier to interpret the distribution and trends of token counts.
+    - Descriptive statistics offer a concise summary of the model's performance in generating texts of appropriate length.
+
+    **Limitations:**
+    - Token counts alone do not provide a complete assessment of text quality and should be supplemented with other metrics and qualitative analysis.
     """
 
-    name = "token_disparity"
-    required_inputs = ["model", "dataset"]
-
-    def run(self):
-        y_true = list(itertools.chain.from_iterable(self.inputs.dataset.y))
-        y_pred = self.inputs.dataset.y_pred(self.inputs.model)
-
-        df = pd.DataFrame({"reference_column": y_true, "generated_column": y_pred})
-
-        fig = self.token_disparity_histograms(df)
-        figures = []
-        figures.append(
-            Figure(
-                for_object=self,
-                key=self.key,
-                figure=fig,
-            )
-        )
-        return self.cache_results(figures=figures)
-
-    def token_disparity_histograms(self, df):
-        """
-        Visualize the token counts distribution of two given columns using histograms.
-
-        :param df: DataFrame containing the text columns.
-        :param params: Dictionary with the keys ["reference_column", "generated_column"].
-        """
+    # Extract true and predicted values
+    y_true = dataset.y
+    y_pred = dataset.y_pred(model)
 
-        reference_column = "reference_column"
-        generated_column = "generated_column"
+    # Calculate token counts
+    token_counts_true = [len(text.split()) for text in y_true]
+    token_counts_pred = [len(text.split()) for text in y_pred]
 
-        # Initialize the tokenizer
-        tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+    # Create a dataframe for reference and generated token counts
+    df = pd.DataFrame(
+        {"reference_tokens": token_counts_true, "generated_tokens": token_counts_pred}
+    )
 
-        # Tokenize the columns and get the number of tokens
-        df["tokens_1"] = df[reference_column].apply(
-            lambda x: len(tokenizer.tokenize(x))
-        )
-        df["tokens_2"] = df[generated_column].apply(
-            lambda x: len(tokenizer.tokenize(x))
-        )
+    figures = []
 
-        # Create subplots: 1 row, 2 columns
-        fig = make_subplots(
-            rows=1,
-            cols=2,
-            subplot_titles=(
-                f"Tokens in {reference_column}",
-                f"Tokens in {generated_column}",
-            ),
-        )
+    # Generate histograms and bar charts for reference and generated token counts
+    token_types = ["reference_tokens", "generated_tokens"]
+    token_names = ["Reference Tokens", "Generated Tokens"]
 
-        # Add histograms
-        fig.add_trace(
-            go.Histogram(
-                x=df["tokens_1"],
-                marker_color="blue",
-                name=f"Tokens in {reference_column}",
-            ),
-            row=1,
-            col=1,
+    for token_type, token_name in zip(token_types, token_names):
+        # Histogram
+        hist_fig = go.Figure(data=[go.Histogram(x=df[token_type])])
+        hist_fig.update_layout(
+            title=f"{token_name} Histogram",
+            xaxis_title=token_name,
+            yaxis_title="Count",
         )
-
-        fig.add_trace(
-            go.Histogram(
-                x=df["tokens_2"],
-                marker_color="red",
-                name=f"Tokens in {generated_column}",
-            ),
-            row=1,
-            col=2,
+        figures.append(hist_fig)
+
+        # Bar Chart
+        bar_fig = go.Figure(data=[go.Bar(x=df.index, y=df[token_type])])
+        bar_fig.update_layout(
+            title=f"{token_name} Bar Chart",
+            xaxis_title="Row Index",
+            yaxis_title=token_name,
         )
-
-        # Update layout
-        fig.update_layout(title_text="Token Distributions", bargap=0.1)
-
-        fig.update_yaxes(title_text="Number of Documents")
-        fig.update_xaxes(title_text="Number of Tokens", row=1, col=1)
-        fig.update_xaxes(title_text="Number of Tokens", row=1, col=2)
-
-        return fig
+        figures.append(bar_fig)
+
+    # Calculate statistics for each token count type
+    stats_df = df.describe().loc[["mean", "50%", "max", "min", "std"]]
+    stats_df = stats_df.rename(
+        index={
+            "mean": "Mean Count",
+            "50%": "Median Count",
+            "max": "Max Count",
+            "min": "Min Count",
+            "std": "Standard Deviation",
+        }
+    ).T
+    stats_df["Count"] = len(df)
+
+    # Rename columns for clarity
+    stats_df.index = stats_df.index.map(
+        {"reference_tokens": "Reference Tokens", "generated_tokens": "Generated Tokens"}
+    )
+
+    # Create a DataFrame from all collected statistics
+    result_df = pd.DataFrame(stats_df).reset_index().rename(columns={"index": "Metric"})
+
+    return (result_df, *tuple(figures))

validmind/tests/model_validation/ToxicityScore.py

@@ -2,146 +2,132 @@
 # See the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
-import itertools
-from dataclasses import dataclass
-
 import evaluate
 import pandas as pd
 import plotly.graph_objects as go
-import plotly.subplots as sp
 
-from validmind.vm_models import Figure, Metric
+from validmind import tags, tasks
 
 
-@dataclass
-class ToxicityScore(Metric):
+@tags("nlp", "text_data", "visualization")
+@tasks("text_classification", "text_summarization")
+def ToxicityScore(dataset, model):
     """
+    Computes and visualizes the toxicity score for input text, true text, and predicted text, assessing content quality and potential risk.
+
     **Purpose:**
-    The ToxicityScore metric is designed to present a sequential representation of toxicity scores for various texts.
-    Leveraging line plots, it gives an overview of how toxicity scores evolve across the sequence of texts, highlighting
-    trends and patterns.
+    The ToxicityScore metric is designed to evaluate the toxicity levels of texts generated by models. This is crucial for
+    identifying and mitigating harmful or offensive content in machine-generated texts.
 
     **Test Mechanism:**
-    The mechanism involves fetching texts from specific columns, computing their toxicity scores using a preloaded
-    `toxicity` evaluation tool, and then plotting these scores. A multi-panel visualization is created where each
-    panel is dedicated to a specific text data column. Line plots serve as the primary visual tool, showing the progression of toxicity scores across text sequences. Each
-    line plot corresponds to a specific text data column, illustrating the variation in toxicity scores as one moves
-    from one text segment to the next.
+    The function starts by extracting the input, true, and predicted values from the provided dataset and model. The toxicity score is
+    computed for each text using a preloaded `toxicity` evaluation tool. The scores are compiled into dataframes, and histograms
+    and bar charts are generated to visualize the distribution of toxicity scores. Additionally, a table of descriptive statistics
+    (mean, median, standard deviation, minimum, and maximum) is compiled for the toxicity scores, providing a comprehensive
+    summary of the model's performance.
 
     **Signs of High Risk:**
-    Drastic spikes in the line plots, especially those that reach high toxicity values, indicate potentially toxic
-    content within the associated text segment. If predicted summaries diverge significantly from input or target
-    texts, it could be indicative of issues in the model's generated content.
+    - Drastic spikes in toxicity scores indicate potentially toxic content within the associated text segment.
+    - Persistent high toxicity scores across multiple texts may suggest systemic issues in the model's text generation process.
 
     **Strengths:**
-    The ToxicityScore offers a dynamic view of toxicity trends, enabling users to detect patterns or irregularities
-    across the dataset. This is particularly valuable when comparing predicted content with actual data, helping
-    highlight any inconsistencies or abnormalities in model output.
+    - Provides a clear evaluation of toxicity levels in generated texts, helping to ensure content safety and appropriateness.
+    - Visual representations (histograms and bar charts) make it easier to interpret the distribution and trends of toxicity scores.
+    - Descriptive statistics offer a concise summary of the model's performance in generating non-toxic texts.
 
     **Limitations:**
-    This metric’s accuracy is contingent upon the underlying `toxicity` tool. The line plots provide a broad overview
-    of toxicity trends but do not specify which portions or tokens of the text are responsible for high toxicity scores.
-    Consequently, for granular insights, supplementary, in-depth analysis might be needed.
+    - The accuracy of the toxicity scores is contingent upon the underlying `toxicity` tool.
+    - The scores provide a broad overview but do not specify which portions or tokens of the text are responsible for high toxicity.
+    - Supplementary, in-depth analysis might be needed for granular insights.
     """
 
-    name = "toxicity_line_plot"
-    required_inputs = ["model", "dataset"]
-    metadata = {
-        "task_types": [
-            "text_classification",
-            "text_summarization",
-        ],
-        "tags": ["toxicity_line_plot"],
-    }
-
-    def _get_datasets(self):
-        # Check model attributes
-        if not hasattr(self, "model"):
-            raise AttributeError("The 'model' attribute is missing.")
-
-        y_true = list(itertools.chain.from_iterable(self.inputs.dataset.y))
-        y_pred = self.inputs.dataset.y_pred(self.inputs.model)
-        input_text = self.inputs.dataset.df[self.inputs.dataset.text_column]
-
-        # Ensure consistency in lengths
-        if not len(y_true) == len(y_pred) == len(input_text):
-            raise ValueError(
-                "Inconsistent lengths among input text, true summaries, and predicted summaries."
-            )
-
-        return input_text, y_true, y_pred
-
-    def toxicity_line_plots(self, df):
-        """
-        Compute toxicity scores for texts and then plot line plots for all columns of df.
-
-        Parameters:
-        - df (pd.DataFrame): The dataframe containing texts.
-        """
-
-        # Extract necessary parameters
-        toxicity = evaluate.load("toxicity")
-
-        # Get all columns of df
-        text_columns = df.columns.tolist()
-
-        # Determine the number of rows required based on the number of text columns
-        num_rows = (len(text_columns) + 1) // 2
-
-        # Create a subplot layout
-        fig = sp.make_subplots(rows=num_rows, cols=2, subplot_titles=text_columns)
-
-        subplot_height = 350
-        total_height = num_rows * subplot_height + 200
-
-        for idx, col in enumerate(text_columns, start=1):
-            row = (idx - 1) // 2 + 1
-            col_idx = (idx - 1) % 2 + 1
-
-            # Get list of texts from dataframe
-            texts = df[col].tolist()
-
-            # Compute toxicity for texts
-            toxicity_scores = toxicity.compute(predictions=texts)["toxicity"]
-
-            # Add traces to the corresponding subplot
-            fig.add_trace(
-                go.Scatter(
-                    y=toxicity_scores,
-                    mode="lines+markers",
-                    marker=dict(size=5),
-                    line=dict(width=1.5),
-                    showlegend=False,
-                ),
-                row=row,
-                col=col_idx,
-            )
-
-            # Update xaxes and yaxes titles only for the first subplot
-            if idx == 1:
-                fig.update_xaxes(title_text="Text Index", row=row, col=col_idx)
-                fig.update_yaxes(title_text="Toxicity Score", row=row, col=col_idx)
-
-        # Update layout
-        fig.update_layout(
-            title_text="Line Plots of Toxicity Scores", height=total_height
-        )
-
-        return fig
-
-    def run(self):
-        input_text, y_true, y_pred = self._get_datasets()
-
-        df = pd.DataFrame(
-            {
-                "Input Text": input_text,
-                "Target Text": y_true,
-                "Predicted Summaries": y_pred,
-            }
+    # Extract true, predicted, and input values
+    y_true = dataset.y
+    y_pred = dataset.y_pred(model)
+    input_text = dataset.df[dataset.text_column]
+
+    # Load the toxicity evaluation metric
+    toxicity = evaluate.load("toxicity")
+
+    # Function to calculate toxicity scores
+    def compute_toxicity_scores(texts):
+        scores = []
+        for text in texts:
+            score = toxicity.compute(predictions=[text])
+            scores.append(score["toxicity"])
+        return scores
+
+    # Calculate toxicity scores for input, true, and predicted texts
+    input_toxicity = compute_toxicity_scores(input_text)
+    true_toxicity = compute_toxicity_scores(y_true)
+    pred_toxicity = compute_toxicity_scores(y_pred)
+
+    # Convert scores to dataframes
+    input_df = pd.DataFrame(input_toxicity, columns=["Input Text Toxicity"])
+    true_df = pd.DataFrame(true_toxicity, columns=["True Text Toxicity"])
+    pred_df = pd.DataFrame(pred_toxicity, columns=["Predicted Text Toxicity"])
+
+    figures = []
+
+    # Function to create histogram and bar chart for toxicity scores
+    def create_figures(df, title):
+        # Histogram
+        hist_fig = go.Figure(data=[go.Histogram(x=df.iloc[:, 0])])
+        hist_fig.update_layout(
+            title=f"{title} Histogram",
+            xaxis_title=title,
+            yaxis_title="Count",
         )
-
-        fig = self.toxicity_line_plots(df)
-
-        return self.cache_results(
-            figures=[Figure(for_object=self, key=self.key, figure=fig)]
+        figures.append(hist_fig)
+
+        # Bar Chart
+        bar_fig = go.Figure(data=[go.Bar(x=df.index, y=df.iloc[:, 0])])
+        bar_fig.update_layout(
+            title=f"{title} Bar Chart",
+            xaxis_title="Text Instance Index",
+            yaxis_title=title,
         )
+        figures.append(bar_fig)
+
+    # Create figures for each toxicity score dataframe
+    create_figures(input_df, "Input Text Toxicity")
+    create_figures(true_df, "True Text Toxicity")
+    create_figures(pred_df, "Predicted Text Toxicity")
+
+    # Calculate statistics for each toxicity score dataframe
+    def calculate_stats(df):
+        stats = df.describe().loc[["mean", "50%", "max", "min", "std"]].T
+        stats.columns = [
+            "Mean Score",
+            "Median Score",
+            "Max Score",
+            "Min Score",
+            "Standard Deviation",
+        ]
+        stats["Metric"] = df.columns[0]
+        stats["Count"] = len(df)
+        return stats
+
+    input_stats = calculate_stats(input_df)
+    true_stats = calculate_stats(true_df)
+    pred_stats = calculate_stats(pred_df)
+
+    # Combine statistics into a single dataframe
+    result_df = (
+        pd.concat([input_stats, true_stats, pred_stats])
+        .reset_index()
+        .rename(columns={"index": "Statistic"})
+    )
+    result_df = result_df[
+        [
+            "Metric",
+            "Mean Score",
+            "Median Score",
+            "Max Score",
+            "Min Score",
+            "Standard Deviation",
+            "Count",
+        ]
+    ]
+
+    return (result_df, *tuple(figures))

validmind/tests/model_validation/embeddings/CosineSimilarityComparison.py

@@ -0,0 +1,96 @@
+# 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 itertools import combinations
+
+import numpy as np
+import pandas as pd
+import plotly.express as px
+from sklearn.metrics.pairwise import cosine_similarity
+
+
+def CosineSimilarityComparison(dataset, models):
+    """
+    Computes pairwise cosine similarities between model embeddings and visualizes the results through bar charts,
+    alongside compiling a comprehensive table of descriptive statistics for each model pair.
+
+    **Purpose:**
+    This function is designed to analyze and compare the embeddings produced by different models using Cosine Similarity.
+    Cosine Similarity, a measure calculating the cosine of the angle between two vectors, is widely used to determine
+    the alignment or similarity between vectors in high-dimensional spaces, such as text embeddings. This analysis helps
+    to understand how similar or different the models' predictions are in terms of embedding generation.
+
+    **Test Mechanism:**
+    The function begins by computing the embeddings for each model using the provided dataset. It then calculates the
+    cosine similarity for every possible pair of models, generating a similarity matrix. Each element of this matrix
+    represents the cosine similarity between two model embeddings. The function flattens this matrix and uses it to
+    create a bar chart for each model pair, visualizing their similarity distribution. Additionally, it compiles a table
+    with descriptive statistics (mean, median, standard deviation, minimum, and maximum) for the similarities of each
+    pair, including a reference to the compared models.
+
+    **Signs of High Risk:**
+
+    - A high concentration of cosine similarity values close to 1 could suggest that the models are producing very
+      similar embeddings, which could be a sign of redundancy or lack of diversity in model training or design.
+    - Conversely, very low similarity values near -1 indicate strong dissimilarity, potentially highlighting models
+      that are too divergent, possibly focusing on very different features of the data.
+
+    **Strengths:**
+
+    - Enables detailed comparisons between multiple models' embedding strategies through visual and statistical means.
+    - Helps identify which models produce similar or dissimilar embeddings, useful for tasks requiring model diversity.
+    - Provides quantitative and visual feedback on the degree of similarity, enhancing interpretability of model
+      behavior in embedding spaces.
+
+    **Limitations:**
+
+    - The analysis is confined to the comparison of embeddings and does not assess the overall performance of the models
+      in terms of their primary tasks (e.g., classification, regression).
+    - Assumes that the models are suitable for generating comparable embeddings, which might not always be the case,
+      especially across different types of models.
+    - Interpretation of results is heavily dependent on the understanding of Cosine Similarity and the nature of high-dimensional
+      embedding spaces.
+    """
+
+    figures = []
+    # Initialize a list to store data for the DataFrame
+    all_stats = []
+
+    # Generate all pairs of models for comparison
+    for model_A, model_B in combinations(models, 2):
+        embeddings_A = np.stack(dataset.y_pred(model_A))
+        embeddings_B = np.stack(dataset.y_pred(model_B))
+
+        # Calculate pairwise cosine similarity
+        similarity_matrix = cosine_similarity(embeddings_A, embeddings_B)
+        similarities = similarity_matrix.flatten()
+
+        # Generate statistics and add model combination as a column
+        stats_data = {
+            "Combination": f"{model_A.input_id} vs {model_B.input_id}",
+            "Mean": np.mean(similarities),
+            "Median": np.median(similarities),
+            "Standard Deviation": np.std(similarities),
+            "Minimum": np.min(similarities),
+            "Maximum": np.max(similarities),
+        }
+        all_stats.append(stats_data)
+
+        # Generate an index for each similarity value
+        indices = range(len(similarities))
+
+        # Create the bar chart using Plotly
+        fig = px.bar(
+            x=indices,
+            y=similarities,
+            labels={"x": "Pair Index", "y": "Cosine Similarity"},
+            title=f"Cosine Similarity - {model_A.input_id} vs {model_B.input_id}",
+        )
+        fig.update_layout(xaxis_title="Pair Index", yaxis_title="Cosine Similarity")
+        figures.append(fig)
+
+    # Create a DataFrame from all collected statistics
+    stats_df = pd.DataFrame(all_stats)
+
+    return (stats_df, *tuple(figures))

validmind/tests/model_validation/embeddings/CosineSimilarityHeatmap.py

@@ -0,0 +1,71 @@
+# 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
+
+import numpy as np
+import plotly.express as px
+from sklearn.metrics.pairwise import cosine_similarity
+
+
+def CosineSimilarityHeatmap(
+    dataset,
+    model,
+    title="Cosine Similarity Matrix",
+    color="Cosine Similarity",
+    xaxis_title="Index",
+    yaxis_title="Index",
+    color_scale="Blues",
+):
+    """
+    Generates an interactive heatmap to visualize the cosine similarities among embeddings derived from a given model.
+
+    **Purpose:**
+    This function is designed to visually analyze the cosine similarities of embeddings from a specific model.
+    Cosine similarity, a measure of the cosine of the angle between two vectors, aids in understanding the
+    orientation and similarity of vectors in multi-dimensional space. This is particularly valuable for exploring
+    text embeddings and their relative similarities among documents, words, or phrases.
+
+    **Test Mechanism:**
+    The function operates through a sequence of steps to visualize cosine similarities. Initially,
+    embeddings are extracted for each dataset entry using the designated model. Following this,
+    the function computes the pairwise cosine similarities among these embeddings. The computed similarities
+    are then displayed in an interactive heatmap.
+
+    **Signs of High Risk:**
+    - High similarity values (close to 1) across the heatmap might not always be indicative of a risk;
+    however, in contexts where diverse perspectives or features are desired, this could suggest a lack of
+    diversity in the model's learning process or potential redundancy.
+    - Similarly, low similarity values (close to -1) indicate strong dissimilarity, which could be beneficial in
+    scenarios demanding diverse outputs. However, in cases where consistency is needed, these low values might
+    highlight that the model is unable to capture a coherent set of features from the data, potentially leading to poor performance on related tasks.
+
+    **Strengths:**
+    - Provides an interactive and intuitive visual representation of embedding similarities, facilitating easy exploration and analysis.
+    - Allows customization of visual elements such as title, axis labels, and color scale to suit specific analytical needs and preferences.
+
+    **Limitations:**
+    - As the number of embeddings increases, the effectiveness of the heatmap might diminish due to overcrowding, making it hard to discern detailed similarities.
+    - The interpretation of the heatmap heavily relies on the appropriate setting of the color scale, as incorrect settings can lead to misleading visual interpretations.
+    """
+
+    embeddings = np.stack(dataset.y_pred(model))
+
+    # Calculate pairwise cosine similarity
+    similarity_matrix = cosine_similarity(embeddings)
+
+    # Create the heatmap using Plotly
+    fig = px.imshow(
+        similarity_matrix,
+        labels=dict(x=xaxis_title, y=yaxis_title, color=color),
+        text_auto=True,
+        aspect="auto",
+        color_continuous_scale=color_scale,
+    )
+
+    fig.update_layout(
+        title=f"{title} - {model.input_id}",
+        xaxis_title=xaxis_title,
+        yaxis_title=yaxis_title,
+    )
+
+    return fig