refactor: use enum class for the task type

app.py

@@ -35,6 +35,7 @@ from src.envs import (
     TOKEN,
 )
 from src.loaders import load_eval_results
+from src.models import TaskType
 from src.utils import remove_html, reset_rank, set_listeners, submit_results, update_metric, upload_file
 
 
@@ -75,7 +76,7 @@ def update_qa_metric(
     global datastore
     return update_metric(
         datastore,
-        "qa",
+        TaskType.qa,
         metric,
         domains,
         langs,
@@ -98,7 +99,7 @@ def update_doc_metric(
     global datastore
     return update_metric(
         datastore,
-        "long-doc",
+        TaskType.long_doc,
         metric,
         domains,
         langs,
@@ -181,7 +182,7 @@ with demo:
                         )
 
                         set_listeners(
-                            "qa",
+                            TaskType.qa,
                             qa_df_elem_ret_rerank,
                             qa_df_elem_ret_rerank_hidden,
                             search_bar,
@@ -224,7 +225,7 @@ with demo:
                         )
 
                         set_listeners(
-                            "qa",
+                            TaskType.qa,
                             qa_df_elem_ret,
                             qa_df_elem_ret_hidden,
                             search_bar_ret,
@@ -281,7 +282,7 @@ with demo:
                         )
 
                         set_listeners(
-                            "qa",
+                            TaskType.qa,
                             qa_df_elem_rerank,
                             qa_df_elem_rerank_hidden,
                             qa_search_bar_rerank,
@@ -348,7 +349,7 @@ with demo:
                         )
 
                         set_listeners(
-                            "long-doc",
+                            TaskType.long_doc,
                             doc_df_elem_ret_rerank,
                             doc_df_elem_ret_rerank_hidden,
                             search_bar,
@@ -405,7 +406,7 @@ with demo:
                         )
 
                         set_listeners(
-                            "long-doc",
+                            TaskType.long_doc,
                             doc_df_elem_ret,
                             doc_df_elem_ret_hidden,
                             search_bar_ret,
@@ -462,7 +463,7 @@ with demo:
                         )
 
                         set_listeners(
-                            "long-doc",
+                            TaskType.long_doc,
                             doc_df_elem_rerank,
                             doc_df_elem_rerank_hidden,
                             doc_search_bar_rerank,

src/benchmarks.py

@@ -4,6 +4,7 @@ from enum import Enum
 from air_benchmark.tasks.tasks import BenchmarkTable
 
 from src.envs import BENCHMARK_VERSION_LIST, METRIC_LIST
+from src.models import TaskType
 
 
 def get_safe_name(name: str):
@@ -23,11 +24,11 @@ class Benchmark:
 
 
 # create a function return an enum class containing all the benchmarks
-def get_benchmarks_enum(benchmark_version, task_type):
+def get_benchmarks_enum(benchmark_version: str, task_type: TaskType):
     benchmark_dict = {}
-    if task_type == "qa":
+    if task_type == TaskType.qa:
         for task, domain_dict in BenchmarkTable[benchmark_version].items():
-            if task != task_type:
+            if task != task_type.value:
                 continue
             for domain, lang_dict in domain_dict.items():
                 for lang, dataset_list in lang_dict.items():
@@ -39,9 +40,9 @@ def get_benchmarks_enum(benchmark_version, task_type):
                         benchmark_dict[benchmark_name] = Benchmark(
                             benchmark_name, metric, col_name, domain, lang, task
                         )
-    elif task_type == "long-doc":
+    elif task_type == TaskType.long_doc:
         for task, domain_dict in BenchmarkTable[benchmark_version].items():
-            if task != task_type:
+            if task != task_type.value:
                 continue
             for domain, lang_dict in domain_dict.items():
                 for lang, dataset_list in lang_dict.items():
@@ -62,14 +63,14 @@ qa_benchmark_dict = {}
 for version in BENCHMARK_VERSION_LIST:
     safe_version_name = get_safe_name(version)[-4:]
     qa_benchmark_dict[safe_version_name] = Enum(
-        f"QABenchmarks_{safe_version_name}", get_benchmarks_enum(version, "qa")
+        f"QABenchmarks_{safe_version_name}", get_benchmarks_enum(version, TaskType.qa)
     )
 
 long_doc_benchmark_dict = {}
 for version in BENCHMARK_VERSION_LIST:
     safe_version_name = get_safe_name(version)[-4:]
     long_doc_benchmark_dict[safe_version_name] = Enum(
-        f"LongDocBenchmarks_{safe_version_name}", get_benchmarks_enum(version, "long-doc")
+        f"LongDocBenchmarks_{safe_version_name}", get_benchmarks_enum(version, TaskType.long_doc)
     )
 
 

src/loaders.py

@@ -11,7 +11,7 @@ from src.envs import (
     DEFAULT_METRIC_LONG_DOC,
     DEFAULT_METRIC_QA,
 )
-from src.models import FullEvalResult, LeaderboardDataStore
+from src.models import FullEvalResult, LeaderboardDataStore, TaskType
 from src.utils import get_default_cols, get_leaderboard_df
 
 pd.options.mode.copy_on_write = True
@@ -64,34 +64,27 @@ def get_safe_name(name: str):
 
 def load_leaderboard_datastore(file_path, version) -> LeaderboardDataStore:
     slug = get_safe_name(version)[-4:]
-    lb_data_store = LeaderboardDataStore(version, slug, 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.qa_raw_df = get_leaderboard_df(lb_data_store, task="qa", metric=DEFAULT_METRIC_QA)
-    print(f"QA data loaded: {lb_data_store.qa_raw_df.shape}")
-    lb_data_store.qa_fmt_df = lb_data_store.qa_raw_df.copy()
-    shown_columns_qa, types_qa = get_default_cols("qa", lb_data_store.slug, add_fix_cols=True)
-    lb_data_store.qa_types = types_qa
-    lb_data_store.qa_fmt_df = lb_data_store.qa_fmt_df[
-        ~lb_data_store.qa_fmt_df[COL_NAME_IS_ANONYMOUS]
-    ][shown_columns_qa]
-    lb_data_store.qa_fmt_df.drop([COL_NAME_REVISION, COL_NAME_TIMESTAMP], axis=1, inplace=True)
-
-    lb_data_store.doc_raw_df = get_leaderboard_df(lb_data_store, task="long-doc", metric=DEFAULT_METRIC_LONG_DOC)
-    print(f"Long-Doc data loaded: {len(lb_data_store.doc_raw_df)}")
-    lb_data_store.doc_fmt_df = lb_data_store.doc_raw_df.copy()
-    shown_columns_long_doc, types_long_doc = get_default_cols("long-doc", lb_data_store.slug, add_fix_cols=True)
-    lb_data_store.doc_types = types_long_doc
-    lb_data_store.doc_fmt_df = lb_data_store.doc_fmt_df[
-        ~lb_data_store.doc_fmt_df[COL_NAME_IS_ANONYMOUS]
-    ][shown_columns_long_doc]
-    lb_data_store.doc_fmt_df.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
+    datastore = LeaderboardDataStore(version, slug, None, None, None, None, None, None, None, None)
+    datastore.raw_data = load_raw_eval_results(file_path)
+    print(f"raw data: {len(datastore.raw_data)}")
+
+    datastore.qa_raw_df = get_leaderboard_df(datastore, TaskType.qa, DEFAULT_METRIC_QA)
+    print(f"QA data loaded: {datastore.qa_raw_df.shape}")
+    datastore.qa_fmt_df = datastore.qa_raw_df.copy()
+    qa_cols, datastore.qa_types = get_default_cols(TaskType.qa, datastore.slug, add_fix_cols=True)
+    datastore.qa_fmt_df = datastore.qa_fmt_df[~datastore.qa_fmt_df[COL_NAME_IS_ANONYMOUS]][qa_cols]
+    datastore.qa_fmt_df.drop([COL_NAME_REVISION, COL_NAME_TIMESTAMP], axis=1, inplace=True)
+
+    datastore.doc_raw_df = get_leaderboard_df(datastore, TaskType.long_doc, DEFAULT_METRIC_LONG_DOC)
+    print(f"Long-Doc data loaded: {len(datastore.doc_raw_df)}")
+    datastore.doc_fmt_df = datastore.doc_raw_df.copy()
+    doc_cols, datastore.doc_types = get_default_cols(TaskType.long_doc, datastore.slug, add_fix_cols=True)
+    datastore.doc_fmt_df = datastore.doc_fmt_df[~datastore.doc_fmt_df[COL_NAME_IS_ANONYMOUS]][doc_cols]
+    datastore.doc_fmt_df.drop([COL_NAME_REVISION, COL_NAME_TIMESTAMP], axis=1, inplace=True)
+
+    datastore.reranking_models = \
+        sorted(list(frozenset([eval_result.reranking_model for eval_result in datastore.raw_data])))
+    return datastore
 
 
 def load_eval_results(file_path: str) -> Dict[str, LeaderboardDataStore]:

src/models.py

@@ -1,11 +1,12 @@
 import json
+from enum import Enum
+
 from collections import defaultdict
 from dataclasses import dataclass
 from typing import List, Optional
 
 import pandas as pd
 
-from src.benchmarks import get_safe_name
 from src.display.formatting import make_clickable_model
 from src.envs import (
     COL_NAME_IS_ANONYMOUS,
@@ -17,6 +18,10 @@ from src.envs import (
     COL_NAME_TIMESTAMP,
 )
 
+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 EvalResult:
@@ -147,4 +152,10 @@ class LeaderboardDataStore:
     doc_fmt_df: Optional[pd.DataFrame]
     reranking_models: Optional[list]
     qa_types: Optional[list]
-    doc_types: Optional[list]
+    doc_types: Optional[list]
+
+
+# Define an enum class with the name `TaskType`. There are two types of tasks, `qa` and `long-doc`.
+class TaskType(Enum):
+    qa = "qa"
+    long_doc = "long-doc"

src/utils.py

@@ -6,6 +6,7 @@ from pathlib import Path
 
 import pandas as pd
 
+from src.models import TaskType
 from src.benchmarks import LongDocBenchmarks, QABenchmarks
 from src.display.columns import get_default_col_names_and_types, get_fixed_col_names_and_types
 from src.display.formatting import styled_error, styled_message
@@ -69,12 +70,12 @@ def search_table(df: pd.DataFrame, query: str) -> pd.DataFrame:
     return df[(df[COL_NAME_RETRIEVAL_MODEL].str.contains(query, case=False))]
 
 
-def get_default_cols(task: str, version_slug, add_fix_cols: bool = True) -> tuple:
+def get_default_cols(task: TaskType, version_slug, add_fix_cols: bool = True) -> tuple:
     cols = []
     types = []
-    if task == "qa":
+    if task == TaskType.qa:
         benchmarks = QABenchmarks[version_slug]
-    elif task == "long-doc":
+    elif task == TaskType.long_doc:
         benchmarks = LongDocBenchmarks[version_slug]
     else:
         raise NotImplementedError
@@ -85,7 +86,6 @@ def get_default_cols(task: str, version_slug, add_fix_cols: bool = True) -> tupl
             continue
         cols.append(col_name)
         types.append(col_type)
-
     if add_fix_cols:
         _cols = []
         _types = []
@@ -104,16 +104,16 @@ def select_columns(
     df: pd.DataFrame,
     domain_query: list,
     language_query: list,
-    task: str = "qa",
+    task: TaskType = TaskType.qa,
     reset_ranking: bool = True,
     version_slug: str = None,
 ) -> pd.DataFrame:
     cols, _ = get_default_cols(task=task, version_slug=version_slug, add_fix_cols=False)
     selected_cols = []
     for c in cols:
-        if task == "qa":
+        if task == TaskType.qa:
             eval_col = QABenchmarks[version_slug].value[c].value
-        elif task == "long-doc":
+        elif task == TaskType.long_doc:
             eval_col = LongDocBenchmarks[version_slug].value[c].value
         else:
             raise NotImplementedError
@@ -141,10 +141,10 @@ def get_safe_name(name: str):
     return "".join(character.lower() for character in name if (character.isalnum() or character == "_"))
 
 
-def _update_table(
-    task: str,
+def _update_df_elem(
+    task: TaskType,
     version: str,
-    hidden_df: pd.DataFrame,
+    source_df: pd.DataFrame,
     domains: list,
     langs: list,
     reranking_query: list,
@@ -154,7 +154,7 @@ def _update_table(
     show_revision_and_timestamp: bool = False,
 ):
     version_slug = get_safe_name(version)[-4:]
-    filtered_df = hidden_df.copy()
+    filtered_df = source_df.copy()
     if not show_anonymous:
         filtered_df = filtered_df[~filtered_df[COL_NAME_IS_ANONYMOUS]]
     filtered_df = filter_models(filtered_df, reranking_query)
@@ -165,7 +165,7 @@ def _update_table(
     return filtered_df
 
 
-def update_table_long_doc(
+def update_doc_df_elem(
     version: str,
     hidden_df: pd.DataFrame,
     domains: list,
@@ -176,8 +176,8 @@ def update_table_long_doc(
     show_revision_and_timestamp: bool = False,
     reset_ranking: bool = True,
 ):
-    return _update_table(
-        "long-doc",
+    return _update_df_elem(
+        TaskType.long_doc,
         version,
         hidden_df,
         domains,
@@ -192,7 +192,7 @@ def update_table_long_doc(
 
 def update_metric(
     datastore,
-    task: str,
+    task: TaskType,
     metric: str,
     domains: list,
     langs: list,
@@ -201,33 +201,24 @@ def update_metric(
     show_anonymous: bool = False,
     show_revision_and_timestamp: bool = False,
 ) -> pd.DataFrame:
-    # raw_data = datastore.raw_data
-    if task == "qa":
-        leaderboard_df = get_leaderboard_df(datastore, task=task, metric=metric)
-        version = datastore.version
-        return update_table(
-            version,
-            leaderboard_df,
-            domains,
-            langs,
-            reranking_model,
-            query,
-            show_anonymous,
-            show_revision_and_timestamp,
-        )
-    elif task == "long-doc":
-        leaderboard_df = get_leaderboard_df(datastore, task=task, metric=metric)
-        version = datastore.version
-        return update_table_long_doc(
-            version,
-            leaderboard_df,
-            domains,
-            langs,
-            reranking_model,
-            query,
-            show_anonymous,
-            show_revision_and_timestamp,
-        )
+    if task == TaskType.qa:
+        update_func = update_qa_df_elem
+    elif task == TaskType.long_doc:
+        update_func = update_doc_df_elem
+    else:
+        raise NotImplemented
+    df_elem = get_leaderboard_df(datastore, task=task, metric=metric)
+    version = datastore.version
+    return update_func(
+        version,
+        df_elem,
+        domains,
+        langs,
+        reranking_model,
+        query,
+        show_anonymous,
+        show_revision_and_timestamp,
+    )
 
 
 def upload_file(filepath: str):
@@ -341,7 +332,7 @@ def reset_rank(df):
     return df
 
 
-def get_leaderboard_df(datastore, task: str, metric: str) -> pd.DataFrame:
+def get_leaderboard_df(datastore, task: TaskType, metric: str) -> pd.DataFrame:
     """
     Creates a dataframe from all the individual experiment results
     """
@@ -349,9 +340,9 @@ def get_leaderboard_df(datastore, task: str, metric: str) -> pd.DataFrame:
     cols = [
         COL_NAME_IS_ANONYMOUS,
     ]
-    if task == "qa":
+    if task == TaskType.qa:
         benchmarks = QABenchmarks[datastore.slug]
-    elif task == "long-doc":
+    elif task == TaskType.long_doc:
         benchmarks = LongDocBenchmarks[datastore.slug]
     else:
         raise NotImplementedError
@@ -360,7 +351,7 @@ def get_leaderboard_df(datastore, task: str, metric: str) -> pd.DataFrame:
     benchmark_cols = [t.value.col_name for t in list(benchmarks.value)]
     all_data_json = []
     for v in raw_data:
-        all_data_json += v.to_dict(task=task, metric=metric)
+        all_data_json += v.to_dict(task=task.value, metric=metric)
     df = pd.DataFrame.from_records(all_data_json)
 
     _benchmark_cols = frozenset(benchmark_cols).intersection(frozenset(df.columns.to_list()))
@@ -385,7 +376,7 @@ def get_leaderboard_df(datastore, task: str, metric: str) -> pd.DataFrame:
 
 
 def set_listeners(
-    task,
+    task: TaskType,
     target_df,
     source_df,
     search_bar,
@@ -396,10 +387,10 @@ def set_listeners(
     show_anonymous,
     show_revision_and_timestamp,
 ):
-    if task == "qa":
-        update_table_func = update_table
-    elif task == "long-doc":
-        update_table_func = update_table_long_doc
+    if task == TaskType.qa:
+        update_table_func = update_qa_df_elem
+    elif task == TaskType.long_doc:
+        update_table_func = update_doc_df_elem
     else:
         raise NotImplementedError
     selector_list = [selected_domains, selected_langs, selected_rerankings, search_bar, show_anonymous]
@@ -427,7 +418,7 @@ def set_listeners(
         )
 
 
-def update_table(
+def update_qa_df_elem(
     version: str,
     hidden_df: pd.DataFrame,
     domains: list,
@@ -438,8 +429,8 @@ def update_table(
     show_revision_and_timestamp: bool = False,
     reset_ranking: bool = True,
 ):
-    return _update_table(
-        "qa",
+    return _update_df_elem(
+        TaskType.qa,
         version,
         hidden_df,
         domains,

tests/test_utils.py

@@ -18,7 +18,7 @@ from src.utils import (
     get_iso_format_timestamp,
     search_table,
     select_columns,
-    update_table_long_doc,
+    update_doc_df_elem,
 )
 
 
@@ -90,7 +90,7 @@ def test_select_columns(toy_df):
 
 
 def test_update_table_long_doc(toy_df_long_doc):
-    df_result = update_table_long_doc(
+    df_result = update_doc_df_elem(
         toy_df_long_doc,
         [
             "law",