Update space

app.py

@@ -1,6 +1,8 @@
 import gradio as gr
 import pandas as pd
 import numpy as np
+from typing import List, Dict
+
 
 # Sample data based on your table (you'll need to update this with the full dataset)
 data_synthesized_full = {
@@ -51,105 +53,149 @@ data_human_generated = {
     'STARK-PRIME_MRR': [30.37, 7.05, 10.07, 9.39, 26.35, 24.33, 15.24, 34.28, 32.98, 19.67, 36.32, 34.82]
 }
 
-df_synthesized_full = pd.DataFrame(data_synthesized_full)
-df_synthesized_10 = pd.DataFrame(data_synthesized_10)
-df_human_generated = pd.DataFrame(data_human_generated)
-
-def format_dataframe(df, dataset):
-    # Filter the dataframe for the selected dataset
-    columns = ['Method'] + [col for col in df.columns if dataset in col]
-    filtered_df = df[columns].copy()
-    
-    # Rename columns
-    filtered_df.columns = [col.split('_')[-1] if '_' in col else col for col in filtered_df.columns]
-    
-    # Sort by MRR
-    filtered_df = filtered_df.sort_values('MRR', ascending=False)
-    
-    return filtered_df
-
-model_types = {
-    'Sparse Retriever': ['BM25'],
-    'Small Dense Retrievers': ['DPR (roberta)', 'ANCE (roberta)', 'QAGNN (roberta)'],
-    'LLM-based Dense Retrievers': ['ada-002', 'voyage-l2-instruct', 'LLM2Vec', 'GritLM-7b'],
-    'Multivector Retrievers': ['multi-ada-002', 'ColBERTv2'],
-    'LLM Rerankers': ['Claude3 Reranker', 'GPT4 Reranker']
-}
-
-def filter_by_model_type(df, selected_types):
-    if not selected_types:  # If no types are selected, return an empty DataFrame
-        return df.head(0)
-    selected_models = [model for type in selected_types for model in model_types[type]]
-    return df[df['Method'].isin(selected_models)]
-
-def format_dataframe(df, dataset):
-    columns = ['Method'] + [col for col in df.columns if dataset in col]
-    filtered_df = df[columns].copy()
-    filtered_df.columns = [col.split('_')[-1] if '_' in col else col for col in filtered_df.columns]
-    filtered_df = filtered_df.sort_values('MRR', ascending=False)
-    return filtered_df
-
-def update_tables(selected_types):
-    filtered_df_full = filter_by_model_type(df_synthesized_full, selected_types)
-    filtered_df_10 = filter_by_model_type(df_synthesized_10, selected_types)
-    filtered_df_human = filter_by_model_type(df_human_generated, selected_types)
-    
-    outputs = []
-    for df in [filtered_df_full, filtered_df_10, filtered_df_human]:
-        for dataset in ['AMAZON', 'MAG', 'PRIME']:
-            outputs.append(format_dataframe(df, f"STARK-{dataset}"))
-    
-    return outputs
-
-css = """
-table > thead {
-    white-space: normal
-}
-
-table {
-    --cell-width-1: 250px
-}
-
-table > tbody > tr > td:nth-child(2) > div {
-    overflow-x: auto
-}
-"""
-
-with gr.Blocks(css=css) as demo:
-    gr.Markdown("# Semi-structured Retrieval Benchmark (STaRK) Leaderboard")
-    gr.Markdown("Refer to the [STaRK paper](https://arxiv.org/pdf/2404.13207) for details on metrics, tasks and models.")
+class DataManager:
+    def __init__(self, data_synthesized_full: Dict, data_synthesized_10: Dict, data_human_generated: Dict):
+        self.df_synthesized_full = pd.DataFrame(data_synthesized_full)
+        self.df_synthesized_10 = pd.DataFrame(data_synthesized_10)
+        self.df_human_generated = pd.DataFrame(data_human_generated)
+        
+        self.model_types = {
+            'Sparse Retriever': ['BM25'],
+            'Small Dense Retrievers': ['DPR (roberta)', 'ANCE (roberta)', 'QAGNN (roberta)'],
+            'LLM-based Dense Retrievers': ['ada-002', 'voyage-l2-instruct', 'LLM2Vec', 'GritLM-7b'],
+            'Multivector Retrievers': ['multi-ada-002', 'ColBERTv2'],
+            'LLM Rerankers': ['Claude3 Reranker', 'GPT4 Reranker']
+        }
+        
+        self.metrics = ['Hit@1', 'Hit@5', 'R@20', 'MRR']
+        self.datasets = ['AMAZON', 'MAG', 'PRIME']
+
+    def filter_by_model_type(self, df: pd.DataFrame, selected_types: List[str]) -> pd.DataFrame:
+        if not selected_types:
+            return df.head(0)
+        selected_models = [model for type in selected_types for model in self.model_types[type]]
+        return df[df['Method'].isin(selected_models)]
+
+    def format_dataframe(self, df: pd.DataFrame, dataset: str) -> pd.DataFrame:
+        columns = ['Method'] + [col for col in df.columns if dataset in col]
+        filtered_df = df[columns].copy()
+        filtered_df.columns = [col.split('_')[-1] if '_' in col else col for col in filtered_df.columns]
+        
+        # Format numeric columns to 2 decimal places
+        for col in filtered_df.columns:
+            if col != 'Method':
+                filtered_df[col] = filtered_df[col].round(2)
+        
+        # Sort by MRR by default
+        filtered_df = filtered_df.sort_values('MRR', ascending=False)
+        return filtered_df
+
+    def get_best_model(self, df: pd.DataFrame, metric: str) -> str:
+        return df.loc[df[metric].idxmax(), 'Method']
+
+# Custom components
+def create_metric_summary(df: pd.DataFrame, dataset: str) -> str:
+    best_mrr = df['MRR'].max()
+    best_hit1 = df['Hit@1'].max()
+    best_model_mrr = df.loc[df['MRR'].idxmax(), 'Method']
+    best_model_hit1 = df.loc[df['Hit@1'].idxmax(), 'Method']
     
-    with gr.Row():
-        model_type_filter = gr.CheckboxGroup(
-            choices=list(model_types.keys()),
-            value=list(model_types.keys()),
-            label="Model types",
-            interactive=True
+    return f"""
+    ### {dataset} Dataset Summary
+    - Best MRR: {best_mrr:.2f}% ({best_model_mrr})
+    - Best Hit@1: {best_hit1:.2f}% ({best_model_hit1})
+    """
+
+# Main application
+def create_app(data_manager: DataManager):
+    with gr.Blocks(css="""
+        .metric-summary { margin: 1rem 0; padding: 1rem; background: #f7f7f7; border-radius: 4px; }
+        .table-container { margin-top: 1rem; }
+        .model-filter { margin-bottom: 1rem; }
+        """) as demo:
+        
+        gr.Markdown("# Semi-structured Retrieval Benchmark (STaRK) Leaderboard")
+        gr.Markdown("### An evaluation benchmark for semi-structured text retrieval")
+        
+        with gr.Row():
+            with gr.Column(scale=3):
+                model_type_filter = gr.CheckboxGroup(
+                    choices=list(data_manager.model_types.keys()),
+                    value=list(data_manager.model_types.keys()),
+                    label="Model Types",
+                    interactive=True,
+                    elem_classes=["model-filter"]
+                )
+            
+            with gr.Column(scale=1):
+                sort_by = gr.Radio(
+                    choices=data_manager.metrics,
+                    value="MRR",
+                    label="Sort by Metric",
+                    interactive=True
+                )
+
+        all_dataframes = []
+        
+        with gr.Tabs() as tabs:
+            data_sources = [
+                ("Synthesized (Full)", data_manager.df_synthesized_full),
+                ("Synthesized (10%)", data_manager.df_synthesized_10),
+                ("Human-Generated", data_manager.df_human_generated)
+            ]
+            
+            for source_name, source_df in data_sources:
+                with gr.TabItem(source_name):
+                    for dataset in data_manager.datasets:
+                        with gr.Box():
+                            gr.Markdown(create_metric_summary(
+                                data_manager.format_dataframe(source_df, f"STARK-{dataset}"),
+                                dataset
+                            ))
+                            df_display = gr.DataFrame(
+                                interactive=False,
+                                elem_classes=["table-container"]
+                            )
+                            all_dataframes.append(df_display)
+        
+        def update_tables(selected_types: List[str], sort_metric: str):
+            outputs = []
+            for df_source in [data_manager.df_synthesized_full, 
+                            data_manager.df_synthesized_10,
+                            data_manager.df_human_generated]:
+                filtered_df = data_manager.filter_by_model_type(df_source, selected_types)
+                for dataset in data_manager.datasets:
+                    formatted_df = data_manager.format_dataframe(filtered_df, f"STARK-{dataset}")
+                    formatted_df = formatted_df.sort_values(sort_metric, ascending=False)
+                    outputs.append(formatted_df)
+            return outputs
+
+        # Register event handlers
+        model_type_filter.change(
+            update_tables,
+            inputs=[model_type_filter, sort_by],
+            outputs=all_dataframes
+        )
+        
+        sort_by.change(
+            update_tables,
+            inputs=[model_type_filter, sort_by],
+            outputs=all_dataframes
         )
-    
-    all_dfs = []
-    
-    with gr.Tabs() as outer_tabs:
-        for tab_name, df_source in [("Synthesized (full)", df_synthesized_full), 
-                                    ("Synthesized (10%)", df_synthesized_10), 
-                                    ("Human-Generated", df_human_generated)]:
-            with gr.TabItem(tab_name):
-                with gr.Tabs() as inner_tabs:
-                    for dataset in ['AMAZON', 'MAG', 'PRIME']:
-                        with gr.TabItem(dataset):
-                            df = gr.DataFrame(interactive=False)
-                            all_dfs.append(df)
 
-    model_type_filter.change(
-        update_tables,
-        inputs=[model_type_filter],
-        outputs=all_dfs
-    )
+        # Initial load
+        demo.load(
+            update_tables,
+            inputs=[model_type_filter, sort_by],
+            outputs=all_dataframes
+        )
 
-    demo.load(
-        update_tables,
-        inputs=[model_type_filter],
-        outputs=all_dfs
-    )
+    return demo
 
-demo.launch()
+if __name__ == "__main__":
+    # Initialize data manager with your existing data
+    data_manager = DataManager(data_synthesized_full, data_synthesized_10, data_human_generated)
+    
+    # Create and launch the app
+    demo = create_app(data_manager)
+    demo.launch()