add colors

app.py

@@ -193,7 +193,36 @@ def random_sample(r: gr.Request, subset):
 
 subsets = eval_set.unique("subset")
 
-def regex_table(dataframe, regex, filter_button):
+color_map = {
+    "Generative": "#7497db",
+    "Custom Classifier": "#E8ECF2",
+    "Seq. Classifier": "#ffcd75",
+    "DPO": "#75809c",
+}
+def color_model_type_column(df, color_map):
+    """
+    Apply color to the 'Model Type' column of the DataFrame based on a given color mapping.
+
+    Parameters:
+    df (pd.DataFrame): The DataFrame containing the 'Model Type' column.
+    color_map (dict): A dictionary mapping model types to colors.
+
+    Returns:
+    pd.Styler: The styled DataFrame.
+    """
+    # Function to apply color based on the model type
+    def apply_color(val):
+        color = color_map.get(val, "default")  # Default color if not specified in color_map
+        return f'background-color: {color}'
+    
+    # Format for different columns
+    format_dict = {col: "{:.1f}" for col in df.columns if col not in ['Average', 'Model', 'Model Type']}
+    format_dict['Average'] = "{:.2f}"
+    format_dict[''] = "{:d}"
+
+    return df.style.applymap(apply_color, subset=['Model Type']).format(format_dict, na_rep='')
+
+def regex_table(dataframe, regex, filter_button, style=True):
     """
     Takes a model name as a regex, then returns only the rows that has that in it.
     """
@@ -228,6 +257,8 @@ def regex_table(dataframe, regex, filter_button):
         # sort array by Score column
         data = data.sort_values(by='Score', ascending=False)
 
+    data.reset_index(drop=True, inplace=True)
+
     # replace column '' with count/rank
     data[''] = np.arange(1, 1 + len(data))
 
@@ -242,11 +273,15 @@ def regex_table(dataframe, regex, filter_button):
             # replace any data[col].values == '' with np.NaN
             data[col] = data[col].replace('', np.NaN)
             data[col] = np.round(np.array(data[col].values).astype(float), 1)
+    if style:
+        # apply color
+        data = color_model_type_column(data, color_map)
+
     return data
 
 # import ipdb; ipdb.set_trace()
 
-total_models = len(regex_table(rewardbench_data_avg.copy(), "", ["Seq. Classifiers", "DPO", "Custom Classifiers", "Generative"]).values)
+total_models = len(regex_table(rewardbench_data_avg.copy(), "", ["Seq. Classifiers", "DPO", "Custom Classifiers", "Generative"], style=False).values)
 
 with gr.Blocks(css=custom_css) as app:
     # create tabs for the app, moving the current table to one titled "rewardbench" and the benchmark_text to a tab called "About"
@@ -281,7 +316,7 @@ with gr.Blocks(css=custom_css) as app:
                     visible=False,
                 )
                 rewardbench_table = gr.Dataframe(
-                    regex_table(rewardbench_data_avg.copy(), "", ["Seq. Classifiers", "DPO", "Custom Classifiers",  "Generative", "Prior Sets"]).values,
+                    regex_table(rewardbench_data_avg.copy(), "", ["Seq. Classifiers", "DPO", "Custom Classifiers",  "Generative", "Prior Sets"]),
                     datatype=col_types_rewardbench_avg,
                     headers=rewardbench_data_avg.columns.tolist(),
                     elem_id="rewardbench_dataframe_avg",
@@ -306,7 +341,7 @@ with gr.Blocks(css=custom_css) as app:
                     visible=False,
                 )
                 rewardbench_table_detailed = gr.Dataframe(
-                    regex_table(rewardbench_data.copy(), "", ["Seq. Classifiers", "DPO",  "Generative", "Custom Classifiers"]).values,
+                    regex_table(rewardbench_data.copy(), "", ["Seq. Classifiers", "DPO",  "Generative", "Custom Classifiers"]),
                     datatype=col_types_rewardbench,
                     headers=rewardbench_data.columns.tolist(),
                     elem_id="rewardbench_dataframe",
@@ -351,7 +386,7 @@ with gr.Blocks(css=custom_css) as app:
                     visible=False,
                 )
                 pref_sets_table = gr.Dataframe(
-                    regex_table(prefs_data.copy(), "", ["Seq. Classifiers", "DPO", "Custom Classifiers"]).values,
+                    regex_table(prefs_data.copy(), "", ["Seq. Classifiers", "DPO", "Custom Classifiers"]),
                     datatype=col_types_prefs,
                     headers=prefs_data.columns.tolist(),
                     elem_id="prefs_dataframe",