add data split tab + refactor

app.py

@@ -7,6 +7,7 @@ import plotly.express as px
 from utils import (
     KEY_TO_CATEGORY_NAME,
     PROPRIETARY_LICENSES,
+    CAT_NAME_TO_EXPLANATION,
     download_latest_data_from_space,
 )
 
@@ -55,30 +56,66 @@ for k, v in merged_dfs.items():
         merged_dfs[k], release_date_mapping[["key", "Release Date"]], on="key"
     )
 
-df = merged_dfs["Overall"]
-df["License"] = df["License"].apply(
-    lambda x: "Proprietary LLM" if x in PROPRIETARY_LICENSES else "Open LLM"
-)
-df["Release Date"] = pd.to_datetime(df["Release Date"])
-df["Month-Year"] = df["Release Date"].dt.to_period("M")
-df["rating"] = df["rating"].round()
+
+# format dataframes
+def format_data(df):
+    df["License"] = df["License"].apply(
+        lambda x: "Proprietary LLM" if x in PROPRIETARY_LICENSES else "Open LLM"
+    )
+    df["Release Date"] = pd.to_datetime(df["Release Date"])
+    df["Month-Year"] = df["Release Date"].dt.to_period("M")
+    df["rating"] = df["rating"].round()
+    return df.reset_index(drop=True)
+
+
+merged_dfs = {k: format_data(v) for k, v in merged_dfs.items()}
+
+
+# get constants
+filter_ranges = {}
+for k, df in merged_dfs.items():
+    filter_ranges[k] = {
+        "min_elo_score": df["rating"].min().round(),
+        "max_elo_score": df["rating"].max().round(),
+        "upper_models_per_month": int(
+            df.groupby(["Month-Year", "License"])["rating"]
+            .apply(lambda x: x.count())
+            .max()
+        ),
+    }
+
+min_elo_score = float("inf")
+max_elo_score = float("-inf")
+upper_models_per_month = 0
+
+for key, value in filter_ranges.items():
+    min_elo_score = min(min_elo_score, value["min_elo_score"])
+    max_elo_score = max(max_elo_score, value["max_elo_score"])
+    upper_models_per_month = max(
+        upper_models_per_month, value["upper_models_per_month"]
+    )
+
+
+date_updated = elo_results["full"]["last_updated_datetime"].split(" ")[0]
+
+
+def get_data_split(dfs, set_name):
+    df = dfs[set_name].copy(deep=True)
+    return df.reset_index(drop=True)
 
 
 ###################
 ### Plot Data
 ###################
 
-date_updated = elo_results["full"]["last_updated_datetime"].split(" ")[0]
-min_elo_score = df["rating"].min().round()
-max_elo_score = df["rating"].max().round()
-upper_models_per_month = int(
-    df.groupby(["Month-Year", "License"])["rating"].apply(lambda x: x.count()).max()
-)
 
+def build_plot(min_score, max_models_per_month, toggle_annotations, set_selector):
 
-def build_plot(min_score, max_models_per_month, toggle_annotations):
+    df = get_data_split(merged_dfs, set_name=set_selector)
 
+    # filter data
     filtered_df = df[(df["rating"] >= min_score)]
+
     filtered_df = (
         filtered_df.groupby(["Month-Year", "License"])
         .apply(lambda x: x.nlargest(max_models_per_month, "rating"))
@@ -91,11 +128,11 @@ def build_plot(min_score, max_models_per_month, toggle_annotations):
         y="rating",
         color="License",
         hover_name="Model",
-        hover_data=["Organization", "License"],
+        hover_data=["Organization", "License", "Link"],
         trendline="ols",
         title=f"Proprietary vs Open LLMs (LMSYS Arena ELO as of {date_updated})",
         labels={"rating": "Arena ELO", "Release Date": "Release Date"},
-        height=700,
+        height=800,
         template="seaborn",
     )
 
@@ -143,45 +180,58 @@ with gr.Blocks(
         </div>
         """
     )
+
     with gr.Row():
-        min_score = gr.Slider(
-            minimum=min_elo_score,
-            maximum=max_elo_score,
-            value=800,
-            step=50,
-            label="Minimum ELO Score",
-        )
-        max_models_per_month = gr.Slider(
-            value=upper_models_per_month,
-            minimum=1,
-            maximum=upper_models_per_month,
-            step=1,
-            label="Max Models per Month (per License)",
-        )
-        toggle_annotations = gr.Radio(
-            choices=[True, False], label="Overlay Best Model Name", value=False
-        )
+        with gr.Column():
+            toggle_annotations = gr.Radio(
+                choices=[True, False], label="Overlay Best Model Name", value=True
+            )
+            set_selector = gr.Dropdown(
+                choices=list(CAT_NAME_TO_EXPLANATION.keys()),
+                label="Select Dataset",
+                value="Overall",
+            )
+        with gr.Column():
+            min_score = gr.Slider(
+                minimum=min_elo_score,
+                maximum=max_elo_score,
+                value=(max_elo_score - min_elo_score) * 0.3 + min_elo_score,
+                step=50,
+                label="Minimum ELO Score",
+            )
+            max_models_per_month = gr.Slider(
+                value=upper_models_per_month - 2,
+                minimum=1,
+                maximum=upper_models_per_month,
+                step=1,
+                label="Max Models per Month (per License)",
+            )
 
     # Show plot
     plot = gr.Plot()
     demo.load(
         fn=build_plot,
-        inputs=[min_score, max_models_per_month, toggle_annotations],
+        inputs=[min_score, max_models_per_month, toggle_annotations, set_selector],
         outputs=plot,
     )
     min_score.change(
         fn=build_plot,
-        inputs=[min_score, max_models_per_month, toggle_annotations],
+        inputs=[min_score, max_models_per_month, toggle_annotations, set_selector],
         outputs=plot,
     )
     max_models_per_month.change(
         fn=build_plot,
-        inputs=[min_score, max_models_per_month, toggle_annotations],
+        inputs=[min_score, max_models_per_month, toggle_annotations, set_selector],
         outputs=plot,
     )
     toggle_annotations.change(
         fn=build_plot,
-        inputs=[min_score, max_models_per_month, toggle_annotations],
+        inputs=[min_score, max_models_per_month, toggle_annotations, set_selector],
+        outputs=plot,
+    )
+    set_selector.change(
+        fn=build_plot,
+        inputs=[min_score, max_models_per_month, toggle_annotations, set_selector],
         outputs=plot,
     )