test: add unit tests for models

src/models.py

@@ -109,18 +109,19 @@ class FullEvalResult:
                 continue
             if eval_result.task != task:
                 continue
-            results[eval_result.eval_name]["eval_name"] = eval_result.eval_name
-            results[eval_result.eval_name][COL_NAME_RETRIEVAL_MODEL] = make_clickable_model(
+            eval_name = eval_result.eval_name
+            results[eval_name]["eval_name"] = eval_name
+            results[eval_name][COL_NAME_RETRIEVAL_MODEL] = make_clickable_model(
                 self.retrieval_model, self.retrieval_model_link
             )
-            results[eval_result.eval_name][COL_NAME_RERANKING_MODEL] = make_clickable_model(
+            results[eval_name][COL_NAME_RERANKING_MODEL] = make_clickable_model(
                 self.reranking_model, self.reranking_model_link
             )
-            results[eval_result.eval_name][COL_NAME_RETRIEVAL_MODEL_LINK] = self.retrieval_model_link
-            results[eval_result.eval_name][COL_NAME_RERANKING_MODEL_LINK] = self.reranking_model_link
-            results[eval_result.eval_name][COL_NAME_REVISION] = self.revision
-            results[eval_result.eval_name][COL_NAME_TIMESTAMP] = self.timestamp
-            results[eval_result.eval_name][COL_NAME_IS_ANONYMOUS] = self.is_anonymous
+            results[eval_name][COL_NAME_RETRIEVAL_MODEL_LINK] = self.retrieval_model_link
+            results[eval_name][COL_NAME_RERANKING_MODEL_LINK] = self.reranking_model_link
+            results[eval_name][COL_NAME_REVISION] = self.revision
+            results[eval_name][COL_NAME_TIMESTAMP] = self.timestamp
+            results[eval_name][COL_NAME_IS_ANONYMOUS] = self.is_anonymous
 
             for result in eval_result.results:
                 # add result for each domain, language, and dataset
@@ -132,7 +133,7 @@ class FullEvalResult:
                     benchmark_name = f"{domain}_{lang}"
                 else:
                     benchmark_name = f"{domain}_{lang}_{dataset}"
-                results[eval_result.eval_name][get_safe_name(benchmark_name)] = value
+                results[eval_name][get_safe_name(benchmark_name)] = value
         return [v for v in results.values()]
 
 

tests/src/test_models.py

@@ -6,6 +6,23 @@ from src.models import EvalResult, FullEvalResult
 cur_fp = Path(__file__)
 
 
+# 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 |
+NUM_QA_BENCHMARKS_24_05 = 53
+NUM_DOC_BENCHMARKS_24_05 = 11
+NUM_QA_BENCHMARKS_24_04 = 13
+NUM_DOC_BENCHMARKS_24_04 = 15
+
 def test_eval_result():
     eval_result = EvalResult(
         eval_name="eval_name",
@@ -41,9 +58,21 @@ def test_full_eval_result_init_from_json_file(file_path):
     assert len(full_eval_result.results) == 70
 
 
-def test_full_eval_result_to_dict():
-    json_fp = cur_fp.parents[1] / "toydata/eval_results/" / "AIR-Bench_24.05/bge-m3/NoReranker/results.json"
+@pytest.mark.parametrize(
+    'file_path, task, expected_num_results',
+    [
+        ("AIR-Bench_24.04/bge-m3/jina-reranker-v2-base-multilingual/results.json", "qa", NUM_QA_BENCHMARKS_24_04),
+        ("AIR-Bench_24.04/bge-m3/jina-reranker-v2-base-multilingual/results.json", "long-doc", NUM_DOC_BENCHMARKS_24_04),
+        ("AIR-Bench_24.05/bge-m3/NoReranker/results.json", "qa", NUM_QA_BENCHMARKS_24_05),
+        ("AIR-Bench_24.05/bge-m3/NoReranker/results.json", "long-doc", NUM_DOC_BENCHMARKS_24_05),
+    ])
+def test_full_eval_result_to_dict(file_path, task, expected_num_results):
+    json_fp = cur_fp.parents[1] / "toydata/eval_results/" / file_path
     full_eval_result = FullEvalResult.init_from_json_file(json_fp)
-    result_dict_list = full_eval_result.to_dict()
+    result_dict_list = full_eval_result.to_dict(task)
     assert len(result_dict_list) == 1
-    print(len(result_dict_list[0]))
+    result = result_dict_list[0]
+    attr_list = frozenset([
+        'eval_name', 'Retrieval Method', 'Reranking Model', 'Retrieval Model LINK', 'Reranking Model LINK', 'Revision', 'Submission Date', 'Anonymous Submission'])
+    result_cols = list(result.keys())
+    assert len(result_cols) == (expected_num_results + len(attr_list))