test: add unit tests for utils

src/utils.py

@@ -98,14 +98,7 @@ def get_default_cols(task: TaskType, version_slug, add_fix_cols: bool = True) ->
     return cols, types
 
 
-def select_columns(
-        df: pd.DataFrame,
-        domain_query: list,
-        language_query: list,
-        task: TaskType = TaskType.qa,
-        reset_ranking: bool = True,
-        version_slug: str = None,
-) -> pd.DataFrame:
+def get_selected_cols(task, version_slug, domains, languages):
     cols, _ = get_default_cols(task=task, version_slug=version_slug, add_fix_cols=False)
     selected_cols = []
     for c in cols:
@@ -115,21 +108,35 @@ def select_columns(
             eval_col = LongDocBenchmarks[version_slug].value[c].value
         else:
             raise NotImplementedError
-        if eval_col.domain not in domain_query:
+        if eval_col.domain not in domains:
             continue
-        if eval_col.lang not in language_query:
+        if eval_col.lang not in languages:
             continue
         selected_cols.append(c)
     # We use COLS to maintain sorting
+    return selected_cols
+
+
+def select_columns(
+        df: pd.DataFrame,
+        domains: list,
+        languages: list,
+        task: TaskType = TaskType.qa,
+        reset_ranking: bool = True,
+        version_slug: str = None,
+) -> pd.DataFrame:
+    selected_cols = get_selected_cols(
+        task, version_slug, domains, languages)
     fixed_cols, _ = get_fixed_col_names_and_types()
     filtered_df = df[fixed_cols + selected_cols]
     filtered_df.replace({"": pd.NA}, inplace=True)
     if reset_ranking:
-        filtered_df[COL_NAME_AVG] = filtered_df[selected_cols].apply(calculate_mean, axis=1).round(decimals=2)
-        filtered_df.sort_values(by=[COL_NAME_AVG], ascending=False, inplace=True)
+        filtered_df[COL_NAME_AVG] = \
+            filtered_df[selected_cols].apply(calculate_mean, axis=1).round(decimals=2)
+        filtered_df.sort_values(
+            by=[COL_NAME_AVG], ascending=False, inplace=True)
         filtered_df.reset_index(inplace=True, drop=True)
         filtered_df = reset_rank(filtered_df)
-
     return filtered_df
 
 

tests/src/test_utils.py

@@ -1,26 +1,157 @@
-from src.display.utils import (
-    COLS_LONG_DOC,
-    COLS_QA,
-    TYPES_LONG_DOC,
-    TYPES_QA,
-    AutoEvalColumnQA,
-    fields,
-    get_default_auto_eval_column_dict,
-)
+import pytest
+import pandas as pd
+
+from src.utils import remove_html, calculate_mean, filter_models, filter_queries, get_default_cols, select_columns, get_selected_cols
+from src.models import model_hyperlink, TaskType
+from src.columns import COL_NAME_RERANKING_MODEL, COL_NAME_RETRIEVAL_MODEL
+
+
+NUM_QA_BENCHMARKS_24_05 = 53
+NUM_DOC_BENCHMARKS_24_05 = 11
+NUM_QA_BENCHMARKS_24_04 = 13
+NUM_DOC_BENCHMARKS_24_04 = 15
+
+
+@pytest.fixture
+def toy_df():
+    return pd.DataFrame(
+        {
+            "Retrieval Method": [
+                "bge-m3",
+                "bge-m3",
+                "jina-embeddings-v2-base",
+                "jina-embeddings-v2-base"
+            ],
+            "Reranking Model": [
+                "bge-reranker-v2-m3",
+                "NoReranker",
+                "bge-reranker-v2-m3",
+                "NoReranker"
+            ],
+            "Rank 🏆": [1, 2, 3, 4],
+            "Revision": ["", "", "", ""],
+            "Submission Date": ["", "", "", ""],
+            "Average ⬆️": [0.6, 0.4, 0.3, 0.2],
+            "wiki_en": [0.8, 0.7, 0.2, 0.1],
+            "wiki_zh": [0.4, 0.1, 0.4, 0.3],
+            "news_en": [0.8, 0.7, 0.2, 0.1],
+            "news_zh": [0.4, 0.1, 0.4, 0.3],
+        }
+    )
+
+
+def test_remove_html():
+    model_name = "jina-embeddings-v3"
+    html_str = model_hyperlink(
+        "https://jina.ai", model_name)
+    output_str = remove_html(html_str)
+    assert output_str == model_name
 
 
-def test_fields():
-    for c in fields(AutoEvalColumnQA):
-        print(c)
+def test_calculate_mean():
+    valid_row = [1, 3]
+    invalid_row = [2, pd.NA]
+    df = pd.DataFrame([valid_row, invalid_row], columns=["a", "b"])
+    result = list(df.apply(calculate_mean, axis=1))
+    assert result[0] == sum(valid_row) / 2
+    assert result[1] == -1
 
 
-def test_macro_variables():
-    print(f"COLS_QA: {COLS_QA}")
-    print(f"COLS_LONG_DOC: {COLS_LONG_DOC}")
-    print(f"TYPES_QA: {TYPES_QA}")
-    print(f"TYPES_LONG_DOC: {TYPES_LONG_DOC}")
+@pytest.mark.parametrize("models, expected", [
+    (["model1", "model3"], 2),
+    (["model1", "model_missing"], 1),
+    (["model1", "model2", "model3"], 3),
+    (["model1", ], 1),
+    ([], 3),
+])
+def test_filter_models(models, expected):
+    df = pd.DataFrame(
+        {
+            COL_NAME_RERANKING_MODEL: ["model1", "model2", "model3", ],
+            "col2": [1, 2, 3],
+        }
+    )
+    output_df = filter_models(df, models)
+    assert len(output_df) == expected
+
+
+@pytest.mark.parametrize("query, expected", [
+    ("model1;model3", 2),
+    ("model1;model4", 1),
+    ("model1;model2;model3", 3),
+    ("model1", 1),
+    ("", 3),
+])
+def test_filter_queries(query, expected):
+    df = pd.DataFrame(
+        {
+            COL_NAME_RETRIEVAL_MODEL: ["model1", "model2", "model3", ],
+            COL_NAME_RERANKING_MODEL: ["model4", "model5", "model6", ],
+        }
+    )
+    output_df = filter_queries(query, df)
+    assert len(output_df) == expected
+
+
+@pytest.mark.parametrize(
+    "task_type, slug, expected",
+    [
+        (TaskType.qa, "air_bench_2404", NUM_QA_BENCHMARKS_24_04),
+        (TaskType.long_doc, "air_bench_2404", NUM_DOC_BENCHMARKS_24_04),
+        (TaskType.qa, "air_bench_2405", NUM_QA_BENCHMARKS_24_05),
+        (TaskType.long_doc, "air_bench_2405", NUM_DOC_BENCHMARKS_24_05),
+    ]
+)
+def test_get_default_cols(task_type, slug, expected):
+    attr_cols = ['Rank 🏆', 'Retrieval Method', 'Reranking Model', 'Revision', 'Submission Date', 'Average ⬆️']
+    cols, types = get_default_cols(task_type, slug)
+    benchmark_cols = list(frozenset(cols).difference(frozenset(attr_cols)))
+    assert len(benchmark_cols) == expected
+
+
+@pytest.mark.parametrize(
+    "task_type, domains, languages, expected",
+    [
+        (TaskType.qa, ["wiki", "news"], ["zh",], ["wiki_zh", "news_zh"]),
+        (TaskType.qa, ["law",], ["zh", "en"], ["law_en"]),
+        (
+            TaskType.long_doc,
+            ["healthcare"],
+            ["zh", "en"],
+            [
+                'healthcare_en_pubmed_100k_200k_1',
+                'healthcare_en_pubmed_100k_200k_2',
+                'healthcare_en_pubmed_100k_200k_3',
+                'healthcare_en_pubmed_40k_50k_5_merged',
+                'healthcare_en_pubmed_30k_40k_10_merged'
+            ]
+        )
+    ]
+)
+def test_get_selected_cols(task_type, domains, languages, expected):
+    slug = "air_bench_2404"
+    cols = get_selected_cols(task_type, slug, domains, languages)
+    assert sorted(cols) == sorted(expected)
 
 
-def test_get_default_auto_eval_column_dict():
-    auto_eval_column_dict_list = get_default_auto_eval_column_dict()
-    assert len(auto_eval_column_dict_list) == 9
+def test_select_columns(toy_df):
+    expected = [
+        'Rank 🏆',
+        'Retrieval Method',
+        'Reranking Model',
+        'Revision',
+        'Submission Date',
+        'Average ⬆️',
+        'news_zh']
+    df_result = select_columns(
+        toy_df,
+        [
+            "news",
+        ],
+        [
+            "zh",
+        ],
+        version_slug="air_bench_2404",
+    )
+    assert len(df_result.columns) == len(expected)
+    assert df_result["Average ⬆️"].equals(df_result["news_zh"])

tests/test_utils.py

@@ -75,18 +75,7 @@ def test_filter_queries(toy_df):
     assert df_result.iloc[0]["Retrieval Model"] == "jina-embeddings-v2-base"
 
 
-def test_select_columns(toy_df):
-    df_result = select_columns(
-        toy_df,
-        [
-            "news",
-        ],
-        [
-            "zh",
-        ],
-    )
-    assert len(df_result.columns) == 4
-    assert df_result["Average ⬆️"].equals(df_result["news_zh"])
+
 
 
 def test_update_table_long_doc(toy_df_long_doc):