File size: 1,432 Bytes
828458d 78588de 828458d 78588de 828458d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
from pathlib import Path
import json
from pprint import pprint
from transformers import AutoModelForCausalLM
def get_num_parameters(model_name: str) -> int:
return AutoModelForCausalLM.from_pretrained(model_name).num_parameters()
def main():
results = {}
for pfin in Path(__file__).parent.joinpath("evals").rglob("*.json"):
if pfin.stem == "models":
continue
short_name = pfin.stem.split("_")[2]
if short_name not in results:
results[short_name] = {}
data = json.loads(pfin.read_text(encoding="utf-8"))
if "config" not in data:
continue
config = data["config"]
if "model_args" not in config:
continue
model_args = dict(params.split("=") for params in config["model_args"].split(","))
if "pretrained" not in model_args:
continue
results[short_name]["model_name"] = model_args["pretrained"]
results[short_name]["compute_dtype"] = model_args.get("dtype", None)
results[short_name]["quantization"] = None
if "load_in_8bit" in model_args:
results[short_name]["quantization"] = "8-bit"
elif "load_in_4bit" in model_args:
results[short_name]["quantization"] = "4-bit"
results[short_name]["num_parameters"] = get_num_parameters(model_args["pretrained"])
pprint(results)
if __name__ == '__main__':
main() |