Update app.py

app.py

@@ -109,32 +109,36 @@ def load_leaderboard_data(csv_file_path):
         return pd.DataFrame(columns=UGI_COLS + WRITING_STYLE_COLS + ANIME_RATING_COLS + ADDITIONAL_COLS)
 
 # Update the leaderboard table based on the search query and parameter range filters
-def update_table(df: pd.DataFrame, query: str, param_ranges: list, columns: list, w10_range: tuple, additional_cols: list) -> pd.DataFrame:
+def update_table(df: pd.DataFrame, query: str, param_ranges: list, is_foundation: bool, columns: list, w10_range: tuple, additional_cols: list) -> pd.DataFrame:
     filtered_df = df.copy()
+    
+    # Apply model size filter
     if param_ranges:
-        param_mask = pd.Series(False, index=filtered_df.index)
+        size_mask = pd.Series(False, index=filtered_df.index)
         for param_range in param_ranges:
             if param_range == '~2':
-                param_mask |= (filtered_df['Total Params'] < 2.5)
+                size_mask |= (filtered_df['Total Params'] < 2.5)
             elif param_range == '~4':
-                param_mask |= ((filtered_df['Total Params'] >= 2.5) & (filtered_df['Total Params'] < 6))
+                size_mask |= ((filtered_df['Total Params'] >= 2.5) & (filtered_df['Total Params'] < 6))
             elif param_range == '~8':
-                param_mask |= ((filtered_df['Total Params'] >= 6) & (filtered_df['Total Params'] < 9.5))
+                size_mask |= ((filtered_df['Total Params'] >= 6) & (filtered_df['Total Params'] < 9.5))
             elif param_range == '~13':
-                param_mask |= ((filtered_df['Total Params'] >= 9.5) & (filtered_df['Total Params'] < 16))
+                size_mask |= ((filtered_df['Total Params'] >= 9.5) & (filtered_df['Total Params'] < 16))
             elif param_range == '~20':
-                param_mask |= ((filtered_df['Total Params'] >= 16) & (filtered_df['Total Params'] < 28))
+                size_mask |= ((filtered_df['Total Params'] >= 16) & (filtered_df['Total Params'] < 28))
             elif param_range == '~34':
-                param_mask |= ((filtered_df['Total Params'] >= 28) & (filtered_df['Total Params'] < 40))
+                size_mask |= ((filtered_df['Total Params'] >= 28) & (filtered_df['Total Params'] < 40))
             elif param_range == '~50':
-                param_mask |= ((filtered_df['Total Params'] >= 40) & (filtered_df['Total Params'] < 65))
+                size_mask |= ((filtered_df['Total Params'] >= 40) & (filtered_df['Total Params'] < 65))
             elif param_range == '~70+':
-                param_mask |= (filtered_df['Total Params'] >= 65)
+                size_mask |= (filtered_df['Total Params'] >= 65)
             elif param_range == 'Closed':
-                param_mask |= filtered_df['Total Params'].isna()
-            elif param_range == 'Foundation':
-                param_mask |= (filtered_df['Foundation'] == 1)
-        filtered_df = filtered_df[param_mask]
+                size_mask |= filtered_df['Total Params'].isna()
+        filtered_df = filtered_df[size_mask]
+    
+    # Apply foundation model filter
+    if is_foundation:
+        filtered_df = filtered_df[filtered_df['Foundation'] == 1]
     
     if query:
         filtered_df = filtered_df[filtered_df['Model'].str.contains(query, case=False, na=False)]
@@ -179,15 +183,22 @@ with GraInter:
         with gr.Row():
             search_bar = gr.Textbox(placeholder=" 🔍 Search for a model...", show_label=False, elem_id="search-bar")
         with gr.Row():
-            with gr.Column(scale=5):
+            with gr.Column(scale=7):
                 filter_columns_size = gr.CheckboxGroup(
                     label="Model sizes (in billions of parameters)",
-                    choices=['~2', '~4', '~8', '~13', '~20', '~34', '~50', '~70+', 'Closed', 'Foundation'],
+                    choices=['~2', '~4', '~8', '~13', '~20', '~34', '~50', '~70+', 'Closed'],
                     value=[],
                     interactive=True,
                     elem_id="filter-columns-size",
                 )
-            with gr.Column(scale=2):
+            with gr.Column(min_width=200, scale=0):
+                model_type = gr.Checkbox(
+                    label="Foundation Models Only",
+                    value=False,
+                    interactive=True,
+                    elem_id="model-type",
+                )
+            with gr.Column(scale=3):
                 w10_range = RangeSlider(minimum=0, maximum=10, value=(0, 10), step=0.1, label="W/10 Range")
         with gr.Row():
             additional_columns = gr.CheckboxGroup(
@@ -325,42 +336,56 @@ with GraInter:
             **NA:** When models either reply with one number for every anime, give ratings not between 1 and 10, or don't give every anime in the list a rating.
             """)
 
-    def update_all_tables(query, param_ranges, w10_range, additional_cols):
-        ugi_table = update_table(leaderboard_df, query, param_ranges, UGI_COLS, w10_range, additional_cols)
+    def update_all_tables(query, param_ranges, is_foundation, w10_range, additional_cols):
+        ugi_table = update_table(leaderboard_df, query, param_ranges, is_foundation, UGI_COLS, w10_range, additional_cols)
         
         ws_df = leaderboard_df.sort_values(by='Reg+MyScore 🏆', ascending=False)
-        ws_table = update_table(ws_df, query, param_ranges, WRITING_STYLE_COLS, w10_range, additional_cols)
+        ws_table = update_table(ws_df, query, param_ranges, is_foundation, WRITING_STYLE_COLS, w10_range, additional_cols)
         
         arp_df = leaderboard_df.sort_values(by='Score 🏆', ascending=False)
         arp_df_na = arp_df[arp_df[['Dif', 'Cor']].isna().any(axis=1)]
         arp_df = arp_df[~arp_df[['Dif', 'Cor']].isna().any(axis=1)]
         
-        arp_table = update_table(arp_df, query, param_ranges, ANIME_RATING_COLS, w10_range, additional_cols)
-        arp_na_table = update_table(arp_df_na, query, param_ranges, ANIME_RATING_COLS, w10_range, additional_cols).fillna('NA')
+        arp_table = update_table(arp_df, query, param_ranges, is_foundation, ANIME_RATING_COLS, w10_range, additional_cols)
+        arp_na_table = update_table(arp_df_na, query, param_ranges, is_foundation, ANIME_RATING_COLS, w10_range, additional_cols).fillna('NA')
         
         return ugi_table, ws_table, arp_table, arp_na_table
 
+    # Update the event handlers
+    for component in [search_bar, filter_columns_size, model_type, w10_range, additional_columns]:
+        component.change(
+            fn=update_all_tables,
+            inputs=[search_bar, filter_columns_size, model_type, w10_range, additional_columns],
+            outputs=[leaderboard_table_ugi, leaderboard_table_ws, leaderboard_table_arp, leaderboard_table_arp_na]
+        )
+    
     search_bar.change(
         fn=update_all_tables,
-        inputs=[search_bar, filter_columns_size, w10_range, additional_columns],
+        inputs=[search_bar, filter_columns_size, model_type, w10_range, additional_columns],
         outputs=[leaderboard_table_ugi, leaderboard_table_ws, leaderboard_table_arp, leaderboard_table_arp_na]
     )
     
     filter_columns_size.change(
         fn=update_all_tables,
-        inputs=[search_bar, filter_columns_size, w10_range, additional_columns],
+        inputs=[search_bar, filter_columns_size, model_type, w10_range, additional_columns],
         outputs=[leaderboard_table_ugi, leaderboard_table_ws, leaderboard_table_arp, leaderboard_table_arp_na]
     )
-
+    
+    model_type.change(
+        fn=update_all_tables,
+        inputs=[search_bar, filter_columns_size, model_type, w10_range, additional_columns],
+        outputs=[leaderboard_table_ugi, leaderboard_table_ws, leaderboard_table_arp, leaderboard_table_arp_na]
+    )
+    
     w10_range.change(
         fn=update_all_tables,
-        inputs=[search_bar, filter_columns_size, w10_range, additional_columns],
+        inputs=[search_bar, filter_columns_size, model_type, w10_range, additional_columns],
         outputs=[leaderboard_table_ugi, leaderboard_table_ws, leaderboard_table_arp, leaderboard_table_arp_na]
     )
-
+    
     additional_columns.change(
         fn=update_all_tables,
-        inputs=[search_bar, filter_columns_size, w10_range, additional_columns],
+        inputs=[search_bar, filter_columns_size, model_type, w10_range, additional_columns],
         outputs=[leaderboard_table_ugi, leaderboard_table_ws, leaderboard_table_arp, leaderboard_table_arp_na]
     )