feat : metrics dropdown added to gradio

app.py

@@ -1,263 +1,86 @@
-import json
-import os
-
+from data.model_handler import ModelHandler
+from app.utils import add_rank_and_format, get_refresh_function
 import gradio as gr
-import pandas as pd
-from huggingface_hub import HfApi, hf_hub_download, get_collection
-from huggingface_hub.repocard import metadata_load
-from typing import Dict
-
-
-def get_datasets_nickname() -> Dict: 
-
-    datasets_nickname = {}
-
-    collection = get_collection("vidore/vidore-benchmark-667173f98e70a1c0fa4db00d")
-
-    collection_items = collection.items
-
-    for item in collection_items:
-        dataset_name = item.item_id
-
-        if 'arxivqa' in dataset_name:
-            datasets_nickname[dataset_name] = 'ArxivQA'
-            datasets_nickname[dataset_name + '_ocr_chunk'] = 'ArxivQA'
-            datasets_nickname[dataset_name + '_captioning'] = 'ArxivQA'
-        
-        elif 'docvqa' in dataset_name:
-            datasets_nickname[dataset_name] = 'DocVQA'
-            datasets_nickname[dataset_name + '_ocr_chunk'] = 'DocVQA'
-            datasets_nickname[dataset_name + '_captioning'] = 'DocVQA'
-
-        elif 'infovqa' in dataset_name:
-            datasets_nickname[dataset_name] = 'InfoVQA'
-            datasets_nickname[dataset_name + '_ocr_chunk'] = 'InfoVQA'
-            datasets_nickname[dataset_name + '_captioning'] = 'InfoVQA'
-        
-        elif 'tabfquad' in dataset_name:
-            datasets_nickname[dataset_name] = 'TabFQuad'
-            datasets_nickname[dataset_name + '_ocr_chunk'] = 'TabFQuad'
-            datasets_nickname[dataset_name + '_captioning'] = 'TabFQuad'
-        
-        elif 'tatdqa' in dataset_name:
-            datasets_nickname[dataset_name] = 'TATDQA'
-            datasets_nickname[dataset_name + '_ocr_chunk'] = 'TATDQA'
-            datasets_nickname[dataset_name + '_captioning'] = 'TATDQA'
-        
-        elif 'shiftproject' in dataset_name:
-            datasets_nickname[dataset_name] = 'ShiftProject'
-            datasets_nickname[dataset_name + '_ocr_chunk'] = 'ShiftProject'
-            datasets_nickname[dataset_name + '_captioning'] = 'ShiftProject'
-        
-        elif 'artificial_intelligence' in dataset_name:
-            datasets_nickname[dataset_name] = 'Artificial Intelligence'
-            datasets_nickname[dataset_name + '_ocr_chunk'] = 'Artificial Intelligence'
-            datasets_nickname[dataset_name + '_captioning'] = 'Artificial Intelligence'
-        
-        elif 'energy' in dataset_name:
-            datasets_nickname[dataset_name] = 'Energy'
-            datasets_nickname[dataset_name + '_ocr_chunk'] = 'Energy'
-            datasets_nickname[dataset_name + '_captioning'] = 'Energy'
-        
-        elif 'government_reports' in dataset_name:
-            datasets_nickname[dataset_name] = 'Government Reports'
-            datasets_nickname[dataset_name + '_ocr_chunk'] = 'Government Reports'
-            datasets_nickname[dataset_name + '_captioning'] = 'Government Reports'
-
-        elif 'healthcare' in dataset_name:
-            datasets_nickname[dataset_name] = 'Healthcare'
-            datasets_nickname[dataset_name + '_ocr_chunk'] = 'Healthcare'
-            datasets_nickname[dataset_name + '_captioning'] = 'Healthcare'
-
-    return datasets_nickname
-
-
-def make_clickable_model(model_name, link=None):
-    
-    if link is None:
-        desanitized_model_name = model_name.replace("_", "/")
-        
-        if '/captioning' in desanitized_model_name:
-            desanitized_model_name = desanitized_model_name.replace('/captioning', '')
-        if '/ocr' in desanitized_model_name:
-            desanitized_model_name = desanitized_model_name.replace('/ocr', '')
-
-        link = "https://huggingface.co/" + desanitized_model_name
-        
-    # Remove user from model name
-    # return (
-    #     f'<a target="_blank" style="text-decoration: underline" href="{link}">{model_name.split("/")[-1]}</a>'
-    # )
-    return f'<a target="_blank" style="text-decoration: underline" href="{link}">{model_name}</a>'
-
-
-def add_rank(df):
-    cols_to_rank = [
-        col
-        for col in df.columns
-        if col
-        not in [
-            "Model",
-            "Model Size (Million Parameters)",
-            "Memory Usage (GB, fp32)",
-            "Embedding Dimensions",
-            "Max Tokens",
-        ]
-    ]
-    if len(cols_to_rank) == 1:
-        df.sort_values(cols_to_rank[0], ascending=False, inplace=True)
-    else:
-        df.insert(len(df.columns) - len(cols_to_rank), "Average", df[cols_to_rank].mean(axis=1, skipna=False))
-        df.sort_values("Average", ascending=False, inplace=True)
-    df.insert(0, "Rank", list(range(1, len(df) + 1)))
-    df = df.round(2)
-    # Fill NaN after averaging
-    df.fillna("", inplace=True)
-    return df
-
-
-def get_vidore_data():
-    api = HfApi()
 
-    # local cache path
-    model_infos_path = "model_infos.json"
-    metric = "ndcg_at_5"
+METRICS = ["ndcg_at_5", "recall_at_1", "recall_at_5", "mrr_at_5"]
 
-
-    MODEL_INFOS = {}
-    if os.path.exists(model_infos_path):
-        with open(model_infos_path) as f:
-            MODEL_INFOS = json.load(f)
-
-    models = api.list_models(filter="vidore")
-    repositories = [model.modelId for model in models]  
-
-    datasets_nickname = get_datasets_nickname()
-    for repo_id in repositories:
-        files = [f for f in api.list_repo_files(repo_id) if f.endswith('_metrics.json')]
-        if len(files) == 0:
-            continue
-        else :
-            for file in files:
-                model_name = file.split('_metrics.json')[0]
-
-                if model_name not in MODEL_INFOS:
-                    readme_path = hf_hub_download(repo_id, filename="README.md")
-                    meta = metadata_load(readme_path)
-                    try:
-                        result_path = hf_hub_download(repo_id, filename= file)
-
-                        with open(result_path) as f:
-                            results = json.load(f)
-
-                        for dataset in results:
-                            results[dataset] = {key: value for key, value in results[dataset].items() if metric in key}
-
-                        MODEL_INFOS[model_name] = {"meta":meta, "results": results}
-                    except Exception as e:  
-                        print(f"Error loading {model_name} - {e}")
-                        continue
+def main(): 
+    model_handler = ModelHandler()
+    initial_metric = "ndcg_at_5"
     
-    model_res = {}
-    df = None
-    if len(MODEL_INFOS) > 0:
-        for model in MODEL_INFOS.keys():
-            res = MODEL_INFOS[model]["results"]
-            dataset_res = {}
-            for dataset in res.keys():
-                if "validation_set" == dataset:
-                    continue
-                dataset_res[datasets_nickname[dataset]] = res[dataset][metric]
-            model_res[model] = dataset_res
-
-        df = pd.DataFrame(model_res).T
-        
-    #save model_infos
-    with open(model_infos_path, "w") as f:
-        json.dump(MODEL_INFOS, f)
-    
-    return df
-
+    data = model_handler.get_vidore_data(initial_metric)
+    data = add_rank_and_format(data)
 
-def add_rank_and_format(df):
-    df = df.reset_index()
-    df = df.rename(columns={"index": "Model"})
-    df = add_rank(df)
-    df["Model"] = df["Model"].apply(make_clickable_model)
-    return df
+    NUM_DATASETS = len(data.columns) - 3
+    NUM_SCORES = len(data) * NUM_DATASETS
+    NUM_MODELS = len(data)
 
+    css = """
+    table > thead {
+        white-space: normal
+    }
 
-# 1. Force headers to wrap
-# 2. Force model column (maximum) width
-# 3. Prevent model column from overflowing, scroll instead
-# 4. Prevent checkbox groups from taking up too much space
+    table {
+        --cell-width-1: 250px
+    }
 
-css = """
-table > thead {
-    white-space: normal
-}
+    table > tbody > tr > td:nth-child(2) > div {
+        overflow-x: auto
+    }
 
-table {
-    --cell-width-1: 250px
-}
+    .filter-checkbox-group {
+        max-width: max-content;
+    }
+    """
 
-table > tbody > tr > td:nth-child(2) > div {
-    overflow-x: auto
-}
+    with gr.Blocks(css=css) as block:
+        gr.Markdown("# ViDoRe: The Visual Document Retrieval Benchmark 📚🔍")
+        gr.Markdown("## From the paper - ColPali: Efficient Document Retrieval with Vision Language Models 👀")
 
-.filter-checkbox-group {
-    max-width: max-content;
-}
-"""
+        gr.Markdown(
+            """
+        Visual Document Retrieval Benchmark leaderboard. To submit, refer to the <a href="https://github.com/tonywu71/vidore-benchmark/" target="_blank" style="text-decoration: underline">ViDoRe GitHub repository</a>.  Refer to the [ColPali paper](https://arxiv.org/abs/XXXX.XXXXX) for details on metrics, tasks and models.
+        """
+        )
+        #all_columns = list(data.columns)
+        #default_columns = all_columns
 
+        with gr.Row():
+            metric_dropdown = gr.Dropdown(choices=METRICS, value=initial_metric, label="Select Metric")
+            #column_checkboxes = gr.CheckboxGroup(choices=all_columns, value=default_columns, label="Select Columns to Display")
 
-def get_refresh_function():
-    def _refresh():
-        data_task_category = get_vidore_data()
-        return add_rank_and_format(data_task_category)
+        with gr.Row():
+            datatype = ["number", "markdown"] + ["number"] * (NUM_DATASETS + 1)
+            dataframe = gr.Dataframe(data, datatype=datatype, type="pandas")
 
-    return _refresh
+        with gr.Row():
+            refresh_button = gr.Button("Refresh")
+            refresh_button.click(get_refresh_function(), inputs=[metric_dropdown], outputs=dataframe, concurrency_limit=20)
 
-data = get_vidore_data()
-data = add_rank_and_format(data)
 
-NUM_DATASETS = len(data.columns) - 3
-NUM_SCORES = len(data) * NUM_DATASETS
-NUM_MODELS = len(data)
+        # Automatically refresh the dataframe when the dropdown value changes
+        metric_dropdown.change(get_refresh_function(), inputs=[metric_dropdown], outputs=dataframe)
+        #column_checkboxes.change(get_refresh_function(), inputs=[metric_dropdown, column_checkboxes], outputs=dataframe)
 
-with gr.Blocks(css=css) as block:
-    gr.Markdown("# ViDoRe: The Visual Document Retrieval Benchmark 📚🔍")
-    gr.Markdown("## From the paper - ColPali: Efficient Document Retrieval with Vision Language Models 👀")
 
-    gr.Markdown(
+        gr.Markdown(
+            f"""
+        - **Total Datasets**: {NUM_DATASETS}
+        - **Total Scores**: {NUM_SCORES}
+        - **Total Models**: {NUM_MODELS}
         """
-    Visual Document Retrieval Benchmark leaderboard. To submit, refer to the <a href="https://github.com/tonywu71/vidore-benchmark/" target="_blank" style="text-decoration: underline">ViDoRe GitHub repository</a>.  Refer to the [ColPali paper](https://arxiv.org/abs/XXXX.XXXXX) for details on metrics, tasks and models.
-    """
-    )
+            + r"""
+        Please consider citing:
 
-    with gr.Row():
-        datatype = ["number", "markdown"] + ["number"] * (NUM_DATASETS + 1)
-        dataframe = gr.Dataframe(data, datatype=datatype, type="pandas", height=500)
-
-    with gr.Row():
-        refresh_button = gr.Button("Refresh")
-        refresh_button.click(get_refresh_function(), inputs=None, outputs=dataframe, concurrency_limit=20)
+        ```bibtex
+        INSERT LATER
+        ```
+        """
+        )
 
-    gr.Markdown(
-        f"""
-    - **Total Datasets**: {NUM_DATASETS}
-    - **Total Scores**: {NUM_SCORES}
-    - **Total Models**: {NUM_MODELS}
-    """
-        + r"""
-    Please consider citing:
+    block.queue(max_size=10).launch(debug=True)
 
-    ```bibtex
-    INSERT LATER
-    ```
-    """
-    )
 
+if __name__ == "__main__":    
+    main()
 
-if __name__ == "__main__":
-    block.queue(max_size=10).launch(debug=True)