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