Changed to Plotly for interactive graphs!

suggestion by Nathan Habib implemented. So much better than the dumb printed graphs!

src/display_models/plot_results.py

@@ -1,8 +1,9 @@
 import pandas as pd
-import altair as alt
+import plotly.express as px
+from plotly.graph_objs import Figure
 import pickle
 from datetime import datetime, timezone
-from typing import List, Dict, Tuple, Any, Union
+from typing import List, Dict, Tuple, Any
 
 # Average ⬆️ human baseline is 0.897 (source: averaging human baselines below)
 # ARC human baseline is 0.80 (source: https://lab42.global/arc/)
@@ -82,6 +83,7 @@ def create_scores_df(results_df: pd.DataFrame) -> pd.DataFrame:
         "MMLU": [],
         "TruthfulQA": [],
         "Result Date": [],
+        "Model Name": [],
     }
 
     # Step 3: Iterate over the rows of the DataFrame and update the scores dictionary
@@ -92,6 +94,9 @@ def create_scores_df(results_df: pd.DataFrame) -> pd.DataFrame:
                 if not scores[column] or scores[column][-1] <= date:
                     scores[column].append(date)
                 continue
+            if column == "Model Name":
+                scores[column].append(row["model_name_for_query"])
+                continue
             current_max = scores[column][-1] if scores[column] else float("-inf")
             scores[column].append(max(current_max, row[column]))
 
@@ -114,7 +119,7 @@ def create_plot_df(scores_df: pd.DataFrame) -> pd.DataFrame:
 
     # Iterate over the cols and create a new DataFrame for each column
     for col in cols:
-        d = scores_df[[col, "Result Date"]].copy().reset_index(drop=True)
+        d = scores_df[[col, "Model Name", "Result Date"]].copy().reset_index(drop=True)
         d["Metric Name"] = col
         d.rename(columns={col: "Metric Value"}, inplace=True)
         dfs.append(d)
@@ -133,37 +138,77 @@ def create_plot_df(scores_df: pd.DataFrame) -> pd.DataFrame:
     return concat_df
 
 
-def create_metric_plot_obj(df: pd.DataFrame, metrics: List[str], human_baselines: Dict[str, float]) -> alt.LayerChart:
+def create_metric_plot_obj(
+    df: pd.DataFrame, metrics: List[str], human_baselines: Dict[str, float], title: str
+) -> Figure:
     """
-    Creates a visualization of metrics over time compared to human baselines.
-
-    :param df: A DataFrame containing 'Metric Name', 'Metric Value', and 'Result Date' columns.
-    :param metrics: A list of metric names to be included in the plot.
-    :param human_baselines: A dictionary mapping metric names to their corresponding human baseline values.
-    :return: An Altair LayerChart object visualizing the metrics over time.
+    Create a Plotly figure object with lines representing different metrics
+    and horizontal dotted lines representing human baselines.
+
+    :param df: The DataFrame containing the metric values, names, and dates.
+    :param metrics: A list of strings representing the names of the metrics
+                    to be included in the plot.
+    :param human_baselines: A dictionary where keys are metric names
+                            and values are human baseline values for the metrics.
+    :param title: A string representing the title of the plot.
+    :return: A Plotly figure object with lines representing metrics and
+             horizontal dotted lines representing human baselines.
     """
-    # Filter the DataFrame based on the metrics parameter
+
+    # Filter the DataFrame based on the specified metrics
     df = df[df["Metric Name"].isin(metrics)]
 
-    # Filter the human_baselines dictionary to include only the specified metrics
+    # Filter the human baselines based on the specified metrics
     filtered_human_baselines = {k: v for k, v in human_baselines.items() if k in metrics}
 
-    # Create a DataFrame from filtered human baselines
-    human_baselines_df = pd.DataFrame(list(filtered_human_baselines.items()), columns=["Metric Name", "Metric Value"])
-
-    # Create the lines chart for each metric over time.
-    base = alt.Chart(df).encode(x="Result Date:T")
-    lines = base.mark_line().encode(
-        alt.Y("Metric Value:Q", scale=alt.Scale(domain=[0, 100])),
-        color="Metric Name:N",
+    # Create a line figure using plotly express with specified markers and custom data
+    fig = px.line(
+        df,
+        x="Result Date",
+        y="Metric Value",
+        color="Metric Name",
+        markers=True,
+        custom_data=["Metric Name", "Metric Value", "Model Name"],
+        title=title,
     )
 
-    # Create the rules (horizontal lines) chart for the human baselines.
-    yrules = (
-        alt.Chart(human_baselines_df)
-        .mark_rule(strokeDash=[12, 6], size=2)
-        .encode(y="Metric Value:Q", color="Metric Name:N")
+    # Update hovertemplate for better hover interaction experience
+    fig.update_traces(
+        hovertemplate="<br>".join(
+            [
+                "Model Name: %{customdata[2]}",
+                "Metric Name: %{customdata[0]}",
+                "Date: %{x}",
+                "Metric Value: %{y}",
+            ]
+        )
     )
 
-    # Combine lines with yrules and return the chart.
-    return lines + yrules
+    # Create a dictionary to hold the color mapping for each metric
+    metric_color_mapping = {}
+
+    # Map each metric name to its color in the figure
+    for trace in fig.data:
+        metric_color_mapping[trace.name] = trace.line.color
+
+    # Iterate over filtered human baselines and add horizontal lines to the figure
+    for metric, value in filtered_human_baselines.items():
+        color = metric_color_mapping.get(metric, "blue")  # Retrieve color from mapping; default to blue if not found
+        location = "top left" if metric == "HellaSwag" else "bottom left"  # Set annotation position
+        # Add horizontal line with matched color and positioned annotation
+        fig.add_hline(
+            y=value,
+            line_dash="dot",
+            annotation_text=f"{metric} human baseline",
+            annotation_position=location,
+            annotation_font_size=10,
+            annotation_font_color=color,
+            line_color=color,
+        )
+
+    return fig
+
+
+# Example Usage:
+# human_baselines dictionary is defined.
+# chart = create_metric_plot_obj(scores_df, ["ARC", "HellaSwag", "MMLU", "TruthfulQA"], human_baselines, "Graph Title")