add subset to old format (#33)

embeddings-benchmark · Sep 16, 2024 · 1b85142 · 1b85142
1 parent 7e0fe6f
commit 1b85142
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/results.py b/results.py
@@ -371,7 +371,7 @@ def _generate_examples(self, filepath):
                         # Old MTEB format
                         if isinstance(res_dict.get(split), dict):
                             for k, v in res_dict.get(split, {}).items():
-                                if key in ["hf_subset", "languages"]:
+                                if k in ["hf_subset", "languages"]:
                                     res_dict[k] = v
 
                                 v /= len(TEST_AVG_SPLIT[ds_name])
@@ -383,7 +383,7 @@ def _generate_examples(self, filepath):
                         elif isinstance(res_dict.get(split), list):
                             assert len(res_dict[split]) == 1, "Only single-lists supported for now"
                             for k, v in res_dict[split][0].items():
-                                if key in ["hf_subset", "languages"]:
+                                if k in ["hf_subset", "languages"]:
                                     res_dict[k] = v
                                 if not isinstance(v, float): continue
                                 v /= len(TEST_AVG_SPLIT[ds_name])
@@ -451,6 +451,8 @@ def _generate_examples(self, filepath):
                             if lang in SKIP_KEYS: continue
                             test_result_lang = res_dict.get(lang) if is_multilingual else res_dict
                             subset = test_result_lang.pop("hf_subset", "")
+                            if subset == "" and is_multilingual:
+                                subset = lang
                             for metric, score in test_result_lang.items():
                                 if not isinstance(score, dict):
                                     score = {metric: score}