Skip to content

Commit

Permalink
Print traceback for safetensors metadata collection
Browse files Browse the repository at this point in the history
  • Loading branch information
tomaarsen committed Jul 29, 2024
1 parent 9f85f1d commit 19c9748
Showing 1 changed file with 4 additions and 2 deletions.
6 changes: 4 additions & 2 deletions utils/model_size.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import json
import re
import traceback
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

Expand All @@ -13,10 +14,11 @@ 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)
except Exception as e:
print(traceback.format_exc())
else:
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:
Expand Down

0 comments on commit 19c9748

Please sign in to comment.