"""Print safetensors key shapes."""

import argparse

from safetensors import safe_open


parser = argparse.ArgumentParser()
parser.add_argument("path")
parser.add_argument("--head", type=int, default=80)
parser.add_argument("--tail", type=int, default=30)
args = parser.parse_args()

with safe_open(args.path, framework="pt", device="cpu") as checkpoint:
    keys = list(checkpoint.keys())
    print({"path": args.path, "keys": len(keys)})
    for key in keys[:args.head]:
        print(key, checkpoint.get_slice(key).get_shape())
    if args.tail:
        print("last")
        for key in keys[-args.tail:]:
            print(key, checkpoint.get_slice(key).get_shape())