Spaces:
Running
Running
File size: 1,916 Bytes
4af3178 abf58d4 4af3178 abf58d4 4af3178 8168003 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 |
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 = {
"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}")
|