22 lines
672 B
Python
22 lines
672 B
Python
"""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())
|