Spaces:
Running
Running
| import json | |
| import os | |
| from typing import Any, Dict | |
| import pandas as pd | |
| from huggingface_hub import HfApi, hf_hub_download, metadata_load | |
| from .dataset_handler import VIDORE_2_DATASETS_KEYWORDS, VIDORE_DATASETS_KEYWORDS, get_datasets_nickname | |
| BLOCKLIST = ["impactframes"] | |
| class ModelHandler: | |
| def __init__(self, model_infos_path="model_infos.json"): | |
| self.api = HfApi() | |
| self.model_infos_path = model_infos_path | |
| self.model_infos = self._load_model_infos() | |
| def _load_model_infos(self) -> Dict: | |
| if os.path.exists(self.model_infos_path): | |
| with open(self.model_infos_path) as f: | |
| return json.load(f) | |
| return {} | |
| def _save_model_infos(self): | |
| with open(self.model_infos_path, "w") as f: | |
| json.dump(self.model_infos, f) | |
| def _are_results_in_new_vidore_format(self, results: Dict[str, Any]) -> bool: | |
| return "metadata" in results and "metrics" in results | |
| def _is_baseline_repo(self, repo_id: str) -> bool: | |
| return repo_id == "vidore/baseline-results" | |
| def sanitize_model_name(self, model_name): | |
| return model_name.replace("/", "_").replace(".", "-thisisapoint-") | |
| def fuze_model_infos(self, model_name, results): | |
| for dataset, metrics in results.items(): | |
| if dataset not in self.model_infos[model_name]["results"].keys(): | |
| self.model_infos[model_name]["results"][dataset] = metrics | |
| else: | |
| continue | |
| def get_vidore_data(self, metric="ndcg_at_5"): | |
| models = self.api.list_models(filter="vidore") | |
| repositories = [model.modelId for model in models] # type: ignore | |
| # Sort repositories to process non-baseline repos first (to prioritize their results) | |
| repositories.sort(key=lambda x: self._is_baseline_repo(x)) | |
| for repo_id in repositories: | |
| org_name = repo_id.split("/")[0] | |
| if org_name in BLOCKLIST: | |
| continue | |
| 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: | |
| if file.endswith("results.json"): | |
| model_name = repo_id.replace("/", "_").replace(".", "-thisisapoint-") | |
| else: | |
| model_name = file.split("_metrics.json")[0] | |
| model_name = model_name.replace("/", "_").replace(".", "-thisisapoint-") | |
| # Skip if the model is from baseline and we already have results | |
| readme_path = hf_hub_download(repo_id, filename="README.md") | |
| meta = metadata_load(readme_path) | |
| try: | |
| result_path = hf_hub_download(repo_id, filename=file) | |
| with open(result_path) as f: | |
| results = json.load(f) | |
| if self._are_results_in_new_vidore_format(results): | |
| metadata = results["metadata"] | |
| results = results["metrics"] | |
| # Handles the case where the model is both in baseline and outside of it | |
| # (prioritizes the non-baseline results) | |
| if self._is_baseline_repo(repo_id) and self.sanitize_model_name(model_name) in self.model_infos: | |
| self.fuze_model_infos(model_name, results) | |
| self.model_infos[model_name] = {"meta": meta, "results": results} | |
| except Exception as e: | |
| print(f"Error loading {model_name} - {e}") | |
| continue | |
| # In order to keep only models relevant to a benchmark | |
| def filter_models_by_benchmark(self, benchmark_version=1): | |
| filtered_model_infos = {} | |
| keywords = VIDORE_DATASETS_KEYWORDS if benchmark_version == 1 else VIDORE_2_DATASETS_KEYWORDS | |
| for model, info in self.model_infos.items(): | |
| results = info["results"] | |
| if any(any(keyword in dataset for keyword in keywords) for dataset in results.keys()): | |
| filtered_model_infos[model] = info | |
| return filtered_model_infos | |
| # Compute the average of a metric for each model, | |
| def compute_averages(self, metric="ndcg_at_5", benchmark_version=1): | |
| model_res = {} | |
| filtered_model_infos = self.filter_models_by_benchmark(benchmark_version) | |
| if len(filtered_model_infos) > 0: | |
| for model in filtered_model_infos.keys(): | |
| res = filtered_model_infos[model]["results"] | |
| dataset_res = {} | |
| keywords = VIDORE_DATASETS_KEYWORDS if benchmark_version == 1 else VIDORE_2_DATASETS_KEYWORDS | |
| for dataset in res.keys(): | |
| if not any(keyword in dataset for keyword in keywords): | |
| continue | |
| 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() | |
| def add_rank(df, benchmark_version=1): | |
| df.fillna(0.0, inplace=True) | |
| cols_to_rank = [ | |
| col | |
| for col in df.columns | |
| if col | |
| not in [ | |
| "Model", | |
| "Model Size (Million Parameters)", | |
| "Memory Usage (GB, fp32)", | |
| "Embedding Dimensions", | |
| "Max Tokens", | |
| ] | |
| ] | |
| if len(cols_to_rank) == 1: | |
| df.sort_values(cols_to_rank[0], ascending=False, inplace=True) | |
| else: | |
| 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))) | |
| # 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)) | |
| return df | |