Add safetensors shape inspector

tools/inspect_safetensors.py

@@ -1,40 +1,22 @@
-"""Report safetensors metadata without materializing model weights."""
+"""Print safetensors key shapes."""
 
 import argparse
-import json
-import struct
-from collections import Counter
-from pathlib import Path
+
+from safetensors import safe_open
 
 
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("checkpoint", type=Path)
-    parser.add_argument("--output", type=Path)
-    args = parser.parse_args()
+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 args.checkpoint.open("rb") as file:
-        header_size = struct.unpack("<Q", file.read(8))[0]
-        header = json.loads(file.read(header_size))
-
-    tensors = {name: spec for name, spec in header.items() if name != "__metadata__"}
-    report = {
-        "checkpoint": str(args.checkpoint),
-        "metadata": header.get("__metadata__", {}),
-        "tensor_count": len(tensors),
-        "dtypes": dict(sorted(Counter(spec["dtype"] for spec in tensors.values()).items())),
-        "tensors": {
-            name: {key: spec[key] for key in ("dtype", "shape", "data_offsets") if key in spec}
-            for name, spec in tensors.items()
-        },
-    }
-    rendered = json.dumps(report, indent=2, sort_keys=True)
-    if args.output:
-        args.output.parent.mkdir(parents=True, exist_ok=True)
-        args.output.write_text(rendered + "\n", encoding="utf-8")
-    else:
-        print(rendered)
-
-
-if __name__ == "__main__":
-    main()
+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())