149 lines
4.7 KiB
Python
149 lines
4.7 KiB
Python
import argparse
|
|
import os
|
|
|
|
from collections import Counter
|
|
from safetensors.torch import safe_open
|
|
|
|
|
|
def format_bytes(num_bytes: int) -> str:
|
|
units = ["B", "KB", "MB", "GB", "TB"]
|
|
size = float(num_bytes)
|
|
for unit in units:
|
|
if size < 1024.0 or unit == units[-1]:
|
|
if unit == "B":
|
|
return f"{int(size)} {unit}"
|
|
return f"{size:.2f} {unit}"
|
|
size /= 1024.0
|
|
return f"{num_bytes} B"
|
|
|
|
|
|
def estimate_tensor_bytes(tensor) -> int:
|
|
return int(tensor.numel()) * int(tensor.element_size())
|
|
|
|
|
|
def common_prefix(a: str, b: str) -> str:
|
|
n = min(len(a), len(b))
|
|
i = 0
|
|
while i < n and a[i] == b[i]:
|
|
i += 1
|
|
return a[:i]
|
|
|
|
|
|
def derive_schema_prefixes(keys, max_depth: int = 4) -> Counter:
|
|
separators = [".", "/", ":", "_"]
|
|
counts = Counter()
|
|
|
|
for key in keys:
|
|
parts = [key]
|
|
for sep in separators:
|
|
new_parts = []
|
|
for p in parts:
|
|
new_parts.extend(p.split(sep))
|
|
parts = new_parts
|
|
|
|
parts = [p for p in parts if p]
|
|
if not parts:
|
|
continue
|
|
|
|
depth = min(max_depth, len(parts))
|
|
prefix = ".".join(parts[:depth])
|
|
counts[prefix] += 1
|
|
|
|
return counts
|
|
|
|
|
|
def inspect_and_write_report(
|
|
path: str,
|
|
limit: int = 30,
|
|
schema_depth: int = 4,
|
|
schema_top: int = 25,
|
|
) -> None:
|
|
if not os.path.isfile(path):
|
|
raise FileNotFoundError(path)
|
|
|
|
base_name = os.path.splitext(os.path.basename(path))[0]
|
|
output_path = os.path.join(
|
|
os.path.dirname(path),
|
|
f"{base_name}.schematics.txt"
|
|
)
|
|
|
|
file_size = os.path.getsize(path)
|
|
|
|
with open(output_path, "w", encoding="utf-8") as out:
|
|
out.write("Tensor Container Schematic Report\n")
|
|
out.write("=" * 80 + "\n\n")
|
|
out.write(f"Source file: {path}\n")
|
|
out.write(f"File size: {format_bytes(file_size)}\n\n")
|
|
|
|
with safe_open(path, framework="pt", device="cpu") as f:
|
|
keys = list(f.keys())
|
|
total = len(keys)
|
|
|
|
out.write(f"Total tensors: {total}\n")
|
|
out.write(f"Showing first {min(limit, total)} entries\n\n")
|
|
|
|
for idx, key in enumerate(keys[:limit], start=1):
|
|
tensor = f.get_tensor(key)
|
|
|
|
shape = tuple(tensor.shape)
|
|
ndim = int(tensor.ndim)
|
|
numel = int(tensor.numel())
|
|
dtype = str(tensor.dtype).replace("torch.", "")
|
|
bytes_est = estimate_tensor_bytes(tensor)
|
|
|
|
out.write(f"{idx:02d}. {key}\n")
|
|
out.write(f" shape: {shape}\n")
|
|
out.write(f" ndim: {ndim}\n")
|
|
out.write(f" numel: {numel}\n")
|
|
out.write(f" dtype: {dtype}\n")
|
|
out.write(f" bytes: {format_bytes(bytes_est)}\n")
|
|
out.write("-" * 80 + "\n")
|
|
|
|
if total > 0:
|
|
out.write("\nLexical Schema Summary (Purely Name-Based)\n")
|
|
out.write("=" * 80 + "\n")
|
|
out.write(
|
|
f"Grouping token depth: {schema_depth} | "
|
|
f"Top groups shown: {schema_top}\n\n"
|
|
)
|
|
|
|
counts = derive_schema_prefixes(keys, max_depth=schema_depth)
|
|
for prefix, count in counts.most_common(schema_top):
|
|
out.write(f"{count:6d} {prefix}\n")
|
|
|
|
keys_sorted = sorted(keys)
|
|
shared = keys_sorted[0]
|
|
for k in keys_sorted[1:]:
|
|
shared = common_prefix(shared, k)
|
|
if not shared:
|
|
break
|
|
|
|
out.write("\nCommon character prefix across all keys:\n")
|
|
out.write(shared if shared else "(none)")
|
|
out.write("\n")
|
|
|
|
print(f"Report written to:\n {output_path}")
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(
|
|
description="Generate a schematic report for a .safetensors file."
|
|
)
|
|
parser.add_argument("path", type=str, help="Path to .safetensors file")
|
|
parser.add_argument("--limit", type=int, default=30, help="Number of keys to list (default: 30)")
|
|
parser.add_argument("--schema-depth", type=int, default=4, help="Token depth for schema grouping")
|
|
parser.add_argument("--schema-top", type=int, default=25, help="Number of schema groups to show")
|
|
|
|
args = parser.parse_args()
|
|
|
|
inspect_and_write_report(
|
|
path=args.path,
|
|
limit=args.limit,
|
|
schema_depth=args.schema_depth,
|
|
schema_top=args.schema_top,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|