feat : support for results.json

data/dataset_handler.py

@@ -1,64 +1,45 @@
-from typing import Dict
-from huggingface_hub import get_collection
-
-
-def get_datasets_nickname() -> Dict:
-    datasets_nickname = {}
-
-    collection = get_collection("vidore/vidore-benchmark-667173f98e70a1c0fa4db00d")
-    collection_items = collection.items
-
-    for item in collection_items:
-        dataset_name = item.item_id
-
-        if 'arxivqa' in dataset_name:
-            datasets_nickname[dataset_name] = 'ArxivQA'
-            datasets_nickname[dataset_name + '_ocr_chunk'] = 'ArxivQA'
-            datasets_nickname[dataset_name + '_captioning'] = 'ArxivQA'
-
-        elif 'docvqa' in dataset_name:
-            datasets_nickname[dataset_name] = 'DocVQA'
-            datasets_nickname[dataset_name + '_ocr_chunk'] = 'DocVQA'
-            datasets_nickname[dataset_name + '_captioning'] = 'DocVQA'
-
-        elif 'infovqa' in dataset_name:
-            datasets_nickname[dataset_name] = 'InfoVQA'
-            datasets_nickname[dataset_name + '_ocr_chunk'] = 'InfoVQA'
-            datasets_nickname[dataset_name + '_captioning'] = 'InfoVQA'
-
-        elif 'tabfquad' in dataset_name:
-            datasets_nickname[dataset_name] = 'TabFQuad'
-            datasets_nickname[dataset_name + '_ocr_chunk'] = 'TabFQuad'
-            datasets_nickname[dataset_name + '_captioning'] = 'TabFQuad'
-
-        elif 'tatdqa' in dataset_name:
-            datasets_nickname[dataset_name] = 'TATDQA'
-            datasets_nickname[dataset_name + '_ocr_chunk'] = 'TATDQA'
-            datasets_nickname[dataset_name + '_captioning'] = 'TATDQA'
-
-        elif 'shiftproject' in dataset_name:
-            datasets_nickname[dataset_name] = 'ShiftProject'
-            datasets_nickname[dataset_name + '_ocr_chunk'] = 'ShiftProject'
-            datasets_nickname[dataset_name + '_captioning'] = 'ShiftProject'
-
-        elif 'artificial_intelligence' in dataset_name:
-            datasets_nickname[dataset_name] = 'Artificial Intelligence'
-            datasets_nickname[dataset_name + '_ocr_chunk'] = 'Artificial Intelligence'
-            datasets_nickname[dataset_name + '_captioning'] = 'Artificial Intelligence'
-
-        elif 'energy' in dataset_name:
-            datasets_nickname[dataset_name] = 'Energy'
-            datasets_nickname[dataset_name + '_ocr_chunk'] = 'Energy'
-            datasets_nickname[dataset_name + '_captioning'] = 'Energy'
-
-        elif 'government_reports' in dataset_name:
-            datasets_nickname[dataset_name] = 'Government Reports'
-            datasets_nickname[dataset_name + '_ocr_chunk'] = 'Government Reports'
-            datasets_nickname[dataset_name + '_captioning'] = 'Government Reports'
-
-        elif 'healthcare' in dataset_name:
-            datasets_nickname[dataset_name] = 'Healthcare'
-            datasets_nickname[dataset_name + '_ocr_chunk'] = 'Healthcare'
-            datasets_nickname[dataset_name + '_captioning'] = 'Healthcare'
-
-    return datasets_nickname
+VIDORE_DATASETS_KEYWORDS = [
+    "arxivqa",
+    "docvqa",
+    "infovqa",
+    "tabfquad",
+    "tatdqa",
+    "shift",
+    "artificial_intelligence",
+    "energy",
+    "government_reports",
+    "healthcare_industry",
+]
+
+
+def get_datasets_nickname(dataset_name) -> str:
+    if "arxivqa" in dataset_name:
+        return "ArxivQA"
+
+    elif "docvqa" in dataset_name:
+        return "DocVQA"
+    elif "infovqa" in dataset_name:
+        return "InfoVQA"
+    elif "tabfquad" in dataset_name:
+        return "TabFQuad"
+
+    elif "tatdqa" in dataset_name:
+        return "TAT-DQA"
+
+    elif "shift" in dataset_name:
+        return "Shift Project"
+
+    elif "artificial_intelligence" in dataset_name:
+        return "Artificial Intelligence"
+
+    elif "energy" in dataset_name:
+        return "Energy"
+
+    elif "government_reports" in dataset_name:
+        return "Government Reports"
+
+    elif "healthcare_industry" in dataset_name:
+        return "Healthcare Industry"
+
+    else:
+        raise ValueError(f"Dataset {dataset_name} not found in ViDoRe")

data/model_handler.py

@@ -3,7 +3,7 @@ import os
 from typing import Dict
 from huggingface_hub import HfApi, hf_hub_download, metadata_load
 import pandas as pd
-from .dataset_handler import get_datasets_nickname
+from .dataset_handler import get_datasets_nickname, VIDORE_DATASETS_KEYWORDS
 
 class ModelHandler:
     def __init__(self, model_infos_path="model_infos.json"):
@@ -25,14 +25,18 @@ class ModelHandler:
         models = self.api.list_models(filter="vidore")
         repositories = [model.modelId for model in models]  # type: ignore
 
-        datasets_nickname = get_datasets_nickname()
         for repo_id in repositories:
-            files = [f for f in self.api.list_repo_files(repo_id) if f.endswith('_metrics.json')]
+            files = [f for f in self.api.list_repo_files(repo_id) if f.endswith('_metrics.json') or f == 'results.json']
+
+            
             if len(files) == 0:
                 continue
             else:
                 for file in files:
-                    model_name = file.split('_metrics.json')[0]
+                    if file.endswith('results.json'):
+                        model_name = repo_id.replace('/', '_')
+                    else:
+                        model_name = file.split('_metrics.json')[0]
 
                     if model_name not in self.model_infos:
                         readme_path = hf_hub_download(repo_id, filename="README.md")
@@ -51,7 +55,7 @@ class ModelHandler:
                             print(f"Error loading {model_name} - {e}")
                             continue
 
-        #self._save_model_infos()
+        self._save_model_infos()
 
         model_res = {}
         if len(self.model_infos) > 0:
@@ -59,12 +63,17 @@ class ModelHandler:
                 res = self.model_infos[model]["results"]
                 dataset_res = {}
                 for dataset in res.keys():
-                    if "validation_set" == dataset:
+                    #for each keyword check if it is in the dataset name if not continue
+                    if not any(keyword in dataset for keyword in VIDORE_DATASETS_KEYWORDS):
+                        print(f"{dataset} not found in ViDoRe datasets. Skipping ...")
                         continue
-                    dataset_res[datasets_nickname[dataset]] = res[dataset][metric]
+
+                    dataset_nickname = get_datasets_nickname(dataset)
+                    dataset_res[dataset_nickname] = res[dataset][metric]
                 model_res[model] = dataset_res
                 
             df = pd.DataFrame(model_res).T
+            
             return df
         return pd.DataFrame()
 
@@ -88,7 +97,10 @@ class ModelHandler:
             df.insert(len(df.columns) - len(cols_to_rank), "Average", df[cols_to_rank].mean(axis=1, skipna=False))
             df.sort_values("Average", ascending=False, inplace=True)
         df.insert(0, "Rank", list(range(1, len(df) + 1)))
-        df = df.round(2)
+        #multiply values by 100 if they are floats and round to 1 decimal place
+        for col in df.columns:
+            if df[col].dtype == "float64":
+                df[col] = df[col].apply(lambda x: round(x * 100, 1))
         # Fill NaN after averaging
         df.fillna("", inplace=True)
         return df