Add task count into table column name

app.py

@@ -55,8 +55,8 @@ with gr.Blocks() as block:
                 )
 
             # Define different captions for each table
-            default_caption = "**Table 1: MEGA-Bench full results.** <br> The Core set contains $N_{\\text{core}} = 440$ tasks evaluated by rule-based metrics, and the Open-ended set contains $N_{\\text{open}} = 65$ tasks evaluated by a VLM judge (we use GPT-4o-0806). <br> $\\text{Overall} \\ = \\ \\frac{\\max(\\text{Core w/o CoT}, \\ \\text{Core w/ CoT}) \\ \\cdot \\ N_{\\text{core}} \\ + \\ \\text{Open-ended} \\ \\cdot \\ N_{\\text{open}}}{N_{\\text{core}} \\ + \\ N_{\\text{open}}}$"
-            core_single_image_caption = "**Table 2: MEGA-Bench Core Single-image results.** <br> This subset contains 273 single-image tasks from the Core set of the benchmark. For open-source models, we drop the image input in the 1-shot demonstration example so that the entire query contains a single image only. Compared to the default table, some models with only single-image support are added."
+            default_caption = "**Table 1: MEGA-Bench full results.** The number in the parentheses is the number of tasks of each keyword. <br> The Core set contains $N_{\\text{core}} = 440$ tasks evaluated by rule-based metrics, and the Open-ended set contains $N_{\\text{open}} = 65$ tasks evaluated by a VLM judge (we use GPT-4o-0806). <br> $\\text{Overall} \\ = \\ \\frac{\\max(\\text{Core w/o CoT}, \\ \\text{Core w/ CoT}) \\ \\cdot \\ N_{\\text{core}} \\ + \\ \\text{Open-ended} \\ \\cdot \\ N_{\\text{open}}}{N_{\\text{core}} \\ + \\ N_{\\text{open}}}$"
+            core_single_image_caption = "**Table 2: MEGA-Bench Core Single-image results.** The number in the parentheses is the number of tasks in each keyword. <br> This subset contains 273 single-image tasks from the Core set of the benchmark. For open-source models, we drop the image input in the 1-shot demonstration example so that the entire query contains a single image only. <br> Compared to the default table, some models with only single-image support are added."
 
             caption_component = gr.Markdown(
                 value=default_caption,

utils.py

@@ -102,13 +102,35 @@ class BaseDataLoader:
         self.MODEL_GROUPS = self._initialize_model_groups()
 
     def _initialize_super_groups(self):
-        # Define the desired order of super groups
+        # Get a sample model to access the structure
+        sample_model = next(iter(self.MODEL_DATA))
         
-        groups = {DIMENSION_NAME_MAP[dim]: [KEYWORD_NAME_MAP.get(k, k) for k in self.MODEL_DATA[next(iter(self.MODEL_DATA))][dim].keys()] 
-                 for dim in self.MODEL_DATA[next(iter(self.MODEL_DATA))]}
+        # Create groups with task counts
+        groups = {}
+        self.keyword_display_map = {}  # Add this map to store display-to-original mapping
         
-        order = ["Skills", "Application", "Output Format", "Input Format", "Visual Input Number"]
-        # Sort the dictionary based on the predefined order
+        for dim in self.MODEL_DATA[sample_model]:
+            dim_name = DIMENSION_NAME_MAP[dim]
+            # Create a list of tuples (display_name, count, keyword) for sorting
+            keyword_info = []
+            
+            for keyword in self.MODEL_DATA[sample_model][dim]:
+                # Get the task count for this keyword
+                task_count = self.MODEL_DATA[sample_model][dim][keyword]["count"]
+                original_name = KEYWORD_NAME_MAP.get(keyword, keyword)
+                display_name = f"{original_name}({task_count})"
+                keyword_info.append((display_name, task_count, keyword))
+            
+            # Sort by count (descending) and then by display name (for ties)
+            keyword_info.sort(key=lambda x: (-x[1], x[0]))
+            
+            # Store sorted display names and update mapping
+            groups[dim_name] = [info[0] for info in keyword_info]
+            for display_name, _, keyword in keyword_info:
+                self.keyword_display_map[display_name] = keyword
+        
+        # Sort based on predefined order
+        order = ["Application", "Skills", "Output Format", "Input Format", "Visual Input Number"]
         return {k: groups[k] for k in order if k in groups}
 
     def _initialize_model_groups(self) -> Dict[str, list]:
@@ -167,12 +189,12 @@ class DefaultDataLoader(BaseDataLoader):
                 "Core(w/ CoT)": round(core_cot_score * 100, 2),
                 "Open-ended": round(summary["open"]["macro_mean_score"] * 100, 2)
             }
-            for keyword in self.SUPER_GROUPS[selected_super_group]:
-                original_keyword = get_original_keyword(keyword)
+            for display_name in self.SUPER_GROUPS[selected_super_group]:
+                original_keyword = self.keyword_display_map[display_name]
                 if original_dimension in model_data and original_keyword in model_data[original_dimension]:
-                    row[keyword] = round(model_data[original_dimension][original_keyword]["average_score"] * 100, 2)
+                    row[display_name] = round(model_data[original_dimension][original_keyword]["average_score"] * 100, 2)
                 else:
-                    row[keyword] = None
+                    row[display_name] = None
             data.append(row)
         
         df = pd.DataFrame(data)
@@ -209,12 +231,12 @@ class CoreSingleDataLoader(BaseDataLoader):
                 "Models": get_display_model_name(model),
                 "Core SI": round(core_si_score * 100, 2),
             }
-            for keyword in self.SUPER_GROUPS[selected_super_group]:
-                original_keyword = get_original_keyword(keyword)
+            for display_name in self.SUPER_GROUPS[selected_super_group]:
+                original_keyword = self.keyword_display_map[display_name]
                 if original_dimension in model_data and original_keyword in model_data[original_dimension]:
-                    row[keyword] = round(model_data[original_dimension][original_keyword]["average_score"] * 100, 2)
+                    row[display_name] = round(model_data[original_dimension][original_keyword]["average_score"] * 100, 2)
                 else:
-                    row[keyword] = None
+                    row[display_name] = None
             data.append(row)
         
         df = pd.DataFrame(data)