refactor: refactor the benchmarks

app.py

@@ -6,8 +6,8 @@ from src.about import (
     TITLE
 )
 from src.benchmarks import (
-    qa_benchmark_dict,
-    long_doc_benchmark_dict
+    BenchmarksQA,
+    BenchmarksLongDoc
 )
 from src.display.css_html_js import custom_css
 from src.envs import (
@@ -76,11 +76,11 @@ 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_QA = list(frozenset([c.value.domain for c in list(BenchmarksQA)]))
+LANG_COLS_QA = list(frozenset([c.value.lang for c in list(BenchmarksQA)]))
 
-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()]))
+DOMAIN_COLS_LONG_DOC = list(frozenset([c.value.domain for c in list(BenchmarksLongDoc)]))
+LANG_COLS_LONG_DOC = list(frozenset([c.value.lang for c in list(BenchmarksLongDoc)]))
 
 demo = gr.Blocks(css=custom_css)
 

src/benchmarks.py

@@ -25,25 +25,34 @@ class Benchmark:
     task: str
 
 
-qa_benchmark_dict = {}
-long_doc_benchmark_dict = {}
-for task, domain_dict in BenchmarkTable['AIR-Bench_24.04'].items():
-    for domain, lang_dict in domain_dict.items():
-        for lang, dataset_list in lang_dict.items():
-            if task == "qa":
-                benchmark_name = f"{domain}_{lang}"
-                benchmark_name = get_safe_name(benchmark_name)
-                col_name = benchmark_name
-                for metric in dataset_list:
-                    qa_benchmark_dict[benchmark_name] = Benchmark(benchmark_name, metric, col_name, domain, lang, task)
-            elif task == "long-doc":
-                for dataset in dataset_list:
-                    benchmark_name = f"{domain}_{lang}_{dataset}"
+# create a function return an enum class containing all the benchmarks
+def get_benchmarks_enum(benchmark_version):
+    qa_benchmark_dict = {}
+    long_doc_benchmark_dict = {}
+    for task, domain_dict in BenchmarkTable[benchmark_version].items():
+        for domain, lang_dict in domain_dict.items():
+            for lang, dataset_list in lang_dict.items():
+                if task == "qa":
+                    benchmark_name = f"{domain}_{lang}"
                     benchmark_name = get_safe_name(benchmark_name)
                     col_name = benchmark_name
-                    for metric in METRIC_LIST:
-                        long_doc_benchmark_dict[benchmark_name] = Benchmark(benchmark_name, metric, col_name, domain,
+                    for metric in dataset_list:
+                        qa_benchmark_dict[benchmark_name] = \
+                            Benchmark(
+                                benchmark_name, metric, col_name, domain, lang, task)
+                elif task == "long-doc":
+                    for dataset in dataset_list:
+                        benchmark_name = f"{domain}_{lang}_{dataset}"
+                        benchmark_name = get_safe_name(benchmark_name)
+                        col_name = benchmark_name
+                        for metric in METRIC_LIST:
+                            long_doc_benchmark_dict[benchmark_name] = \
+                                Benchmark(
+                                    benchmark_name, metric, col_name, domain,
                                                                             lang, task)
+    return qa_benchmark_dict, long_doc_benchmark_dict
 
-BenchmarksQA = Enum('BenchmarksQA', qa_benchmark_dict)
-BenchmarksLongDoc = Enum('BenchmarksLongDoc', long_doc_benchmark_dict)
+_qa_benchmark_dict, _long_doc_benchmark_dict = get_benchmarks_enum('AIR-Bench_24.04')
+
+BenchmarksQA = Enum('BenchmarksQA', _qa_benchmark_dict)
+BenchmarksLongDoc = Enum('BenchmarksLongDoc', _long_doc_benchmark_dict)

src/utils.py

@@ -6,7 +6,7 @@ from typing import List
 
 import pandas as pd
 
-from src.benchmarks import qa_benchmark_dict, long_doc_benchmark_dict, BenchmarksQA, BenchmarksLongDoc
+from src.benchmarks import BenchmarksQA, BenchmarksLongDoc
 from src.display.formatting import styled_message, styled_error
 from src.display.columns 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, COLS_QA, TYPES_QA, COLS_LONG_DOC, TYPES_LONG_DOC, \
@@ -68,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 = [c.col_name for c in qa_benchmark_dict.values()]
+        benchmark_list = [c.value.col_name for c in list(BenchmarksQA)]
     elif task == "long-doc":
         cols_list = COLS_LONG_DOC
         types_list = TYPES_LONG_DOC
-        benchmark_list = [c.col_name for c in long_doc_benchmark_dict.values()]
+        benchmark_list = [c.value.col_name for c in list(BenchmarksLongDoc)]
     else:
         raise NotImplemented
     for col_name, col_type in zip(cols_list, types_list):

tests/src/test_benchmarks.py

@@ -3,6 +3,8 @@ from src.benchmarks import BenchmarksQA, BenchmarksLongDoc
 
 def test_qabenchmarks():
     print(list(BenchmarksQA))
+    for benchmark in list(BenchmarksQA):
+        print(benchmark.name, benchmark.metric, benchmark.col_name, benchmark.domain, benchmark.lang, benchmark.task)
 
 
 def test_longdocbenchmarks():