Tom Aarsen commited on
Commit
abf58d4
1 Parent(s): 9a53e83

Add dwzhu/e5-base-4k as a pytorch_model.bin model with fp16

Browse files
Files changed (1) hide show
  1. utils/model_size.py +6 -3
utils/model_size.py CHANGED
@@ -5,8 +5,9 @@ from huggingface_hub import hf_hub_download
5
 
6
  # Map model IDs to the number of bytes used for one parameter. So, 4 bytes for fp32, 2 bytes for fp16, etc.
7
  # By default, we assume that the model is stored in fp32.
8
- KNOWN_BYTES_PER_PARAM = {}
9
-
 
10
 
11
  def get_model_parameters_memory(model_info: ModelInfo):
12
  '''Get the size of the model in million of parameters.'''
@@ -22,7 +23,9 @@ def get_model_parameters_memory(model_info: ModelInfo):
22
  url = hf_hub_url(model_info.id, filename="pytorch_model.bin")
23
  meta = get_hf_file_metadata(url)
24
  bytes_per_param = KNOWN_BYTES_PER_PARAM.get(model_info.id, 4)
25
- return round(meta.size / bytes_per_param / 1e6), round(meta.size / 1024**3, 2)
 
 
26
 
27
  if "pytorch_model.bin.index.json" in filenames:
28
  index_path = hf_hub_download(model_info.id, filename="pytorch_model.bin.index.json")
 
5
 
6
  # Map model IDs to the number of bytes used for one parameter. So, 4 bytes for fp32, 2 bytes for fp16, etc.
7
  # By default, we assume that the model is stored in fp32.
8
+ KNOWN_BYTES_PER_PARAM = {
9
+ "dwzhu/e5-base-4k": 2,
10
+ }
11
 
12
  def get_model_parameters_memory(model_info: ModelInfo):
13
  '''Get the size of the model in million of parameters.'''
 
23
  url = hf_hub_url(model_info.id, filename="pytorch_model.bin")
24
  meta = get_hf_file_metadata(url)
25
  bytes_per_param = KNOWN_BYTES_PER_PARAM.get(model_info.id, 4)
26
+ num_params = round(meta.size / bytes_per_param / 1e6)
27
+ size_gb = round(meta.size * (4 / bytes_per_param) / 1024**3, 2)
28
+ return num_params, size_gb
29
 
30
  if "pytorch_model.bin.index.json" in filenames:
31
  index_path = hf_hub_download(model_info.id, filename="pytorch_model.bin.index.json")