Bram Vanroy commited on
Commit
6467c6f
1 Parent(s): 98e77c3

optimize parameter calculation

Browse files
generate_overview_json.py CHANGED
@@ -1,12 +1,17 @@
1
- from pathlib import Path
2
  import json
 
3
 
 
4
  from tqdm import tqdm
5
  from transformers import AutoModelForCausalLM
6
 
7
 
8
  def get_num_parameters(model_name: str) -> int:
9
- return AutoModelForCausalLM.from_pretrained(model_name).num_parameters()
 
 
 
 
10
 
11
 
12
  def main():
 
 
1
  import json
2
+ from pathlib import Path
3
 
4
+ from huggingface_hub import get_safetensors_metadata
5
  from tqdm import tqdm
6
  from transformers import AutoModelForCausalLM
7
 
8
 
9
  def get_num_parameters(model_name: str) -> int:
10
+ try:
11
+ metadata = get_safetensors_metadata(model_name)
12
+ return sum(metadata.parameter_count.values())
13
+ except Exception:
14
+ return AutoModelForCausalLM.from_pretrained(model_name).num_parameters()
15
 
16
 
17
  def main():
generate_overview_requirements.txt CHANGED
@@ -1 +1,2 @@
 
1
  transformers
 
1
+ huggingface_hub
2
  transformers