Compute model size based on number of parameters

app.py

@@ -4,11 +4,13 @@ import os
 
 from datasets import load_dataset
 import gradio as gr
-from huggingface_hub import get_hf_file_metadata, HfApi, hf_hub_download, hf_hub_url
+from huggingface_hub import HfApi, hf_hub_download
 from huggingface_hub.repocard import metadata_load
 import pandas as pd
 from tqdm.autonotebook import tqdm
 
+from utils.model_size import get_model_size
+
 TASKS = [
     "BitextMining",
     "Classification",
@@ -786,94 +788,94 @@ EXTERNAL_MODEL_TO_SEQLEN = {
 }
 
 EXTERNAL_MODEL_TO_SIZE = {
-    "allenai-specter": 0.44,
-    "all-MiniLM-L12-v2": 0.13,
-    "all-MiniLM-L6-v2": 0.09,
-    "all-mpnet-base-v2": 0.44,
-    "bert-base-10lang-cased": 0.61,
-    "bert-base-15lang-cased": 0.61,
-    "bert-base-25lang-cased": 0.61,
-    "bert-base-multilingual-cased": 0.71,
-    "bert-base-multilingual-uncased": 0.67,
-    "bert-base-uncased": 0.44,
-    "bert-base-swedish-cased": 0.50,
-    "bge-base-zh-v1.5": 0.41,
-    "bge-large-zh-v1.5": 1.30,
-    "bge-large-zh-noinstruct": 1.30,
-    "bge-small-zh-v1.5": 0.10,
-    "camembert-base": 0.45,
-    "camembert-large": 1.35,
-    "cross-en-de-roberta-sentence-transformer": 1.11,
-    "contriever-base-msmarco": 0.44,
-    "distilbert-base-25lang-cased": 0.44,
-    "distilbert-base-en-fr-cased": 0.44,
-    "distilbert-base-en-fr-es-pt-it-cased": 0.44,
-    "distilbert-base-fr-cased": 0.44,
-    "distilbert-base-uncased": 0.44,
-    "DanskBERT": 0.50,
-    "distiluse-base-multilingual-cased-v2": 0.54,
-    "dfm-encoder-large-v1": 1.42,
-    "dfm-sentence-encoder-large-1": 1.63,
-    "e5-base": 0.44,
-    "e5-large": 1.34,
-    "e5-mistral-7b-instruct": 14.22,
-    "e5-small": 0.13,
-    "electra-small-nordic": 0.09,
-    "electra-small-swedish-cased-discriminator": 0.06,
-    "flaubert_base_cased": 0.55,
-    "flaubert_base_uncased": 0.55,
-    "flaubert_large_cased": 1.49,
-    "gbert-base": 0.44,
-    "gbert-large": 1.35,
-    "gelectra-base": 0.44,
-    "gelectra-large": 1.34,
-    "glove.6B.300d": 0.48,
-    "gottbert-base": 0.51,
-    "gtr-t5-base": 0.22,
-    "gtr-t5-large": 0.67,
-    "gtr-t5-xl": 2.48,
-    "gtr-t5-xxl": 9.73,
-    "herbert-base-retrieval-v2": 0.50,
-    "komninos": 0.27,
-    "luotuo-bert-medium": 1.31,    
-    "LASER2": 0.17,
-    "LaBSE": 1.88,
-    "m3e-base": 0.41,
-    "m3e-large": 0.41,
-    "msmarco-bert-co-condensor": 0.44,
-    "multi-qa-MiniLM-L6-cos-v1": 0.09,
-    "multilingual-e5-base": 1.11,
-    "multilingual-e5-small": 0.47,
-    "multilingual-e5-large": 2.24,
-    "nb-bert-base": 0.71,
-    "nb-bert-large": 1.42,
-    "nomic-embed-text-v1.5-64": 0.55,
-    "nomic-embed-text-v1.5-128": 0.55,
-    "nomic-embed-text-v1.5-256": 0.55,
-    "nomic-embed-text-v1.5-512": 0.55,
-    "norbert3-base": 0.52,
-    "norbert3-large": 1.47,
-    "paraphrase-multilingual-mpnet-base-v2": 1.11,
-    "paraphrase-multilingual-MiniLM-L12-v2": 0.47,
-    "sentence-camembert-base": 0.44,
-    "sentence-camembert-large": 1.35,
-    "sentence-croissant-llm-base": 5.12,
-    "sentence-bert-swedish-cased": 0.50,
-    "sentence-t5-base": 0.22,
-    "sentence-t5-large": 0.67,
-    "sentence-t5-xl": 2.48,
-    "sentence-t5-xxl": 9.73,
-    "silver-retriever-base-v1": 0.50,
-    "sup-simcse-bert-base-uncased": 0.44,
-    "st-polish-paraphrase-from-distilroberta": 0.50,
-    "st-polish-paraphrase-from-mpnet": 0.50,    
-    "text2vec-base-chinese": 0.41,
-    "text2vec-large-chinese": 1.30,
-    "unsup-simcse-bert-base-uncased": 0.44,
-    "use-cmlm-multilingual": 1.89,
-    "voyage-lite-02-instruct": 2.45,
-    "xlm-roberta-base": 1.12,
-    "xlm-roberta-large": 2.24,
+    "allenai-specter": 110,
+    "all-MiniLM-L12-v2": 33,
+    "all-MiniLM-L6-v2": 23,
+    "all-mpnet-base-v2": 110,
+    "bert-base-10lang-cased": 138,
+    "bert-base-15lang-cased": 138,
+    "bert-base-25lang-cased": 138,
+    "bert-base-multilingual-cased": 179,
+    "bert-base-multilingual-uncased": 168,
+    "bert-base-uncased": 110,
+    "bert-base-swedish-cased": 125,
+    "bge-base-zh-v1.5": 102,
+    "bge-large-zh-v1.5": 326,
+    "bge-large-zh-noinstruct": 326,
+    "bge-small-zh-v1.5": 24,
+    "camembert-base": 111,
+    "camembert-large": 338,
+    "cross-en-de-roberta-sentence-transformer": 278,
+    "contriever-base-msmarco": 110,
+    "distilbert-base-25lang-cased": 110,
+    "distilbert-base-en-fr-cased": 110,
+    "distilbert-base-en-fr-es-pt-it-cased": 110,
+    "distilbert-base-fr-cased": 110,
+    "distilbert-base-uncased": 110,
+    "DanskBERT": 125,
+    "distiluse-base-multilingual-cased-v2": 135,
+    "dfm-encoder-large-v1": 355,
+    "dfm-sentence-encoder-large-1": 355,
+    "e5-base": 110,
+    "e5-large": 335,
+    "e5-mistral-7b-instruct": 7110,
+    "e5-small": 33,
+    "electra-small-nordic": 23,
+    "electra-small-swedish-cased-discriminator": 16,
+    "flaubert_base_cased": 138,
+    "flaubert_base_uncased": 138,
+    "flaubert_large_cased": 372,
+    "gbert-base": 110,
+    "gbert-large": 337,
+    "gelectra-base": 110,
+    "gelectra-large": 335,
+    "glove.6B.300d": 120,
+    "gottbert-base": 127,
+    "gtr-t5-base": 110,
+    "gtr-t5-large": 168,
+    "gtr-t5-xl": 1240,
+    "gtr-t5-xxl": 4865,
+    "herbert-base-retrieval-v2": 125,
+    "komninos": 134,
+    "luotuo-bert-medium": 328,    
+    "LASER2": 43,
+    "LaBSE": 471,
+    "m3e-base": 102,
+    "m3e-large": 102,
+    "msmarco-bert-co-condensor": 110,
+    "multi-qa-MiniLM-L6-cos-v1": 23,
+    "multilingual-e5-base": 278,
+    "multilingual-e5-small": 118,
+    "multilingual-e5-large": 560,
+    "nb-bert-base": 179,
+    "nb-bert-large": 355,
+    "nomic-embed-text-v1.5-64": 138,
+    "nomic-embed-text-v1.5-128": 138,
+    "nomic-embed-text-v1.5-256": 138,
+    "nomic-embed-text-v1.5-512": 138,
+    "norbert3-base": 131,
+    "norbert3-large": 368,
+    "paraphrase-multilingual-mpnet-base-v2": 278,
+    "paraphrase-multilingual-MiniLM-L12-v2": 118,
+    "sentence-camembert-base": 110,
+    "sentence-camembert-large": 337,
+    "sentence-croissant-llm-base": 1280,
+    "sentence-bert-swedish-cased": 125,
+    "sentence-t5-base": 110,
+    "sentence-t5-large": 168,
+    "sentence-t5-xl": 1240,
+    "sentence-t5-xxl": 4865,
+    "silver-retriever-base-v1": 125,
+    "sup-simcse-bert-base-uncased": 110,
+    "st-polish-paraphrase-from-distilroberta": 125,
+    "st-polish-paraphrase-from-mpnet": 125,    
+    "text2vec-base-chinese": 102,
+    "text2vec-large-chinese": 326,
+    "unsup-simcse-bert-base-uncased": 110,
+    "use-cmlm-multilingual": 472,
+    "voyage-lite-02-instruct": 613,
+    "xlm-roberta-base": 279,
+    "xlm-roberta-large": 560,
 }
 
 MODELS_TO_SKIP = {
@@ -997,6 +999,7 @@ MODELS_TO_SKIP = {
     "beademiguelperez/sentence-transformers-multilingual-e5-small",
     "arcdev/SFR-Embedding-Mistral",
     "arcdev/e5-mistral-7b-instruct",
+    "Koat/gte-tiny",
 }
 
 def add_lang(examples):
@@ -1079,36 +1082,7 @@ def get_dim_seq_size(model):
             dim = config.get("hidden_dim", config.get("hidden_size", config.get("d_model", "")))
         seq = config.get("n_positions", config.get("max_position_embeddings", config.get("n_ctx", config.get("seq_length", ""))))
     # Get model file size without downloading
-    if "pytorch_model.bin" in filenames:
-        url = hf_hub_url(model.modelId, filename="pytorch_model.bin")
-        meta = get_hf_file_metadata(url)
-        size = round(meta.size / 1e9, 2)
-    elif "pytorch_model.bin.index.json" in filenames:
-        index_path = hf_hub_download(model.modelId, filename="pytorch_model.bin.index.json")
-        """
-        {
-        "metadata": {
-            "total_size": 28272820224
-        },....
-        """
-        size = json.load(open(index_path))
-        if ("metadata" in size) and ("total_size" in size["metadata"]):
-            size = round(size["metadata"]["total_size"] / 1e9, 2)
-    elif "model.safetensors" in filenames:
-        url = hf_hub_url(model.modelId, filename="model.safetensors")
-        meta = get_hf_file_metadata(url)
-        size = round(meta.size / 1e9, 2)
-    elif "model.safetensors.index.json" in filenames:
-        index_path = hf_hub_download(model.modelId, filename="model.safetensors.index.json")
-        """
-        {
-        "metadata": {
-            "total_size": 14483464192
-        },....
-        """
-        size = json.load(open(index_path))
-        if ("metadata" in size) and ("total_size" in size["metadata"]):
-            size = round(size["metadata"]["total_size"] / 1e9, 2)
+    size = get_model_size(model)
     return dim, seq, size
 
 def make_datasets_clickable(df):
@@ -1120,7 +1094,7 @@ def make_datasets_clickable(df):
     return df
 
 def add_rank(df):
-    cols_to_rank = [col for col in df.columns if col not in ["Model", "Model Size (GB)", "Embedding Dimensions", "Max Tokens"]]
+    cols_to_rank = [col for col in df.columns if col not in ["Model", "Model Size (Million Parameters)", "Embedding Dimensions", "Max Tokens"]]
     if len(cols_to_rank) == 1:
         df.sort_values(cols_to_rank[0], ascending=False, inplace=True)
     else:
@@ -1150,7 +1124,7 @@ def get_mteb_data(tasks=["Clustering"], langs=[], datasets=[], fillna=True, add_
         # Model & at least one result
         if len(res) > 1:
             if add_emb_dim:
-                res["Model Size (GB)"] = EXTERNAL_MODEL_TO_SIZE.get(model, "")
+                res["Model Size (Million Parameters)"] = EXTERNAL_MODEL_TO_SIZE.get(model, "")
                 res["Embedding Dimensions"] = EXTERNAL_MODEL_TO_DIM.get(model, "")
                 res["Max Tokens"] = EXTERNAL_MODEL_TO_SEQLEN.get(model, "")
             df_list.append(res)
@@ -1191,7 +1165,7 @@ def get_mteb_data(tasks=["Clustering"], langs=[], datasets=[], fillna=True, add_
             if add_emb_dim:
                 try:
                     # Fails on gated repos, so we only include scores for them
-                    out["Embedding Dimensions"], out["Max Tokens"], out["Model Size (GB)"] = get_dim_seq_size(model)
+                    out["Embedding Dimensions"], out["Max Tokens"], out["Model Size (Million Parameters)"] = get_dim_seq_size(model)
                 except:
                     pass
             df_list.append(out)
@@ -1268,7 +1242,7 @@ def get_mteb_average():
     # Fill NaN after averaging
     DATA_OVERALL.fillna("", inplace=True)
 
-    DATA_OVERALL = DATA_OVERALL[["Rank", "Model", "Model Size (GB)", "Embedding Dimensions", "Max Tokens", f"Average ({len(TASK_LIST_EN)} datasets)", f"Classification Average ({len(TASK_LIST_CLASSIFICATION)} datasets)", f"Clustering Average ({len(TASK_LIST_CLUSTERING)} datasets)", f"Pair Classification Average ({len(TASK_LIST_PAIR_CLASSIFICATION)} datasets)", f"Reranking Average ({len(TASK_LIST_RERANKING)} datasets)", f"Retrieval Average ({len(TASK_LIST_RETRIEVAL)} datasets)", f"STS Average ({len(TASK_LIST_STS)} datasets)", f"Summarization Average ({len(TASK_LIST_SUMMARIZATION)} dataset)"]]
+    DATA_OVERALL = DATA_OVERALL[["Rank", "Model", "Model Size (Million Parameters)", "Embedding Dimensions", "Max Tokens", f"Average ({len(TASK_LIST_EN)} datasets)", f"Classification Average ({len(TASK_LIST_CLASSIFICATION)} datasets)", f"Clustering Average ({len(TASK_LIST_CLUSTERING)} datasets)", f"Pair Classification Average ({len(TASK_LIST_PAIR_CLASSIFICATION)} datasets)", f"Reranking Average ({len(TASK_LIST_RERANKING)} datasets)", f"Retrieval Average ({len(TASK_LIST_RETRIEVAL)} datasets)", f"STS Average ({len(TASK_LIST_STS)} datasets)", f"Summarization Average ({len(TASK_LIST_SUMMARIZATION)} dataset)"]]
     DATA_OVERALL = DATA_OVERALL[DATA_OVERALL.iloc[:, 5:].ne("").any(axis=1)]
 
     return DATA_OVERALL
@@ -1327,7 +1301,7 @@ def get_mteb_average_zh():
     # Fill NaN after averaging
     DATA_OVERALL_ZH.fillna("", inplace=True)
 
-    DATA_OVERALL_ZH = DATA_OVERALL_ZH[["Rank", "Model", "Model Size (GB)", "Embedding Dimensions", "Max Tokens", f"Average ({len(TASK_LIST_ZH)} datasets)", f"Classification Average ({len(TASK_LIST_CLASSIFICATION_ZH)} datasets)", f"Clustering Average ({len(TASK_LIST_CLUSTERING_ZH)} datasets)", f"Pair Classification Average ({len(TASK_LIST_PAIR_CLASSIFICATION_ZH)} datasets)", f"Reranking Average ({len(TASK_LIST_RERANKING_ZH)} datasets)", f"Retrieval Average ({len(TASK_LIST_RETRIEVAL_ZH)} datasets)", f"STS Average ({len(TASK_LIST_STS_ZH)} datasets)"]]
+    DATA_OVERALL_ZH = DATA_OVERALL_ZH[["Rank", "Model", "Model Size (Million Parameters)", "Embedding Dimensions", "Max Tokens", f"Average ({len(TASK_LIST_ZH)} datasets)", f"Classification Average ({len(TASK_LIST_CLASSIFICATION_ZH)} datasets)", f"Clustering Average ({len(TASK_LIST_CLUSTERING_ZH)} datasets)", f"Pair Classification Average ({len(TASK_LIST_PAIR_CLASSIFICATION_ZH)} datasets)", f"Reranking Average ({len(TASK_LIST_RERANKING_ZH)} datasets)", f"Retrieval Average ({len(TASK_LIST_RETRIEVAL_ZH)} datasets)", f"STS Average ({len(TASK_LIST_STS_ZH)} datasets)"]]
     DATA_OVERALL_ZH = DATA_OVERALL_ZH[DATA_OVERALL_ZH.iloc[:, 5:].ne("").any(axis=1)]
 
     return DATA_OVERALL_ZH
@@ -1389,7 +1363,7 @@ def get_mteb_average_fr():
     # Fill NaN after averaging
     DATA_OVERALL_FR.fillna("", inplace=True)
 
-    DATA_OVERALL_FR = DATA_OVERALL_FR[["Rank", "Model", "Model Size (GB)", "Embedding Dimensions", "Max Tokens", f"Average ({len(TASK_LIST_FR)} datasets)", f"Classification Average ({len(TASK_LIST_CLASSIFICATION_FR)} datasets)", f"Clustering Average ({len(TASK_LIST_CLUSTERING_FR)} datasets)", f"Pair Classification Average ({len(TASK_LIST_PAIR_CLASSIFICATION_FR)} datasets)", f"Reranking Average ({len(TASK_LIST_RERANKING_FR)} datasets)", f"Retrieval Average ({len(TASK_LIST_RETRIEVAL_FR)} datasets)", f"STS Average ({len(TASK_LIST_STS_FR)} datasets)", f"Summarization Average ({len(TASK_LIST_SUMMARIZATION_FR)} dataset)"]]
+    DATA_OVERALL_FR = DATA_OVERALL_FR[["Rank", "Model", "Model Size (Million Parameters)", "Embedding Dimensions", "Max Tokens", f"Average ({len(TASK_LIST_FR)} datasets)", f"Classification Average ({len(TASK_LIST_CLASSIFICATION_FR)} datasets)", f"Clustering Average ({len(TASK_LIST_CLUSTERING_FR)} datasets)", f"Pair Classification Average ({len(TASK_LIST_PAIR_CLASSIFICATION_FR)} datasets)", f"Reranking Average ({len(TASK_LIST_RERANKING_FR)} datasets)", f"Retrieval Average ({len(TASK_LIST_RETRIEVAL_FR)} datasets)", f"STS Average ({len(TASK_LIST_STS_FR)} datasets)", f"Summarization Average ({len(TASK_LIST_SUMMARIZATION_FR)} dataset)"]]
     DATA_OVERALL_FR = DATA_OVERALL_FR[DATA_OVERALL_FR.iloc[:, 5:].ne("").any(axis=1)]
 
     return DATA_OVERALL_FR
@@ -1443,7 +1417,7 @@ def get_mteb_average_pl():
     # Fill NaN after averaging
     DATA_OVERALL_PL.fillna("", inplace=True)
 
-    DATA_OVERALL_PL = DATA_OVERALL_PL[["Rank", "Model", "Model Size (GB)", "Embedding Dimensions", "Max Tokens", f"Average ({len(TASK_LIST_PL)} datasets)", f"Classification Average ({len(TASK_LIST_CLASSIFICATION_PL)} datasets)", f"Clustering Average ({len(TASK_LIST_CLUSTERING_PL)} datasets)", f"Pair Classification Average ({len(TASK_LIST_PAIR_CLASSIFICATION_PL)} datasets)", f"Retrieval Average ({len(TASK_LIST_RETRIEVAL_PL)} datasets)", f"STS Average ({len(TASK_LIST_STS_PL)} datasets)"]]
+    DATA_OVERALL_PL = DATA_OVERALL_PL[["Rank", "Model", "Model Size (Million Parameters)", "Embedding Dimensions", "Max Tokens", f"Average ({len(TASK_LIST_PL)} datasets)", f"Classification Average ({len(TASK_LIST_CLASSIFICATION_PL)} datasets)", f"Clustering Average ({len(TASK_LIST_CLUSTERING_PL)} datasets)", f"Pair Classification Average ({len(TASK_LIST_PAIR_CLASSIFICATION_PL)} datasets)", f"Retrieval Average ({len(TASK_LIST_RETRIEVAL_PL)} datasets)", f"STS Average ({len(TASK_LIST_STS_PL)} datasets)"]]
     DATA_OVERALL_PL = DATA_OVERALL_PL[DATA_OVERALL_PL.iloc[:, 5:].ne("").any(axis=1)]
 
     return DATA_OVERALL_PL

utils/model_size.py

@@ -0,0 +1,39 @@
+import json
+import re
+from huggingface_hub.hf_api import ModelInfo, get_safetensors_metadata, model_info as get_model_info, get_hf_file_metadata, hf_hub_url
+from huggingface_hub import hf_hub_download
+
+# Map model IDs to the number of bytes used for one parameter. So, 4 bytes for fp32, 2 bytes for fp16, etc.
+# By default, we assume that the model is stored in fp32.
+KNOWN_BYTES_PER_PARAM = {}
+
+
+def get_model_size(model_info: ModelInfo):
+    '''Get the size of the model in million of parameters.'''
+    try:
+        safetensors = get_safetensors_metadata(model_info.id)
+        return round(sum(safetensors.parameter_count.values()) / 1e6)
+    except Exception as e:
+        pass
+
+    filenames = [sib.rfilename for sib in model_info.siblings]
+    if "pytorch_model.bin" in filenames:
+        url = hf_hub_url(model_info.id, filename="pytorch_model.bin")
+        meta = get_hf_file_metadata(url)
+        bytes_per_param = KNOWN_BYTES_PER_PARAM.get(model_info.id, 4)
+        return round(meta.size / bytes_per_param / 1e6)
+    
+    if "pytorch_model.bin.index.json" in filenames:
+        index_path = hf_hub_download(model_info.id, filename="pytorch_model.bin.index.json")
+        """
+        {
+        "metadata": {
+            "total_size": 28272820224
+        },....
+        """
+        size = json.load(open(index_path))
+        bytes_per_param = KNOWN_BYTES_PER_PARAM.get(model_info.id, 4)
+        if ("metadata" in size) and ("total_size" in size["metadata"]):
+            return round(size["metadata"]["total_size"] / bytes_per_param / 1e6)
+
+    return None