#!/usr/bin/env python3
"""CLI to open and explore .safetensors files."""
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
import numpy as np
from safetensors import safe_open
DTYPE_SIZE = {
"BOOL": 1,
"U8": 1,
"I8": 1,
"F8_E4M3": 1,
"F8_E5M2": 1,
"F16": 2,
"BF16": 2,
"I16": 2,
"U16": 2,
"F32": 4,
"I32": 4,
"U32": 4,
"F64": 8,
"I64": 8,
"U64": 8,
}
def human_size(n: int) -> str:
units = ["B", "KB", "MB", "GB", "TB"]
size = float(n)
for unit in units:
if size < 1024 or unit == units[-1]:
return f"{size:.1f} {unit}" if unit != "B" else f"{int(size)} B"
size /= 1024
return f"{n} B"
def tensor_nbytes(dtype: str, shape: list[int]) -> int:
elems = int(np.prod(shape)) if shape else 1
return elems * DTYPE_SIZE.get(dtype.upper(), 4)
def collect_info(path: Path, framework: str = "np") -> dict:
tensors: list[dict] = []
metadata: dict | None = None
with safe_open(str(path), framework=framework) as f:
metadata = f.metadata()
for name in f.keys():
slice_ = f.get_slice(name)
shape = list(slice_.get_shape())
dtype = str(slice_.get_dtype())
nbytes = tensor_nbytes(dtype, shape)
tensors.append(
{
"name": name,
"dtype": dtype,
"shape": shape,
"ndim": len(shape),
"numel": int(np.prod(shape)) if shape else 1,
"nbytes": nbytes,
}
)
tensors.sort(key=lambda t: t["name"])
total_bytes = sum(t["nbytes"] for t in tensors)
return {
"path": str(path.resolve()),
"file_size": path.stat().st_size,
"tensor_count": len(tensors),
"total_tensor_bytes": total_bytes,
"metadata": metadata or {},
"tensors": tensors,
}
def filter_tensors(tensors: list[dict], pattern: str | None) -> list[dict]:
if not pattern:
return tensors
pat = pattern.lower()
return [t for t in tensors if pat in t["name"].lower()]
def print_summary(info: dict, pattern: str | None = None) -> None:
tensors = filter_tensors(info["tensors"], pattern)
print(f"File: {info['path']}")
print(f"File size: {human_size(info['file_size'])} ({info['file_size']:,} bytes)")
print(f"Tensors: {len(tensors)}" + (f" / {info['tensor_count']} (filtered)" if pattern else ""))
print(f"Tensor bytes: {human_size(info['total_tensor_bytes'])}")
meta = info["metadata"]
if meta:
print(f"Metadata keys: {len(meta)}")
for k, v in meta.items():
preview = str(v)
if len(preview) > 80:
preview = preview[:77] + "..."
print(f" {k}: {preview}")
else:
print("Metadata: (none)")
print()
def print_table(tensors: list[dict], sort_by: str = "name") -> None:
if not tensors:
print("No tensors matched.")
return
key_map = {
"name": lambda t: t["name"],
"size": lambda t: t["nbytes"],
"numel": lambda t: t["numel"],
"dtype": lambda t: t["dtype"],
}
rows = sorted(tensors, key=key_map.get(sort_by, key_map["name"]), reverse=sort_by in {"size", "numel"})
name_w = max(4, max(len(t["name"]) for t in rows))
dtype_w = max(5, max(len(t["dtype"]) for t in rows))
shape_w = max(5, max(len(str(t["shape"])) for t in rows))
header = f"{'NAME':<{name_w}} {'DTYPE':<{dtype_w}} {'SHAPE':<{shape_w}} {'NUMEL':>12} {'SIZE':>10}"
print(header)
print("-" * len(header))
for t in rows:
print(
f"{t['name']:<{name_w}} {t['dtype']:<{dtype_w}} {str(t['shape']):<{shape_w}} "
f"{t['numel']:>12,} {human_size(t['nbytes']):>10}"
)
def inspect_tensor(path: Path, name: str, sample: int = 8, framework: str = "np") -> None:
with safe_open(str(path), framework=framework) as f:
keys = list(f.keys())
if name not in keys:
matches = [k for k in keys if name.lower() in k.lower()]
print(f"Tensor '{name}' not found.", file=sys.stderr)
if matches:
print("Did you mean:", file=sys.stderr)
for m in matches[:10]:
print(f" - {m}", file=sys.stderr)
sys.exit(1)
tensor = f.get_tensor(name)
arr = np.asarray(tensor)
print(f"Name: {name}")
print(f"Dtype: {arr.dtype}")
print(f"Shape: {list(arr.shape)}")
print(f"Numel: {arr.size:,}")
print(f"Size: {human_size(arr.nbytes)}")
if arr.size == 0:
print("Empty tensor.")
return
flat = arr.astype(np.float64).ravel()
print(f"Min: {flat.min()}")
print(f"Max: {flat.max()}")
print(f"Mean: {flat.mean()}")
print(f"Std: {flat.std()}")
print(f"AbsMax: {np.abs(flat).max()}")
n = min(sample, flat.size)
print(f"Sample ({n} values):")
print(f" {flat[:n].tolist()}")
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
description="Explore the content of a .safetensors file (CLI).",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""\
examples:
%(prog)s model.safetensors
%(prog)s model.safetensors --filter attn
%(prog)s model.safetensors --sort size
%(prog)s model.safetensors --inspect model.layers.0.weight
%(prog)s model.safetensors --json
""",
)
parser.add_argument("file", type=Path, help="Path to a .safetensors file")
parser.add_argument("-f", "--filter", metavar="SUBSTR", help="Only show tensors whose name contains SUBSTR")
parser.add_argument(
"-s",
"--sort",
choices=["name", "size", "numel", "dtype"],
default="name",
help="Sort tensor table (default: name)",
)
parser.add_argument("-i", "--inspect", metavar="NAME", help="Load one tensor and print stats + sample values")
parser.add_argument("--sample", type=int, default=8, help="Number of sample values when inspecting (default: 8)")
parser.add_argument("--json", action="store_true", help="Dump structured summary as JSON")
parser.add_argument("--metadata-only", action="store_true", help="Only print file metadata")
return parser
def main(argv: list[str] | None = None) -> int:
parser = build_parser()
args = parser.parse_args(argv)
path: Path = args.file
if not path.exists():
print(f"Error: file not found: {path}", file=sys.stderr)
return 1
if not path.is_file():
print(f"Error: not a file: {path}", file=sys.stderr)
return 1
try:
if args.inspect:
inspect_tensor(path, args.inspect, sample=args.sample)
return 0
info = collect_info(path)
if args.json:
if args.filter:
info = {**info, "tensors": filter_tensors(info["tensors"], args.filter)}
print(json.dumps(info, indent=2, ensure_ascii=False))
return 0
print_summary(info, pattern=args.filter)
if args.metadata_only:
return 0
print_table(filter_tensors(info["tensors"], args.filter), sort_by=args.sort)
return 0
except Exception as exc: # noqa: BLE001 — CLI boundary
print(f"Error: {exc}", file=sys.stderr)
return 1
if __name__ == "__main__":
raise SystemExit(main())
Requirement
safetensors>=0.4.0
numpy>=1.24.0
Articles
