refactor: move the data model

src/read_evals.py

@@ -1,18 +1,13 @@
-import json
 import os.path
-from collections import defaultdict
-from dataclasses import dataclass
 from typing import List
 
 import pandas as pd
 
-from src.benchmarks import get_safe_name, BenchmarksQA, BenchmarksLongDoc
+from src.benchmarks import BenchmarksQA, BenchmarksLongDoc
 from src.display.utils import COLS_QA, COLS_LONG_DOC
-from src.display.column_names import COL_NAME_AVG, COL_NAME_RETRIEVAL_MODEL, COL_NAME_RERANKING_MODEL, \
-    COL_NAME_RETRIEVAL_MODEL_LINK, COL_NAME_RERANKING_MODEL_LINK, COL_NAME_RANK, COL_NAME_REVISION, COL_NAME_TIMESTAMP, \
-    COL_NAME_IS_ANONYMOUS
+from src.display.column_names import COL_NAME_AVG, COL_NAME_RANK, COL_NAME_REVISION, COL_NAME_IS_ANONYMOUS
 
-from src.display.formatting import make_clickable_model
+from src.models import FullEvalResult
 
 pd.options.mode.copy_on_write = True
 
@@ -24,121 +19,6 @@ def calculate_mean(row):
         return row.mean()
 
 
-@dataclass
-class EvalResult:
-    """
-    Evaluation result of a single embedding model with a specific reranking model on benchmarks over different
-    domains, languages, and datasets
-    """
-    eval_name: str  # name of the evaluation, [retrieval_model]_[reranking_model]_[metric]
-    retrieval_model: str
-    reranking_model: str
-    results: list  # results on all the benchmarks stored as dict
-    task: str
-    metric: str
-    timestamp: str = ""  # submission timestamp
-    revision: str = ""
-    is_anonymous: bool = False
-
-
-@dataclass
-class FullEvalResult:
-    """
-    Evaluation result of a single embedding model with a specific reranking model on benchmarks over different tasks
-    """
-    eval_name: str  # name of the evaluation, [retrieval_model]_[reranking_model]
-    retrieval_model: str
-    reranking_model: str
-    retrieval_model_link: str
-    reranking_model_link: str
-    results: List[EvalResult]  # results on all the EvalResults over different tasks and metrics.
-    timestamp: str = ""
-    revision: str = ""
-    is_anonymous: bool = False
-
-    @classmethod
-    def init_from_json_file(cls, json_filepath):
-        """
-        Initiate from the result json file for a single model.
-        The json file will be written only when the status is FINISHED.
-        """
-        with open(json_filepath) as fp:
-            model_data = json.load(fp)
-
-        # store all the results for different metrics and tasks
-        result_list = []
-        retrieval_model_link = ""
-        reranking_model_link = ""
-        revision = ""
-        for item in model_data:
-            config = item.get("config", {})
-            # eval results for different metrics
-            results = item.get("results", [])
-            retrieval_model_link = config["retrieval_model_link"]
-            if config["reranking_model_link"] is None:
-                reranking_model_link = ""
-            else:
-                reranking_model_link = config["reranking_model_link"]
-            eval_result = EvalResult(
-                eval_name=f"{config['retrieval_model']}_{config['reranking_model']}_{config['metric']}",
-                retrieval_model=config["retrieval_model"],
-                reranking_model=config["reranking_model"],
-                results=results,
-                task=config["task"],
-                metric=config["metric"],
-                timestamp=config.get("timestamp", "2024-05-12T12:24:02Z"),
-                revision=config.get("revision", "3a2ba9dcad796a48a02ca1147557724e"),
-                is_anonymous=config.get("is_anonymous", False)
-            )
-            result_list.append(eval_result)
-        return cls(
-            eval_name=f"{result_list[0].retrieval_model}_{result_list[0].reranking_model}",
-            retrieval_model=result_list[0].retrieval_model,
-            reranking_model=result_list[0].reranking_model,
-            retrieval_model_link=retrieval_model_link,
-            reranking_model_link=reranking_model_link,
-            results=result_list,
-            timestamp=result_list[0].timestamp,
-            revision=result_list[0].revision,
-            is_anonymous=result_list[0].is_anonymous
-        )
-
-    def to_dict(self, task='qa', metric='ndcg_at_3') -> List:
-        """
-        Convert the results in all the EvalResults over different tasks and metrics. The output is a list of dict compatible with the dataframe UI
-        """
-        results = defaultdict(dict)
-        for eval_result in self.results:
-            if eval_result.metric != metric:
-                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(self.retrieval_model, self.retrieval_model_link))
-            results[eval_result.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
-
-            # print(f'result loaded: {eval_result.eval_name}')
-            for result in eval_result.results:
-                # add result for each domain, language, and dataset
-                domain = result["domain"]
-                lang = result["lang"]
-                dataset = result["dataset"]
-                value = result["value"] * 100
-                if dataset == 'default':
-                    benchmark_name = f"{domain}_{lang}"
-                else:
-                    benchmark_name = f"{domain}_{lang}_{dataset}"
-                results[eval_result.eval_name][get_safe_name(benchmark_name)] = value
-        return [v for v in results.values()]
-
-
 def get_raw_eval_results(results_path: str) -> List[FullEvalResult]:
     """
     Load the evaluation results from a json file

src/utils.py

@@ -12,7 +12,8 @@ from src.display.utils import COLS_QA, TYPES_QA, COLS_LONG_DOC, TYPES_LONG_DOC,
 from src.display.column_names 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
 from src.envs import API, SEARCH_RESULTS_REPO, LATEST_BENCHMARK_VERSION
-from src.read_evals import FullEvalResult, get_leaderboard_df, calculate_mean
+from src.read_evals import get_leaderboard_df, calculate_mean
+from src.models import FullEvalResult
 
 import re
 

tests/src/test_read_evals.py

@@ -1,6 +1,7 @@
 from pathlib import Path
 
-from src.read_evals import FullEvalResult, get_raw_eval_results, get_leaderboard_df
+from src.read_evals import get_raw_eval_results, get_leaderboard_df
+from src.models import FullEvalResult
 
 cur_fp = Path(__file__)