Update space

app.py

@@ -1,104 +1,12 @@
 import gradio as gr
 import pandas as pd
+import numpy as np
 import os
 import re
 from datetime import datetime
 import json
-import pandas as pd
-import torch
-import numpy as np
-from tqdm import tqdm
-from concurrent.futures import ProcessPoolExecutor, as_completed
-
-from stark_qa import load_qa
-from stark_qa.evaluator import Evaluator
-
-
-def process_single_instance(args):
-    idx, eval_csv, qa_dataset, evaluator, eval_metrics = args
-    query, query_id, answer_ids, meta_info = qa_dataset[idx]
-
-    try:
-        pred_rank = eval_csv[eval_csv['query_id'] == query_id]['pred_rank'].item()
-    except IndexError:
-        raise IndexError(f'Error when processing query_id={query_id}, please make sure the predicted results exist for this query.')
-    except Exception as e:
-        raise RuntimeError(f'Unexpected error occurred while fetching prediction rank for query_id={query_id}: {e}')
-
-    if isinstance(pred_rank, str):
-        try:
-            pred_rank = eval(pred_rank)
-        except SyntaxError as e:
-            raise ValueError(f'Failed to parse pred_rank as a list for query_id={query_id}: {e}')
-    
-    if not isinstance(pred_rank, list):
-        raise TypeError(f'Error when processing query_id={query_id}, expected pred_rank to be a list but got {type(pred_rank)}.')
-
-    pred_dict = {pred_rank[i]: -i for i in range(min(100, len(pred_rank)))}
-    answer_ids = torch.LongTensor(answer_ids)
-    result = evaluator.evaluate(pred_dict, answer_ids, metrics=eval_metrics)
-
-    result["idx"], result["query_id"] = idx, query_id
-    return result
-
-
-def compute_metrics(csv_path: str, dataset: str, split: str, num_workers: int = 4):
-    candidate_ids_dict = {
-        'amazon': [i for i in range(957192)],
-        'mag': [i for i in range(1172724, 1872968)],
-        'prime': [i for i in range(129375)]
-    }
-    try:
-        eval_csv = pd.read_csv(csv_path)
-        if 'query_id' not in eval_csv.columns:
-            raise ValueError('No `query_id` column found in the submitted csv.')
-        if 'pred_rank' not in eval_csv.columns:
-            raise ValueError('No `pred_rank` column found in the submitted csv.')
-
-        eval_csv = eval_csv[['query_id', 'pred_rank']]
-
-        if dataset not in candidate_ids_dict:
-            raise ValueError(f"Invalid dataset '{dataset}', expected one of {list(candidate_ids_dict.keys())}.")
-        if split not in ['test', 'test-0.1', 'human_generated_eval']:
-            raise ValueError(f"Invalid split '{split}', expected one of ['test', 'test-0.1', 'human_generated_eval'].")
 
-        evaluator = Evaluator(candidate_ids_dict[dataset])
-        eval_metrics = ['hit@1', 'hit@5', 'recall@20', 'mrr']
-        qa_dataset = load_qa(dataset, human_generated_eval=split == 'human_generated_eval')
-        split_idx = qa_dataset.get_idx_split()
-        all_indices = split_idx[split].tolist()
-
-        results_list = []
-        query_ids = []
-
-        # Prepare args for each worker
-        args = [(idx, eval_csv, qa_dataset, evaluator, eval_metrics) for idx in all_indices]
-
-        with ProcessPoolExecutor(max_workers=num_workers) as executor:
-            futures = [executor.submit(process_single_instance, arg) for arg in args]
-            for future in tqdm(as_completed(futures), total=len(futures)):
-                result = future.result()  # This will raise an error if the worker encountered one
-                results_list.append(result)
-                query_ids.append(result['query_id'])
-
-        # Concatenate results and compute final metrics
-        eval_csv = pd.concat([eval_csv, pd.DataFrame(results_list)], ignore_index=True)
-        final_results = {
-            metric: np.mean(eval_csv[eval_csv['query_id'].isin(query_ids)][metric]) for metric in eval_metrics
-        }
-        return final_results
-
-    except pd.errors.EmptyDataError:
-        return "Error: The CSV file is empty or could not be read. Please check the file and try again."
-    except FileNotFoundError:
-        return f"Error: The file {csv_path} could not be found. Please check the file path and try again."
-    except Exception as error:
-        return f"{error}"
-
-
-
-
-# Sample data based on your table (you'll need to update this with the full dataset)
+# Data dictionaries for leaderboard
 data_synthesized_full = {
     'Method': ['BM25', 'DPR (roberta)', 'ANCE (roberta)', 'QAGNN (roberta)', 'ada-002', 'voyage-l2-instruct', 'LLM2Vec', 'GritLM-7b', 'multi-ada-002', 'ColBERTv2'],
     'STARK-AMAZON_Hit@1': [44.94, 15.29, 30.96, 26.56, 39.16, 40.93, 21.74, 42.08, 40.07, 46.10],
@@ -147,9 +55,21 @@ 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]
 }
 
+# Initialize DataFrames
 df_synthesized_full = pd.DataFrame(data_synthesized_full)
 df_synthesized_10 = pd.DataFrame(data_synthesized_10)
 df_human_generated = pd.DataFrame(data_human_generated)
+
+# Model type definitions
+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']
+}
+
+# Submission form validation functions
 def validate_email(email_str):
     """Validate email format(s)"""
     emails = [e.strip() for e in email_str.split(';')]
@@ -169,11 +89,9 @@ def validate_csv(file_obj):
         df = pd.read_csv(file_obj.name)
         required_cols = ['query_id', 'pred_rank']
         
-        # Check columns
         if not all(col in df.columns for col in required_cols):
             return False, "CSV must contain 'query_id' and 'pred_rank' columns"
             
-        # Check pred_rank format and length
         try:
             first_rank = eval(df['pred_rank'].iloc[0]) if isinstance(df['pred_rank'].iloc[0], str) else df['pred_rank'].iloc[0]
             if not isinstance(first_rank, list) or len(first_rank) < 20:
@@ -190,16 +108,39 @@ def save_submission(submission_data):
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     submission_id = f"{submission_data['team_name']}_{timestamp}"
     
-    # Create submissions directory if it doesn't exist
     os.makedirs("submissions", exist_ok=True)
-    
-    # Save submission data
     submission_path = f"submissions/{submission_id}.json"
     with open(submission_path, 'w') as f:
         json.dump(submission_data, f, indent=4)
     
     return submission_id
 
+# Leaderboard functions
+def filter_by_model_type(df, selected_types):
+    if not selected_types:
+        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
+
 def process_submission(
     method_name, team_name, dataset, split, contact_email,
     code_repo, csv_file, model_description, hardware, paper_link
@@ -270,6 +211,21 @@ def process_submission(
     except Exception as e:
         return f"Error processing submission: {str(e)}"
 
+# CSS styling
+css = """
+table > thead {
+    white-space: normal
+}
+
+table {
+    --cell-width-1: 250px
+}
+
+table > tbody > tr > td:nth-child(2) > div {
+    overflow-x: auto
+}
+"""
+
 def add_submission_form(demo):
     with demo:
         gr.Markdown("---")
@@ -338,102 +294,47 @@ def add_submission_form(demo):
             ],
             outputs=result
         )
-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.")
-    
-    with gr.Row():
-        model_type_filter = gr.CheckboxGroup(
-            choices=list(model_types.keys()),
-            value=list(model_types.keys()),
-            label="Model types",
-            interactive=True
+# Main application
+if __name__ == "__main__":
+    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.")
+        
+        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
+        )
+        
+        # Add submission form
+        add_submission_form(demo)
     
-    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
-    )
-    add_submission_form(demo)
-
-demo.launch()
+    demo.launch()