refactor: refactor the benchmarks

app.py

@@ -1,46 +1,28 @@
 import gradio as gr
 from apscheduler.schedulers.background import BackgroundScheduler
-from huggingface_hub import snapshot_download
 
 from src.about import (
     INTRODUCTION_TEXT,
-    BENCHMARKS_TEXT,
-    TITLE,
-    EVALUATION_QUEUE_TEXT
+    TITLE
 )
 from src.benchmarks import (
-    DOMAIN_COLS_QA,
-    LANG_COLS_QA,
-    DOMAIN_COLS_LONG_DOC,
-    LANG_COLS_LONG_DOC,
+    qa_benchmark_dict,
+    long_doc_benchmark_dict,
     METRIC_LIST,
     DEFAULT_METRIC_QA,
     DEFAULT_METRIC_LONG_DOC
 )
 from src.display.css_html_js import custom_css
-from src.display.column_names import COL_NAME_RETRIEVAL_MODEL, COL_NAME_RERANKING_MODEL, COL_NAME_REVISION, \
-    COL_NAME_TIMESTAMP, COL_NAME_IS_ANONYMOUS
 from src.envs import (
     API,
     EVAL_RESULTS_PATH,
-    REPO_ID,
-    RESULTS_REPO,
-    TOKEN,
-    BM25_LINK,
-    BENCHMARK_VERSION_LIST,
-    LATEST_BENCHMARK_VERSION
+    REPO_ID
 )
-from src.read_evals import (
-    get_raw_eval_results,
-    get_leaderboard_df
+from src.loaders import (
+    load_eval_results
 )
 from src.utils import (
-    update_metric,
-    upload_file,
-    get_default_cols,
-    submit_results,
-    reset_rank,
-    remove_html
+    update_metric
 )
 from src.display.gradio_formatting import (
     get_version_dropdown,
@@ -51,8 +33,7 @@ from src.display.gradio_formatting import (
     get_language_dropdown,
     get_anonymous_checkbox,
     get_revision_and_ts_checkbox,
-    get_leaderboard_table,
-    get_noreranking_dropdown
+    get_leaderboard_table
 )
 from src.display.gradio_listener import set_listeners
 
@@ -69,65 +50,6 @@ def restart_space():
 #     print(f'failed to download')
 #     restart_space()
 
-from dataclasses import dataclass
-import pandas as pd
-from typing import Optional
-
-
-@dataclass
-class LeaderboardDataStore:
-    raw_data: Optional[list]
-    raw_df_qa: Optional[pd.DataFrame]
-    raw_df_long_doc: Optional[pd.DataFrame]
-    leaderboard_df_qa: Optional[pd.DataFrame]
-    leaderboard_df_long_doc: Optional[pd.DataFrame]
-    reranking_models: Optional[list]
-    types_qa: Optional[list]
-    types_long_doc: Optional[list]
-
-
-def load_leaderboard_data(file_path) -> LeaderboardDataStore:
-    lb_data_store = LeaderboardDataStore(None, None, None, None, None, None, None, None)
-    lb_data_store.raw_data = get_raw_eval_results(file_path)
-    print(f'raw data: {len(lb_data_store.raw_data)}')
-
-    lb_data_store.raw_df_qa = get_leaderboard_df(
-        lb_data_store.raw_data, task='qa', metric=DEFAULT_METRIC_QA)
-    lb_data_store.leaderboard_df_qa = lb_data_store.raw_df_qa.copy()
-    # leaderboard_df_qa = leaderboard_df_qa[has_no_nan_values(df, _benchmark_cols)]
-    print(f'QA data loaded: {lb_data_store.raw_df_qa.shape}')
-    shown_columns_qa, types_qa = get_default_cols(
-        'qa', lb_data_store.leaderboard_df_qa.columns, add_fix_cols=True)
-    lb_data_store.types_qa = types_qa
-    lb_data_store.leaderboard_df_qa = \
-    lb_data_store.leaderboard_df_qa[~lb_data_store.leaderboard_df_qa[COL_NAME_IS_ANONYMOUS]][shown_columns_qa]
-    lb_data_store.leaderboard_df_qa.drop([COL_NAME_REVISION, COL_NAME_TIMESTAMP], axis=1, inplace=True)
-
-    lb_data_store.raw_df_long_doc = get_leaderboard_df(
-        lb_data_store.raw_data, task='long-doc', metric=DEFAULT_METRIC_LONG_DOC)
-    print(f'Long-Doc data loaded: {len(lb_data_store.raw_df_long_doc)}')
-    lb_data_store.leaderboard_df_long_doc = lb_data_store.raw_df_long_doc.copy()
-    shown_columns_long_doc, types_long_doc = get_default_cols(
-        'long-doc', lb_data_store.leaderboard_df_long_doc.columns, add_fix_cols=True)
-    lb_data_store.types_long_doc = types_long_doc
-    lb_data_store.leaderboard_df_long_doc = \
-    lb_data_store.leaderboard_df_long_doc[~lb_data_store.leaderboard_df_long_doc[COL_NAME_IS_ANONYMOUS]][
-        shown_columns_long_doc]
-    lb_data_store.leaderboard_df_long_doc.drop([COL_NAME_REVISION, COL_NAME_TIMESTAMP], axis=1, inplace=True)
-
-    lb_data_store.reranking_models = sorted(
-        list(frozenset([eval_result.reranking_model for eval_result in lb_data_store.raw_data])))
-    return lb_data_store
-
-
-def load_eval_results(file_path: str):
-    output = {}
-    versions = ("AIR-Bench_24.04",)
-    for version in versions:
-        fn = f"{file_path}/{version}"
-        output[version] = load_leaderboard_data(fn)
-    return output
-
 
 data = load_eval_results(EVAL_RESULTS_PATH)
 
@@ -157,6 +79,12 @@ def update_metric_long_doc(
     return update_metric(data["AIR-Bench_24.04"].raw_data, "long-doc", metric, domains, langs, reranking_model, query, show_anonymous, show_revision_and_timestamp)
 
 
+DOMAIN_COLS_QA = list(frozenset([c.domain for c in qa_benchmark_dict.values()]))
+LANG_COLS_QA = list(frozenset([c.lang for c in qa_benchmark_dict.values()]))
+
+DOMAIN_COLS_LONG_DOC = list(frozenset([c.domain for c in long_doc_benchmark_dict.values()]))
+LANG_COLS_LONG_DOC = list(frozenset([c.lang for c in long_doc_benchmark_dict.values()]))
+
 demo = gr.Blocks(css=custom_css)
 
 with demo:

src/benchmarks.py

@@ -1,16 +1,10 @@
 from dataclasses import dataclass
 from enum import Enum
-from air_benchmark.tasks.tasks import BenchmarkTable
-
 
-def get_safe_name(name: str):
-    """Get RFC 1123 compatible safe name"""
-    name = name.replace('-', '_')
-    return ''.join(
-        character.lower()
-        for character in name
-        if (character.isalnum() or character == '_'))
+from air_benchmark.tasks.tasks import BenchmarkTable
 
+DEFAULT_METRIC_QA = "ndcg_at_10"
+DEFAULT_METRIC_LONG_DOC = "recall_at_10"
 
 METRIC_LIST = [
     "ndcg_at_1",
@@ -46,6 +40,15 @@ METRIC_LIST = [
 ]
 
 
+def get_safe_name(name: str):
+    """Get RFC 1123 compatible safe name"""
+    name = name.replace('-', '_')
+    return ''.join(
+        character.lower()
+        for character in name
+        if (character.isalnum() or character == '_'))
+
+
 @dataclass
 class Benchmark:
     name: str  # [domain]_[language]_[metric], task_key in the json file,
@@ -78,15 +81,3 @@ for task, domain_dict in BenchmarkTable['AIR-Bench_24.04'].items():
 
 BenchmarksQA = Enum('BenchmarksQA', qa_benchmark_dict)
 BenchmarksLongDoc = Enum('BenchmarksLongDoc', long_doc_benchmark_dict)
-
-BENCHMARK_COLS_QA = [c.col_name for c in qa_benchmark_dict.values()]
-BENCHMARK_COLS_LONG_DOC = [c.col_name for c in long_doc_benchmark_dict.values()]
-
-DOMAIN_COLS_QA = list(frozenset([c.domain for c in qa_benchmark_dict.values()]))
-LANG_COLS_QA = list(frozenset([c.lang for c in qa_benchmark_dict.values()]))
-
-DOMAIN_COLS_LONG_DOC = list(frozenset([c.domain for c in long_doc_benchmark_dict.values()]))
-LANG_COLS_LONG_DOC = list(frozenset([c.lang for c in long_doc_benchmark_dict.values()]))
-
-DEFAULT_METRIC_QA = "ndcg_at_10"
-DEFAULT_METRIC_LONG_DOC = "recall_at_10"

src/display/gradio_formatting.py

@@ -64,7 +64,7 @@ def get_domain_dropdown(domain_list, default_domains):
 def get_language_dropdown(language_list, default_languages):
     return gr.Dropdown(
         choices=language_list,
-        value=language_list,
+        value=default_languages,
         label="Select the languages",
         multiselect=True,
         interactive=True

src/display/utils.py

@@ -57,7 +57,7 @@ def get_default_auto_eval_column_dict():
 
 def make_autoevalcolumn(cls_name="BenchmarksQA", benchmarks=BenchmarksQA):
     auto_eval_column_dict = get_default_auto_eval_column_dict()
-    ## Leaderboard columns
+    # Leaderboard columns
     for benchmark in benchmarks:
         auto_eval_column_dict.append(
             [benchmark.name, ColumnContent, ColumnContent(benchmark.value.col_name, "number", True)]

src/{read_evals.py → loaders.py}

@@ -3,23 +3,18 @@ from typing import List
 
 import pandas as pd
 
-from src.benchmarks import BenchmarksQA, BenchmarksLongDoc
-from src.display.utils import COLS_QA, COLS_LONG_DOC
-from src.display.column_names import COL_NAME_AVG, COL_NAME_RANK, COL_NAME_REVISION, COL_NAME_IS_ANONYMOUS
+from src.benchmarks import DEFAULT_METRIC_QA, DEFAULT_METRIC_LONG_DOC
+from src.display.column_names import COL_NAME_REVISION, COL_NAME_IS_ANONYMOUS, \
+    COL_NAME_TIMESTAMP
 
-from src.models import FullEvalResult
+from src.models import FullEvalResult, LeaderboardDataStore
+from src.utils import get_default_cols, get_leaderboard_df
 
 pd.options.mode.copy_on_write = True
 
 
-def calculate_mean(row):
-    if pd.isna(row).any():
-        return -1
-    else:
-        return row.mean()
 
-
-def get_raw_eval_results(results_path: str) -> List[FullEvalResult]:
+def load_raw_eval_results(results_path: str) -> List[FullEvalResult]:
     """
     Load the evaluation results from a json file
     """
@@ -58,41 +53,44 @@ def get_raw_eval_results(results_path: str) -> List[FullEvalResult]:
     return results
 
 
-def get_leaderboard_df(raw_data: List[FullEvalResult], task: str, metric: str) -> pd.DataFrame:
-    """
-    Creates a dataframe from all the individual experiment results
-    """
-    cols = [COL_NAME_IS_ANONYMOUS, ]
-    if task == "qa":
-        cols += COLS_QA
-        benchmark_cols = [t.value.col_name for t in BenchmarksQA]
-    elif task == "long-doc":
-        cols += COLS_LONG_DOC
-        benchmark_cols = [t.value.col_name for t in BenchmarksLongDoc]
-    else:
-        raise NotImplemented
-    all_data_json = []
-    for v in raw_data:
-        all_data_json += v.to_dict(task=task, metric=metric)
-    df = pd.DataFrame.from_records(all_data_json)
-    # print(f'dataframe created: {df.shape}')
-
-    _benchmark_cols = frozenset(benchmark_cols).intersection(frozenset(df.columns.to_list()))
-
-    # calculate the average score for selected benchmarks
-    df[COL_NAME_AVG] = df[list(_benchmark_cols)].apply(calculate_mean, axis=1).round(decimals=2)
-    df.sort_values(by=[COL_NAME_AVG], ascending=False, inplace=True)
-    df.reset_index(inplace=True, drop=True)
-
-    _cols = frozenset(cols).intersection(frozenset(df.columns.to_list()))
-    df = df[_cols].round(decimals=2)
-
-    # filter out if any of the benchmarks have not been produced
-    df[COL_NAME_RANK] = df[COL_NAME_AVG].rank(ascending=False, method="min")
-
-    # shorten the revision
-    df[COL_NAME_REVISION] = df[COL_NAME_REVISION].str[:6]
-    
-    # # replace "0" with "-" for average score
-    # df[COL_NAME_AVG] = df[COL_NAME_AVG].replace(0, "-")
-    return df
+def load_leaderboard_datastore(file_path) -> LeaderboardDataStore:
+    lb_data_store = LeaderboardDataStore(None, None, None, None, None, None, None, None)
+    lb_data_store.raw_data = load_raw_eval_results(file_path)
+    print(f'raw data: {len(lb_data_store.raw_data)}')
+
+    lb_data_store.raw_df_qa = get_leaderboard_df(
+        lb_data_store.raw_data, task='qa', metric=DEFAULT_METRIC_QA)
+    lb_data_store.leaderboard_df_qa = lb_data_store.raw_df_qa.copy()
+    # leaderboard_df_qa = leaderboard_df_qa[has_no_nan_values(df, _benchmark_cols)]
+    print(f'QA data loaded: {lb_data_store.raw_df_qa.shape}')
+    shown_columns_qa, types_qa = get_default_cols(
+        'qa', lb_data_store.leaderboard_df_qa.columns, add_fix_cols=True)
+    lb_data_store.types_qa = types_qa
+    lb_data_store.leaderboard_df_qa = \
+    lb_data_store.leaderboard_df_qa[~lb_data_store.leaderboard_df_qa[COL_NAME_IS_ANONYMOUS]][shown_columns_qa]
+    lb_data_store.leaderboard_df_qa.drop([COL_NAME_REVISION, COL_NAME_TIMESTAMP], axis=1, inplace=True)
+
+    lb_data_store.raw_df_long_doc = get_leaderboard_df(
+        lb_data_store.raw_data, task='long-doc', metric=DEFAULT_METRIC_LONG_DOC)
+    print(f'Long-Doc data loaded: {len(lb_data_store.raw_df_long_doc)}')
+    lb_data_store.leaderboard_df_long_doc = lb_data_store.raw_df_long_doc.copy()
+    shown_columns_long_doc, types_long_doc = get_default_cols(
+        'long-doc', lb_data_store.leaderboard_df_long_doc.columns, add_fix_cols=True)
+    lb_data_store.types_long_doc = types_long_doc
+    lb_data_store.leaderboard_df_long_doc = \
+    lb_data_store.leaderboard_df_long_doc[~lb_data_store.leaderboard_df_long_doc[COL_NAME_IS_ANONYMOUS]][
+        shown_columns_long_doc]
+    lb_data_store.leaderboard_df_long_doc.drop([COL_NAME_REVISION, COL_NAME_TIMESTAMP], axis=1, inplace=True)
+
+    lb_data_store.reranking_models = sorted(
+        list(frozenset([eval_result.reranking_model for eval_result in lb_data_store.raw_data])))
+    return lb_data_store
+
+
+def load_eval_results(file_path: str):
+    output = {}
+    versions = ("AIR-Bench_24.04",)
+    for version in versions:
+        fn = f"{file_path}/{version}"
+        output[version] = load_leaderboard_datastore(fn)
+    return output

src/utils.py

@@ -6,18 +6,23 @@ from typing import List
 
 import pandas as pd
 
-from src.benchmarks import BENCHMARK_COLS_QA, BENCHMARK_COLS_LONG_DOC, BenchmarksQA, BenchmarksLongDoc
+from src.benchmarks import qa_benchmark_dict, long_doc_benchmark_dict, BenchmarksQA, BenchmarksLongDoc
 from src.display.formatting import styled_message, styled_error
 from src.display.utils import COLS_QA, TYPES_QA, COLS_LONG_DOC, TYPES_LONG_DOC, get_default_auto_eval_column_dict
 from src.display.column_names import COL_NAME_AVG, COL_NAME_RETRIEVAL_MODEL, COL_NAME_RERANKING_MODEL, COL_NAME_RANK, \
     COL_NAME_REVISION, COL_NAME_TIMESTAMP, COL_NAME_IS_ANONYMOUS
 from src.envs import API, SEARCH_RESULTS_REPO, LATEST_BENCHMARK_VERSION
-from src.read_evals import get_leaderboard_df, calculate_mean
 from src.models import FullEvalResult
 
 import re
 
 
+def calculate_mean(row):
+    if pd.isna(row).any():
+        return -1
+    else:
+        return row.mean()
+
 def remove_html(input_str):
     # Regular expression for finding HTML tags
     clean = re.sub(r'<.*?>', '', input_str)
@@ -63,11 +68,11 @@ def get_default_cols(task: str, columns: list=[], add_fix_cols: bool=True) -> li
     if task == "qa":
         cols_list = COLS_QA
         types_list = TYPES_QA
-        benchmark_list = BENCHMARK_COLS_QA
+        benchmark_list = [c.col_name for c in qa_benchmark_dict.values()]
     elif task == "long-doc":
         cols_list = COLS_LONG_DOC
         types_list = TYPES_LONG_DOC
-        benchmark_list = BENCHMARK_COLS_LONG_DOC
+        benchmark_list = [c.col_name for c in long_doc_benchmark_dict.values()]
     else:
         raise NotImplemented
     for col_name, col_type in zip(cols_list, types_list):
@@ -318,3 +323,43 @@ def submit_results(
 def reset_rank(df):
     df[COL_NAME_RANK] = df[COL_NAME_AVG].rank(ascending=False, method="min")
     return df
+
+
+def get_leaderboard_df(raw_data: List[FullEvalResult], task: str, metric: str) -> pd.DataFrame:
+    """
+    Creates a dataframe from all the individual experiment results
+    """
+    cols = [COL_NAME_IS_ANONYMOUS, ]
+    if task == "qa":
+        cols += COLS_QA
+        benchmark_cols = [t.value.col_name for t in BenchmarksQA]
+    elif task == "long-doc":
+        cols += COLS_LONG_DOC
+        benchmark_cols = [t.value.col_name for t in BenchmarksLongDoc]
+    else:
+        raise NotImplemented
+    all_data_json = []
+    for v in raw_data:
+        all_data_json += v.to_dict(task=task, metric=metric)
+    df = pd.DataFrame.from_records(all_data_json)
+    # print(f'dataframe created: {df.shape}')
+
+    _benchmark_cols = frozenset(benchmark_cols).intersection(frozenset(df.columns.to_list()))
+
+    # calculate the average score for selected benchmarks
+    df[COL_NAME_AVG] = df[list(_benchmark_cols)].apply(calculate_mean, axis=1).round(decimals=2)
+    df.sort_values(by=[COL_NAME_AVG], ascending=False, inplace=True)
+    df.reset_index(inplace=True, drop=True)
+
+    _cols = frozenset(cols).intersection(frozenset(df.columns.to_list()))
+    df = df[_cols].round(decimals=2)
+
+    # filter out if any of the benchmarks have not been produced
+    df[COL_NAME_RANK] = df[COL_NAME_AVG].rank(ascending=False, method="min")
+
+    # shorten the revision
+    df[COL_NAME_REVISION] = df[COL_NAME_REVISION].str[:6]
+
+    # # replace "0" with "-" for average score
+    # df[COL_NAME_AVG] = df[COL_NAME_AVG].replace(0, "-")
+    return df

tests/src/test_read_evals.py

@@ -1,6 +1,7 @@
 from pathlib import Path
 
-from src.read_evals import get_raw_eval_results, get_leaderboard_df
+from src.read_evals import load_raw_eval_results
+from src.utils import get_leaderboard_df
 from src.models import FullEvalResult
 
 cur_fp = Path(__file__)
@@ -30,7 +31,7 @@ def test_to_dict():
 
 def test_get_raw_eval_results():
     results_path = cur_fp.parents[2] / "toydata" / "eval_results" / "AIR-Bench_24.04"
-    results = get_raw_eval_results(results_path)
+    results = load_raw_eval_results(results_path)
     # only load the latest results
     assert len(results) == 4
     assert results[0].eval_name == "bge-base-en-v1.5_NoReranker"
@@ -41,7 +42,7 @@ def test_get_raw_eval_results():
 
 def test_get_leaderboard_df():
     results_path = cur_fp.parents[2] / "toydata" / "eval_results" / "AIR-Bench_24.04"
-    raw_data = get_raw_eval_results(results_path)
+    raw_data = load_raw_eval_results(results_path)
     df = get_leaderboard_df(raw_data, 'qa', 'ndcg_at_10')
     assert df.shape[0] == 4
     # the results contain only one embedding model
@@ -56,7 +57,7 @@ def test_get_leaderboard_df():
 
 def test_get_leaderboard_df_long_doc():
     results_path = cur_fp.parents[2] / "toydata" / "test_results"
-    raw_data = get_raw_eval_results(results_path)
+    raw_data = load_raw_eval_results(results_path)
     df = get_leaderboard_df(raw_data, 'long-doc', 'ndcg_at_1')
     assert df.shape[0] == 2
     # the results contain only one embedding model