|
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 |
|
|
|
|
|
|
|
KNOWN_BYTES_PER_PARAM = { |
|
"dwzhu/e5-base-4k": 2, |
|
} |
|
|
|
def get_model_parameters_memory(model_info: ModelInfo): |
|
'''Get the size of the model in million of parameters.''' |
|
try: |
|
safetensors = get_safetensors_metadata(model_info.id) |
|
num_parameters = sum(safetensors.parameter_count.values()) |
|
return round(num_parameters / 1e6), round(num_parameters * 4 / 1024**3, 2) |
|
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) |
|
num_params = round(meta.size / bytes_per_param / 1e6) |
|
size_gb = round(meta.size * (4 / bytes_per_param) / 1024**3, 2) |
|
return num_params, size_gb |
|
|
|
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), round(size["metadata"]["total_size"] / 1024**3, 2) |
|
|
|
raise Exception(f"Could not find the model parameters for {model_info.id}") |
|
|