Bram Vanroy commited on
Commit
6d7ff83
1 Parent(s): 8c4485d

update data collection script

Browse files
Files changed (1) hide show
  1. generate_overview_json.py +7 -5
generate_overview_json.py CHANGED
@@ -33,16 +33,18 @@ def main():
33
  if "pretrained" not in model_args:
34
  continue
35
 
36
- results[short_name]["model_name"] = model_args["pretrained"]
37
- results[short_name]["compute_dtype"] = model_args.get("dtype", None)
38
- results[short_name]["quantization"] = None
 
 
 
 
39
  if "load_in_8bit" in model_args:
40
  results[short_name]["quantization"] = "8-bit"
41
  elif "load_in_4bit" in model_args:
42
  results[short_name]["quantization"] = "4-bit"
43
 
44
- results[short_name]["num_parameters"] = get_num_parameters(model_args["pretrained"])
45
-
46
  pprint(results)
47
  pf_overview.write_text(json.dumps(results, indent=4), encoding="utf-8")
48
 
 
33
  if "pretrained" not in model_args:
34
  continue
35
 
36
+ results[short_name] = {
37
+ "model_name": model_args["pretrained"],
38
+ "compute_dtype": model_args.get("dtype", None),
39
+ "quantization": None,
40
+ "num_parameters": get_num_parameters(model_args["pretrained"])
41
+ }
42
+
43
  if "load_in_8bit" in model_args:
44
  results[short_name]["quantization"] = "8-bit"
45
  elif "load_in_4bit" in model_args:
46
  results[short_name]["quantization"] = "4-bit"
47
 
 
 
48
  pprint(results)
49
  pf_overview.write_text(json.dumps(results, indent=4), encoding="utf-8")
50