Update space

app.py

@@ -1,8 +1,6 @@
 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 = {
@@ -53,152 +51,105 @@ 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]
 }
 
-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']
+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()
     
-    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; }
-        .dataset-section { border: 1px solid #ddd; padding: 1rem; margin: 1rem 0; border-radius: 4px; }
-        """) as demo:
-        
-        gr.Markdown("# Semi-structured Retrieval Benchmark (STaRK) Leaderboard")
-        gr.Markdown("### An evaluation benchmark for semi-structured text retrieval")
-        gr.Markdown("Refer to the [STaRK paper](https://arxiv.org/pdf/2404.13207) for details on metrics, tasks and models.")
-        
-        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.Row(elem_classes=["dataset-section"]):
-                            with gr.Column():
-                                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
-        )
+    # 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
 
-        # Initial load
-        demo.load(
-            update_tables,
-            inputs=[model_type_filter, sort_by],
-            outputs=all_dataframes
-        )
+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
+}
 
-    return demo
+table {
+    --cell-width-1: 250px
+}
+
+table > tbody > tr > td:nth-child(2) > div {
+    overflow-x: auto
+}
+"""
 
-if __name__ == "__main__":
-    # Initialize data manager with your existing data
-    data_manager = DataManager(data_synthesized_full, data_synthesized_10, data_human_generated)
+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.")
     
-    # Create and launch the app
-    demo = create_app(data_manager)
-    demo.launch()
+    with gr.Row():
+        model_type_filter = gr.CheckboxGroup(
+            choices=list(model_types.keys()),
+            value=list(model_types.keys()),
+            label="Model types",
+            interactive=True
+        )
+    
+    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
+    )
+
+    demo.load(
+        update_tables,
+        inputs=[model_type_filter],
+        outputs=all_dfs
+    )
+
+demo.launch()