test: add unit tests for benchmarks

src/benchmarks.py

@@ -10,7 +10,7 @@ from src.models import TaskType, get_safe_name
 @dataclass
 class Benchmark:
     name: str  # [domain]_[language]_[metric], task_key in the json file,
-    metric: str  # ndcg_at_1 ,metric_key in the json file
+    metric: str  # metric_key in the json file
     col_name: str  # [domain]_[language], name to display in the leaderboard
     domain: str
     lang: str
@@ -18,54 +18,61 @@ class Benchmark:
 
 
 # create a function return an enum class containing all the benchmarks
-def get_benchmarks_enum(benchmark_version: str, task_type: TaskType):
+def get_qa_benchmarks_dict(version: str):
     benchmark_dict = {}
-    if task_type == TaskType.qa:
-        for task, domain_dict in BenchmarkTable[benchmark_version].items():
-            if task != task_type.value:
-                continue
-            for domain, lang_dict in domain_dict.items():
-                for lang, dataset_list in lang_dict.items():
-                    benchmark_name = get_safe_name(f"{domain}_{lang}")
+    for task, domain_dict in BenchmarkTable[version].items():
+        if task != TaskType.qa.value:
+            continue
+        for domain, lang_dict in domain_dict.items():
+            for lang, dataset_list in lang_dict.items():
+                benchmark_name = get_safe_name(f"{domain}_{lang}")
+                col_name = benchmark_name
+                for metric in dataset_list:
+                    if "test" not in dataset_list[metric]["splits"]:
+                        continue
+                    benchmark_dict[benchmark_name] = Benchmark(
+                        benchmark_name, metric, col_name, domain, lang, task
+                    )
+    return benchmark_dict
+
+
+def get_doc_benchmarks_dict(version: str):
+    benchmark_dict = {}
+    for task, domain_dict in BenchmarkTable[version].items():
+        if task != TaskType.long_doc.value:
+            continue
+        for domain, lang_dict in domain_dict.items():
+            for lang, dataset_list in lang_dict.items():
+                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 dataset_list:
-                        if "test" not in dataset_list[metric]["splits"]:
-                            continue
+                    if "test" not in dataset_list[dataset]["splits"]:
+                        continue
+                    for metric in METRIC_LIST:
                         benchmark_dict[benchmark_name] = Benchmark(
                             benchmark_name, metric, col_name, domain, lang, task
                         )
-    elif task_type == TaskType.long_doc:
-        for task, domain_dict in BenchmarkTable[benchmark_version].items():
-            if task != task_type.value:
-                continue
-            for domain, lang_dict in domain_dict.items():
-                for lang, dataset_list in lang_dict.items():
-                    for dataset in dataset_list:
-                        benchmark_name = f"{domain}_{lang}_{dataset}"
-                        benchmark_name = get_safe_name(benchmark_name)
-                        col_name = benchmark_name
-                        if "test" not in dataset_list[dataset]["splits"]:
-                            continue
-                        for metric in METRIC_LIST:
-                            benchmark_dict[benchmark_name] = Benchmark(
-                                benchmark_name, metric, col_name, domain, lang, task
-                            )
     return benchmark_dict
 
 
 _qa_benchmark_dict = {}
 for version in BENCHMARK_VERSION_LIST:
     safe_version_name = get_safe_name(version)
-    _qa_benchmark_dict[safe_version_name] = Enum(
-        f"QABenchmarks_{safe_version_name}", get_benchmarks_enum(version, TaskType.qa)
-    )
+    _qa_benchmark_dict[safe_version_name] = \
+        Enum(
+            f"QABenchmarks_{safe_version_name}",
+            get_qa_benchmarks_dict(version)
+        )
 
 _doc_benchmark_dict = {}
 for version in BENCHMARK_VERSION_LIST:
     safe_version_name = get_safe_name(version)
-    _doc_benchmark_dict[safe_version_name] = Enum(
-        f"LongDocBenchmarks_{safe_version_name}", get_benchmarks_enum(version, TaskType.long_doc)
-    )
+    _doc_benchmark_dict[safe_version_name] = \
+        Enum(
+            f"LongDocBenchmarks_{safe_version_name}",
+            get_doc_benchmarks_dict(version)
+        )
 
 
 QABenchmarks = Enum("QABenchmarks", _qa_benchmark_dict)

tests/src/test_benchmarks.py

@@ -1,15 +1,49 @@
+import pytest
+
 from src.benchmarks import LongDocBenchmarks, QABenchmarks
+from src.envs import BENCHMARK_VERSION_LIST
+
 
+# Ref: https://github.com/AIR-Bench/AIR-Bench/blob/4b27b8a8f2047a963805fcf6fb9d74be51ec440c/docs/available_tasks.md
+# 24.05
+# | Task | dev | test |
+# | ---- | --- | ---- |
+# | Long-Doc | 4 | 11 |
+# | QA | 54 | 53 |
+#
+# 24.04
+# | Task | test |
+# | ---- | ---- |
+# | Long-Doc | 15 |
+# | QA | 13 |
 
-def test_qabenchmarks():
+@pytest.mark.parametrize(
+    "num_datasets_dict",
+    [
+        {
+            "air_bench_2404": 13,
+            "air_bench_2405": 53
+        }
+    ]
+)
+def test_qa_benchmarks(num_datasets_dict):
+    assert len(QABenchmarks) == len(BENCHMARK_VERSION_LIST)
     for benchmark_list in list(QABenchmarks):
-        print(benchmark_list.name)
-        for b in list(benchmark_list.value):
-            print(b)
-    qa_benchmarks = QABenchmarks["2404"]
-    l = list(frozenset([c.value.domain for c in list(qa_benchmarks.value)]))
-    print(l)
+        version_slug = benchmark_list.name
+        assert num_datasets_dict[version_slug] == len(benchmark_list.value)
 
 
-def test_longdocbenchmarks():
-    print(list(LongDocBenchmarks))
+@pytest.mark.parametrize(
+    "num_datasets_dict",
+    [
+        {
+            "air_bench_2404": 15,
+            "air_bench_2405": 11
+        }
+    ]
+)
+def test_doc_benchmarks(num_datasets_dict):
+    assert len(LongDocBenchmarks) == len(BENCHMARK_VERSION_LIST)
+    for benchmark_list in list(LongDocBenchmarks):
+        version_slug = benchmark_list.name
+        assert num_datasets_dict[version_slug] == len(benchmark_list.value)