Tom Aarsen
commited on
Commit
·
19c9748
1
Parent(s):
9f85f1d
Print traceback for safetensors metadata collection
Browse files- utils/model_size.py +4 -2
utils/model_size.py
CHANGED
|
@@ -1,5 +1,6 @@
|
|
| 1 |
import json
|
| 2 |
import re
|
|
|
|
| 3 |
from huggingface_hub.hf_api import ModelInfo, get_safetensors_metadata, model_info as get_model_info, get_hf_file_metadata, hf_hub_url
|
| 4 |
from huggingface_hub import hf_hub_download
|
| 5 |
|
|
@@ -13,10 +14,11 @@ def get_model_parameters_memory(model_info: ModelInfo):
|
|
| 13 |
'''Get the size of the model in million of parameters.'''
|
| 14 |
try:
|
| 15 |
safetensors = get_safetensors_metadata(model_info.id)
|
|
|
|
|
|
|
|
|
|
| 16 |
num_parameters = sum(safetensors.parameter_count.values())
|
| 17 |
return round(num_parameters / 1e6), round(num_parameters * 4 / 1024**3, 2)
|
| 18 |
-
except Exception as e:
|
| 19 |
-
pass
|
| 20 |
|
| 21 |
filenames = [sib.rfilename for sib in model_info.siblings]
|
| 22 |
if "pytorch_model.bin" in filenames:
|
|
|
|
| 1 |
import json
|
| 2 |
import re
|
| 3 |
+
import traceback
|
| 4 |
from huggingface_hub.hf_api import ModelInfo, get_safetensors_metadata, model_info as get_model_info, get_hf_file_metadata, hf_hub_url
|
| 5 |
from huggingface_hub import hf_hub_download
|
| 6 |
|
|
|
|
| 14 |
'''Get the size of the model in million of parameters.'''
|
| 15 |
try:
|
| 16 |
safetensors = get_safetensors_metadata(model_info.id)
|
| 17 |
+
except Exception as e:
|
| 18 |
+
print(traceback.format_exc())
|
| 19 |
+
else:
|
| 20 |
num_parameters = sum(safetensors.parameter_count.values())
|
| 21 |
return round(num_parameters / 1e6), round(num_parameters * 4 / 1024**3, 2)
|
|
|
|
|
|
|
| 22 |
|
| 23 |
filenames = [sib.rfilename for sib in model_info.siblings]
|
| 24 |
if "pytorch_model.bin" in filenames:
|