291 lines
9.3 KiB
Python
Executable File
291 lines
9.3 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import re
|
|
import struct
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
from safetensors import safe_open
|
|
|
|
|
|
MODEL_SAFETENSORS_FILE = "model.safetensors"
|
|
MODEL_SAFETENSORS_INDEX = "model.safetensors.index.json"
|
|
|
|
DTYPE_SIZES = {
|
|
"F64": 8, "I64": 8, "U64": 8,
|
|
"F32": 4, "I32": 4, "U32": 4,
|
|
"F16": 2, "BF16": 2, "I16": 2, "U16": 2,
|
|
"I8": 1, "U8": 1, "BOOL": 1,
|
|
"F8_E4M3": 1, "F8_E5M2": 1,
|
|
}
|
|
|
|
SIZE_UNITS = ['B', 'KB', 'MB', 'GB', 'TB']
|
|
|
|
|
|
def get_weight_map(model_path: Path) -> Optional[dict[str, str]]:
|
|
index_file = model_path / MODEL_SAFETENSORS_INDEX
|
|
|
|
if index_file.exists():
|
|
with open(index_file, 'r') as f:
|
|
index = json.load(f)
|
|
return index.get("weight_map", {})
|
|
|
|
return None
|
|
|
|
|
|
def get_all_tensor_names(model_path: Path) -> list[str]:
|
|
weight_map = get_weight_map(model_path)
|
|
|
|
if weight_map is not None:
|
|
return list(weight_map.keys())
|
|
|
|
single_file = model_path / MODEL_SAFETENSORS_FILE
|
|
if single_file.exists():
|
|
try:
|
|
with safe_open(single_file, framework="pt", device="cpu") as f:
|
|
return list(f.keys())
|
|
except Exception as e:
|
|
print(f"Error reading {single_file}: {e}")
|
|
sys.exit(1)
|
|
|
|
print(f"Error: No safetensors files found in {model_path}")
|
|
sys.exit(1)
|
|
|
|
|
|
def find_tensor_file(model_path: Path, tensor_name: str) -> Optional[str]:
|
|
weight_map = get_weight_map(model_path)
|
|
|
|
if weight_map is not None:
|
|
return weight_map.get(tensor_name)
|
|
|
|
single_file = model_path / MODEL_SAFETENSORS_FILE
|
|
if single_file.exists():
|
|
return single_file.name
|
|
|
|
return None
|
|
|
|
|
|
def read_safetensors_header(file_path: Path) -> dict:
|
|
with open(file_path, 'rb') as f:
|
|
header_size = struct.unpack('<Q', f.read(8))[0]
|
|
return json.loads(f.read(header_size))
|
|
|
|
|
|
def get_tensor_size_bytes(tensor_meta: dict) -> int:
|
|
offsets = tensor_meta.get("data_offsets")
|
|
if offsets and len(offsets) == 2:
|
|
return offsets[1] - offsets[0]
|
|
n_elements = 1
|
|
for d in tensor_meta.get("shape", []):
|
|
n_elements *= d
|
|
return n_elements * DTYPE_SIZES.get(tensor_meta.get("dtype", "F32"), 4)
|
|
|
|
|
|
def format_size(size_bytes: int) -> str:
|
|
val = float(size_bytes)
|
|
for unit in SIZE_UNITS[:-1]:
|
|
if val < 1024.0:
|
|
return f"{val:.2f} {unit}"
|
|
val /= 1024.0
|
|
return f"{val:.2f} {SIZE_UNITS[-1]}"
|
|
|
|
|
|
def get_all_tensor_metadata(model_path: Path) -> dict[str, dict]:
|
|
weight_map = get_weight_map(model_path)
|
|
|
|
if weight_map is not None:
|
|
file_to_tensors: dict[str, list[str]] = {}
|
|
for tensor_name, file_name in weight_map.items():
|
|
file_to_tensors.setdefault(file_name, []).append(tensor_name)
|
|
|
|
all_metadata: dict[str, dict] = {}
|
|
for file_name, tensor_names in file_to_tensors.items():
|
|
try:
|
|
header = read_safetensors_header(model_path / file_name)
|
|
for tensor_name in tensor_names:
|
|
if tensor_name in header:
|
|
all_metadata[tensor_name] = header[tensor_name]
|
|
except Exception as e:
|
|
print(f"Warning: Could not read header from {file_name}: {e}", file=sys.stderr)
|
|
return all_metadata
|
|
|
|
single_file = model_path / MODEL_SAFETENSORS_FILE
|
|
if single_file.exists():
|
|
try:
|
|
header = read_safetensors_header(single_file)
|
|
return {k: v for k, v in header.items() if k != "__metadata__"}
|
|
except Exception as e:
|
|
print(f"Error reading {single_file}: {e}")
|
|
sys.exit(1)
|
|
|
|
print(f"Error: No safetensors files found in {model_path}")
|
|
sys.exit(1)
|
|
|
|
|
|
def normalize_tensor_name(tensor_name: str) -> str:
|
|
normalized = re.sub(r'\.\d+\.', '.#.', tensor_name)
|
|
normalized = re.sub(r'\.\d+$', '.#', normalized)
|
|
return normalized
|
|
|
|
|
|
def list_all_tensors(
|
|
model_path: Path,
|
|
short: bool = False,
|
|
show_sizes: bool = False,
|
|
):
|
|
tensor_names = get_all_tensor_names(model_path)
|
|
|
|
metadata: Optional[dict[str, dict]] = None
|
|
if show_sizes:
|
|
metadata = get_all_tensor_metadata(model_path)
|
|
|
|
total_bytes = 0
|
|
|
|
if short:
|
|
seen: dict[str, str] = {}
|
|
for tensor_name in sorted(tensor_names):
|
|
normalized = normalize_tensor_name(tensor_name)
|
|
if normalized not in seen:
|
|
seen[normalized] = tensor_name
|
|
display_pairs = list(sorted(seen.items()))
|
|
name_width = max((len(n) for n, _ in display_pairs), default=0)
|
|
for normalized, first_name in display_pairs:
|
|
if metadata and first_name in metadata:
|
|
m = metadata[first_name]
|
|
size = get_tensor_size_bytes(m)
|
|
total_bytes += size
|
|
print(f"{normalized:{name_width}} {m.get('dtype', '?'):6s} {str(m.get('shape', '')):30s} {format_size(size)}")
|
|
else:
|
|
print(normalized)
|
|
else:
|
|
name_width = max((len(n) for n in tensor_names), default=0)
|
|
for tensor_name in sorted(tensor_names):
|
|
if metadata and tensor_name in metadata:
|
|
m = metadata[tensor_name]
|
|
size = get_tensor_size_bytes(m)
|
|
total_bytes += size
|
|
print(f"{tensor_name:{name_width}} {m.get('dtype', '?'):6s} {str(m.get('shape', '')):30s} {format_size(size)}")
|
|
else:
|
|
print(tensor_name)
|
|
|
|
if show_sizes:
|
|
print(f"\nTotal: {format_size(total_bytes)}")
|
|
|
|
|
|
def print_tensor_info(model_path: Path, tensor_name: str, num_values: Optional[int] = None):
|
|
tensor_file = find_tensor_file(model_path, tensor_name)
|
|
|
|
if tensor_file is None:
|
|
print(f"Error: Could not find tensor '{tensor_name}' in model index")
|
|
print(f"Model path: {model_path}")
|
|
sys.exit(1)
|
|
|
|
file_path = model_path / tensor_file
|
|
|
|
try:
|
|
header = read_safetensors_header(file_path)
|
|
tensor_meta = header.get(tensor_name, {})
|
|
dtype_str = tensor_meta.get("dtype")
|
|
|
|
with safe_open(file_path, framework="pt", device="cpu") as f:
|
|
if tensor_name in f.keys():
|
|
tensor_slice = f.get_slice(tensor_name)
|
|
shape = tensor_slice.get_shape()
|
|
print(f"Tensor: {tensor_name}")
|
|
print(f"File: {tensor_file}")
|
|
print(f"Shape: {shape}")
|
|
if dtype_str:
|
|
print(f"Dtype: {dtype_str}")
|
|
if tensor_meta:
|
|
print(f"Size: {format_size(get_tensor_size_bytes(tensor_meta))}")
|
|
if num_values is not None:
|
|
tensor = f.get_tensor(tensor_name)
|
|
if not dtype_str:
|
|
print(f"Dtype: {tensor.dtype}")
|
|
flat = tensor.flatten()
|
|
n = min(num_values, flat.numel())
|
|
print(f"Values: {flat[:n].tolist()}")
|
|
else:
|
|
print(f"Error: Tensor '{tensor_name}' not found in {tensor_file}")
|
|
sys.exit(1)
|
|
|
|
except FileNotFoundError:
|
|
print(f"Error: The file '{file_path}' was not found.")
|
|
sys.exit(1)
|
|
except Exception as e:
|
|
print(f"An error occurred: {e}")
|
|
sys.exit(1)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Print tensor information from a safetensors model"
|
|
)
|
|
parser.add_argument(
|
|
"tensor_name",
|
|
nargs="?",
|
|
help="Name of the tensor to inspect"
|
|
)
|
|
parser.add_argument(
|
|
"-m", "--model-path",
|
|
type=Path,
|
|
help="Path to the model directory (default: MODEL_PATH environment variable)"
|
|
)
|
|
parser.add_argument(
|
|
"-l", "--list-all-short",
|
|
action="store_true",
|
|
help="List unique tensor patterns (layer numbers replaced with #)"
|
|
)
|
|
parser.add_argument(
|
|
"-la", "--list-all",
|
|
action="store_true",
|
|
help="List all tensor names with actual layer numbers"
|
|
)
|
|
parser.add_argument(
|
|
"-n", "--num-values",
|
|
nargs="?",
|
|
const=10,
|
|
default=None,
|
|
type=int,
|
|
metavar="N",
|
|
help="Print the first N values of the tensor flattened (default: 10 if flag is given without a number)"
|
|
)
|
|
parser.add_argument(
|
|
"-s", "--sizes",
|
|
action="store_true",
|
|
help="Show dtype, shape, and size for each tensor when listing"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
model_path = args.model_path
|
|
if model_path is None:
|
|
model_path_str = os.environ.get("MODEL_PATH")
|
|
if model_path_str is None:
|
|
print("Error: --model-path not provided and MODEL_PATH environment variable not set")
|
|
sys.exit(1)
|
|
model_path = Path(model_path_str)
|
|
|
|
if not model_path.exists():
|
|
print(f"Error: Model path does not exist: {model_path}")
|
|
sys.exit(1)
|
|
|
|
if not model_path.is_dir():
|
|
print(f"Error: Model path is not a directory: {model_path}")
|
|
sys.exit(1)
|
|
|
|
if args.list_all_short or args.list_all:
|
|
list_all_tensors(model_path, short=args.list_all_short, show_sizes=args.sizes)
|
|
else:
|
|
if args.tensor_name is None:
|
|
print("Error: tensor_name is required when not using --list-all-short or --list-all")
|
|
sys.exit(1)
|
|
print_tensor_info(model_path, args.tensor_name, args.num_values)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|