From 02ae4fabc3c0a7d733728c2c645c9e472fca42bb Mon Sep 17 00:00:00 2001
From: Alexey Vatolin <vatolinalex@gmail.com>
Date: Wed, 25 Dec 2024 00:00:55 +0100
Subject: [PATCH] fix: Update results_to_dataframe to use BenchmarkResults
 class (#1628)

---
 mteb/task_selection.py | 25 ++++++++++++-------------
 1 file changed, 12 insertions(+), 13 deletions(-)

diff --git a/mteb/task_selection.py b/mteb/task_selection.py
index 20d91a97b..4f00e542c 100644
--- a/mteb/task_selection.py
+++ b/mteb/task_selection.py
@@ -10,7 +10,7 @@
 from sklearn.preprocessing import StandardScaler
 from tqdm import tqdm
 
-import mteb
+from mteb.load_results.benchmark_results import BenchmarkResults
 
 MODEL_NAME = str
 REVISION = str
@@ -35,7 +35,7 @@ def mse_with_zscore(x: list[float], y: list[float]) -> float:
 
 
 def results_to_dataframe(
-    mteb_results: dict[MODEL_NAME, dict[REVISION, list[mteb.MTEBResults]]],
+    mteb_results: BenchmarkResults,
     drop_na: bool = True,
     **kwargs: Any,
 ) -> pd.DataFrame:
@@ -47,17 +47,16 @@ def results_to_dataframe(
         **kwargs: Additional keyword arguments to be passed to the `get_score` method of the `MTEBResults` class.
     """
     data = []
-    for model_name, revisions in mteb_results.items():
-        for rev, tasks_results in revisions.items():
-            for task_result in tasks_results:
-                data.append(
-                    {
-                        "Model": model_name,
-                        "Revision": rev,
-                        "task": task_result.task_name,
-                        "main_score": task_result.get_score(**kwargs),
-                    }
-                )
+    for model_res in mteb_results:
+        for task_result in model_res.task_results:
+            data.append(
+                {
+                    "Model": model_res.model_name,
+                    "Revision": model_res.model_revision,
+                    "task": task_result.task_name,
+                    "main_score": task_result.get_score(**kwargs),
+                }
+            )
     df = pd.DataFrame(data)
 
     if drop_na: