Tony Wu commited on
Commit
1f00e6d
1 Parent(s): 66d537f

feat: add support for new vidore output format

Browse files
Files changed (1) hide show
  1. data/model_handler.py +18 -3
data/model_handler.py CHANGED
@@ -1,6 +1,6 @@
1
  import json
2
  import os
3
- from typing import Dict
4
 
5
  import pandas as pd
6
  from huggingface_hub import HfApi, hf_hub_download, metadata_load
@@ -26,6 +26,21 @@ class ModelHandler:
26
  with open(self.model_infos_path, "w") as f:
27
  json.dump(self.model_infos, f)
28
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
29
  def get_vidore_data(self, metric="ndcg_at_5"):
30
  models = self.api.list_models(filter="vidore")
31
  repositories = [model.modelId for model in models] # type: ignore
@@ -55,8 +70,8 @@ class ModelHandler:
55
  with open(result_path) as f:
56
  results = json.load(f)
57
 
58
- for dataset in results:
59
- results[dataset] = {key: value for key, value in results[dataset].items()}
60
 
61
  self.model_infos[model_name] = {"meta": meta, "results": results}
62
  except Exception as e:
 
1
  import json
2
  import os
3
+ from typing import Any, Dict, Tuple
4
 
5
  import pandas as pd
6
  from huggingface_hub import HfApi, hf_hub_download, metadata_load
 
26
  with open(self.model_infos_path, "w") as f:
27
  json.dump(self.model_infos, f)
28
 
29
+ def _are_results_in_new_vidore_format(self, results: Dict[str, Any]) -> bool:
30
+ return "metadata" in results and "metrics" in results
31
+
32
+ @staticmethod
33
+ def convert_new_vidore_results_format(results: Dict[str, Any]) -> Tuple[str, str, Dict[str, Any]]:
34
+ if "metadata" not in results:
35
+ raise KeyError("results does not contain a 'metadata' key.")
36
+ if "metrics" not in results:
37
+ raise KeyError("results does not contain a 'metrics' key.")
38
+
39
+ metadata = results["metadata"]
40
+ metrics = results["metrics"]
41
+
42
+ return metadata["timestamp"], metadata["vidore_benchmark_hash"], metrics
43
+
44
  def get_vidore_data(self, metric="ndcg_at_5"):
45
  models = self.api.list_models(filter="vidore")
46
  repositories = [model.modelId for model in models] # type: ignore
 
70
  with open(result_path) as f:
71
  results = json.load(f)
72
 
73
+ if self._are_results_in_new_vidore_format(results):
74
+ timestamp, vidore_hash, results = self.convert_new_vidore_results_format(results)
75
 
76
  self.model_infos[model_name] = {"meta": meta, "results": results}
77
  except Exception as e: