1106 lines
39 KiB
Python
1106 lines
39 KiB
Python
#!/usr/bin/env python3
|
|
"""Print a summary of a Hugging Face checkpoint.
|
|
|
|
Weight bytes are not read. config.json, the tokenizer files, and safetensors
|
|
headers are enough for the architecture, quantization, and tensor list.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import re
|
|
import struct
|
|
import sys
|
|
from collections import Counter
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
# Same round-trip check as nvfp4gguf.bundle.chat_template_has_tools.
|
|
_TOOL_RESULT_MARKERS = (
|
|
"tool_response",
|
|
"tool_results",
|
|
"role'] == 'tool'",
|
|
"role'] == \"tool\"",
|
|
"role\"] == 'tool'",
|
|
"role\"] == \"tool\"",
|
|
"message.role == 'tool'",
|
|
"message.role == \"tool\"",
|
|
"ipython",
|
|
)
|
|
|
|
_NESTED_CONFIGS = (
|
|
"text_config",
|
|
"llm_config",
|
|
"language_config",
|
|
"thinker_config",
|
|
"decoder",
|
|
)
|
|
|
|
_INDEX_NAMES = (
|
|
"model.safetensors.index.json",
|
|
"diffusion_pytorch_model.safetensors.index.json",
|
|
)
|
|
|
|
_MAX_HEADER = 64 * 1024 * 1024
|
|
|
|
|
|
@dataclass
|
|
class TensorInfo:
|
|
name: str
|
|
dtype: str
|
|
shape: list[int]
|
|
nbytes: int | None
|
|
file: str
|
|
|
|
|
|
@dataclass
|
|
class ShardInfo:
|
|
path: Path
|
|
header_size: int
|
|
data_end: int
|
|
metadata: dict
|
|
|
|
|
|
@dataclass
|
|
class Checkpoint:
|
|
path: Path
|
|
root: Path
|
|
config: dict
|
|
config_name: str
|
|
generation: dict
|
|
tokenizer: dict
|
|
quant: dict
|
|
hf_quant: dict | None
|
|
recipe: str | None
|
|
base_model: str | None
|
|
tensors: list[TensorInfo]
|
|
shards: list[ShardInfo]
|
|
shard_total: int
|
|
partial: bool
|
|
index_meta: dict
|
|
template: str
|
|
template_where: str
|
|
template_count: int
|
|
nvfp4: bool
|
|
comfy: bool
|
|
|
|
|
|
def format_size(num_bytes: int) -> str:
|
|
size = float(num_bytes)
|
|
for unit in ("B", "KiB", "MiB", "GiB", "TiB"):
|
|
if size < 1024 or unit == "TiB":
|
|
return f"{int(size)} B" if unit == "B" else f"{size:.2f} {unit}"
|
|
size /= 1024
|
|
return f"{size:.2f} TiB"
|
|
|
|
|
|
def format_params(count: int) -> str:
|
|
if count >= 1_000_000_000:
|
|
return f"{count} ({count / 1_000_000_000:.2f}B)"
|
|
if count >= 1_000_000:
|
|
return f"{count} ({count / 1_000_000:.2f}M)"
|
|
return str(count)
|
|
|
|
|
|
def format_number(value) -> str:
|
|
if isinstance(value, bool):
|
|
return "true" if value else "false"
|
|
if isinstance(value, int):
|
|
return str(value)
|
|
if isinstance(value, float):
|
|
if value.is_integer() and abs(value) < 1e15:
|
|
return str(int(value))
|
|
return f"{value:.6g}"
|
|
return _preview(value)
|
|
|
|
|
|
def _preview(value, limit: int = 160) -> str:
|
|
if isinstance(value, str):
|
|
flat = value.replace("\n", "\\n")
|
|
if len(flat) > limit:
|
|
return f"{flat[:limit]}... ({len(value)} chars)"
|
|
return flat
|
|
if isinstance(value, dict):
|
|
text = json.dumps(value, ensure_ascii=False, default=str, separators=(", ", ": "))
|
|
if len(text) > limit:
|
|
return f"{text[:limit]}... ({len(text)} chars)"
|
|
return text
|
|
if isinstance(value, list):
|
|
text = ", ".join(_preview(item, 40) for item in value[:12])
|
|
if len(value) > 12:
|
|
text += f", ... ({len(value)})"
|
|
return "[" + text + "]"
|
|
if isinstance(value, bool):
|
|
return "true" if value else "false"
|
|
if isinstance(value, float):
|
|
return format_number(value)
|
|
if value is None:
|
|
return "null"
|
|
return str(value)
|
|
|
|
|
|
def _line(label: str, value) -> None:
|
|
print(f"{label:<16}{value}")
|
|
|
|
|
|
def read_json(path: Path):
|
|
try:
|
|
text = path.read_text(encoding="utf-8")
|
|
except (OSError, UnicodeError) as exc:
|
|
raise ValueError(f"{path}: {exc}") from exc
|
|
try:
|
|
return json.loads(text)
|
|
except json.JSONDecodeError as exc:
|
|
raise ValueError(f"{path}: invalid JSON ({exc.msg})") from exc
|
|
|
|
|
|
def read_header(path: Path) -> tuple[dict, int]:
|
|
with path.open("rb") as handle:
|
|
raw = handle.read(8)
|
|
if len(raw) != 8:
|
|
raise ValueError(f"{path} is not a safetensors file")
|
|
header_size = struct.unpack("<Q", raw)[0]
|
|
if header_size <= 0 or header_size > _MAX_HEADER:
|
|
raise ValueError(f"{path} is not a safetensors file")
|
|
payload = handle.read(header_size)
|
|
if len(payload) != header_size:
|
|
raise ValueError(f"{path} has a truncated safetensors header")
|
|
try:
|
|
header = json.loads(payload)
|
|
except json.JSONDecodeError as exc:
|
|
raise ValueError(f"{path} has an unreadable safetensors header") from exc
|
|
if not isinstance(header, dict):
|
|
raise ValueError(f"{path} has an unreadable safetensors header")
|
|
return header, header_size
|
|
|
|
|
|
def _sections(config: dict) -> list[dict]:
|
|
found = [config]
|
|
for key in _NESTED_CONFIGS:
|
|
nested = config.get(key)
|
|
if isinstance(nested, dict):
|
|
found.append(nested)
|
|
return found
|
|
|
|
|
|
def _present(config: dict, keys: tuple[str, ...]):
|
|
for section in _sections(config):
|
|
for key in keys:
|
|
if key in section and section[key] not in (None, ""):
|
|
return section[key]
|
|
return None
|
|
|
|
|
|
def _positive_int(config: dict, keys: tuple[str, ...], limit: int = 10**12) -> int | None:
|
|
value = _present(config, keys)
|
|
if isinstance(value, bool):
|
|
return None
|
|
if isinstance(value, int) and 0 < value <= limit:
|
|
return value
|
|
if isinstance(value, float) and value.is_integer() and 0 < value <= limit:
|
|
return int(value)
|
|
return None
|
|
|
|
|
|
def _text(value) -> str:
|
|
return value if isinstance(value, str) and value else ""
|
|
|
|
|
|
def chat_template_has_tools(template: str) -> bool:
|
|
if not template:
|
|
return False
|
|
if "tools" not in template and "tool_call" not in template:
|
|
return False
|
|
if "tool_calls" not in template and "assistant_tool_call" not in template:
|
|
return False
|
|
return any(marker in template for marker in _TOOL_RESULT_MARKERS)
|
|
|
|
|
|
def _template_parts(value) -> list[str]:
|
|
if isinstance(value, str):
|
|
return [value] if value.strip() else []
|
|
if not isinstance(value, list):
|
|
return []
|
|
parts = []
|
|
for item in value:
|
|
if isinstance(item, str) and item.strip():
|
|
parts.append(item)
|
|
elif isinstance(item, dict) and isinstance(item.get("template"), str) and item["template"].strip():
|
|
parts.append(item["template"])
|
|
return parts
|
|
|
|
|
|
def load_chat_template(root: Path, tokenizer: dict) -> tuple[str, str, int]:
|
|
parts = _template_parts(tokenizer.get("chat_template")) if tokenizer else []
|
|
if parts:
|
|
count = len(tokenizer.get("chat_template")) if isinstance(tokenizer.get("chat_template"), list) else 1
|
|
return "\n".join(parts), "tokenizer_config.json", count
|
|
path = root / "chat_template.jinja"
|
|
if path.is_file():
|
|
try:
|
|
return path.read_text(encoding="utf-8"), path.name, 1
|
|
except (OSError, UnicodeError) as exc:
|
|
raise ValueError(f"{path}: {exc}") from exc
|
|
return "", "", 0
|
|
|
|
|
|
def _token_text(value) -> str:
|
|
if isinstance(value, str):
|
|
return value
|
|
if isinstance(value, dict) and isinstance(value.get("content"), str):
|
|
return value["content"]
|
|
return ""
|
|
|
|
|
|
def token_label(token_id, tokenizer: dict) -> str:
|
|
if isinstance(token_id, list):
|
|
return ", ".join(token_label(item, tokenizer) for item in token_id)
|
|
name = ""
|
|
decoder = tokenizer.get("added_tokens_decoder") if tokenizer else None
|
|
if isinstance(decoder, dict):
|
|
entry = decoder.get(str(token_id))
|
|
if isinstance(entry, dict):
|
|
name = _token_text(entry.get("content"))
|
|
if name:
|
|
return f"{token_id} ({name})"
|
|
return str(token_id)
|
|
|
|
|
|
def _token_id(source: dict, key: str):
|
|
if not source or key not in source:
|
|
return None
|
|
value = source.get(key)
|
|
if isinstance(value, bool) or value is None or value == "":
|
|
return None
|
|
if isinstance(value, int):
|
|
return value
|
|
if isinstance(value, list) and value and all(isinstance(item, int) and not isinstance(item, bool) for item in value):
|
|
return value
|
|
return None
|
|
|
|
|
|
def read_base_model(path: Path) -> str | None:
|
|
if not path.is_file():
|
|
return None
|
|
try:
|
|
text = path.read_text(encoding="utf-8", errors="replace")[:65536]
|
|
except OSError:
|
|
return None
|
|
if not text.startswith("---"):
|
|
return None
|
|
end = text.find("\n---", 3)
|
|
if end < 0:
|
|
return None
|
|
values: list[str] = []
|
|
grabbing = False
|
|
for line in text[3:end].splitlines():
|
|
if not grabbing:
|
|
if line.startswith("base_model:"):
|
|
rest = line.split(":", 1)[1].strip().strip("'\"")
|
|
if rest:
|
|
values.append(rest)
|
|
else:
|
|
grabbing = True
|
|
continue
|
|
stripped = line.strip()
|
|
if stripped.startswith("- "):
|
|
values.append(stripped[2:].strip().strip("'\""))
|
|
continue
|
|
break
|
|
cleaned = [value for value in values if value]
|
|
return ", ".join(cleaned) if cleaned else None
|
|
|
|
|
|
def read_recipe_scheme(path: Path) -> str | None:
|
|
if not path.is_file():
|
|
return None
|
|
try:
|
|
text = path.read_text(encoding="utf-8", errors="replace")[:65536]
|
|
except OSError:
|
|
return None
|
|
for line in text.splitlines():
|
|
stripped = line.strip()
|
|
if stripped.startswith("scheme:"):
|
|
value = stripped.split(":", 1)[1].strip().strip("'\"")
|
|
return value or None
|
|
return None
|
|
|
|
|
|
def load_quant(config: dict, root: Path) -> tuple[dict, dict | None]:
|
|
raw = config.get("quantization_config")
|
|
quant = dict(raw) if isinstance(raw, dict) else {}
|
|
path = root / "hf_quant_config.json"
|
|
if not path.is_file():
|
|
return quant, None
|
|
extra = read_json(path)
|
|
if not isinstance(extra, dict):
|
|
raise ValueError(f"{path} is not a JSON object")
|
|
block = extra.get("quantization") if isinstance(extra.get("quantization"), dict) else {}
|
|
producer = extra.get("producer") if isinstance(extra.get("producer"), dict) else {}
|
|
for key, value in block.items():
|
|
quant.setdefault(key, value)
|
|
if "quant_method" not in quant and isinstance(producer.get("name"), str) and producer["name"]:
|
|
quant["quant_method"] = producer["name"].lower()
|
|
return quant, extra
|
|
|
|
|
|
def quant_is_nvfp4(quant: dict) -> bool:
|
|
algo = str(quant.get("quant_algo") or "")
|
|
fmt = str(quant.get("format") or "")
|
|
if algo in {"NVFP4", "W4A16_NVFP4"} or algo.endswith("NVFP4"):
|
|
return True
|
|
if "nvfp4" in fmt.lower():
|
|
return True
|
|
groups = quant.get("config_groups")
|
|
if isinstance(groups, dict):
|
|
for group in groups.values():
|
|
if isinstance(group, dict) and "nvfp4" in str(group.get("format") or "").lower():
|
|
return True
|
|
layers = quant.get("quantized_layers")
|
|
if isinstance(layers, dict):
|
|
for entry in layers.values():
|
|
if isinstance(entry, dict) and str(entry.get("quant_algo") or "").endswith("NVFP4"):
|
|
return True
|
|
return False
|
|
|
|
|
|
def _weight_schemes(quant: dict) -> list[tuple]:
|
|
found = []
|
|
groups = quant.get("config_groups")
|
|
if isinstance(groups, dict):
|
|
for group in groups.values():
|
|
if not isinstance(group, dict):
|
|
continue
|
|
weights = group.get("weights") if isinstance(group.get("weights"), dict) else {}
|
|
acts = group.get("input_activations") if isinstance(group.get("input_activations"), dict) else {}
|
|
found.append((weights.get("num_bits"), weights.get("group_size"), acts.get("num_bits")))
|
|
if found:
|
|
return found
|
|
bits = quant.get("bits", quant.get("num_bits", quant.get("w_bit")))
|
|
group_size = quant.get("group_size", quant.get("groupsize"))
|
|
if bits is not None or group_size is not None:
|
|
return [(bits, group_size, None)]
|
|
return []
|
|
|
|
|
|
def find_index(root: Path) -> Path | None:
|
|
for name in _INDEX_NAMES:
|
|
path = root / name
|
|
if path.is_file():
|
|
return path
|
|
found = sorted(path for path in root.glob("*.safetensors.index.json") if path.is_file())
|
|
return found[0] if found else None
|
|
|
|
|
|
def shards_from_index(index_path: Path, payload: dict) -> list[Path]:
|
|
weight_map = payload.get("weight_map")
|
|
if not isinstance(weight_map, dict) or not weight_map:
|
|
raise ValueError(f"{index_path} has no weight_map")
|
|
root = index_path.parent
|
|
names: list[Path] = []
|
|
seen: set[str] = set()
|
|
for name in weight_map.values():
|
|
if not isinstance(name, str) or not name or name.startswith(("/", "\\")) or ".." in Path(name).parts:
|
|
raise ValueError(f"{index_path} has an invalid weight_map entry")
|
|
if name not in seen:
|
|
seen.add(name)
|
|
names.append(root / name)
|
|
return names
|
|
|
|
|
|
def discover_shards(root: Path) -> list[Path]:
|
|
for name in ("model.safetensors", "diffusion_pytorch_model.safetensors"):
|
|
path = root / name
|
|
if path.is_file():
|
|
return [path]
|
|
return sorted(path for path in root.glob("*.safetensors") if path.is_file())
|
|
|
|
|
|
def parse_tensor(name: str, entry, file_name: str) -> tuple[TensorInfo, int]:
|
|
if not isinstance(entry, dict):
|
|
raise ValueError(f"{file_name}: tensor {name!r} has an invalid header entry")
|
|
dtype = entry.get("dtype")
|
|
if not isinstance(dtype, str) or not dtype:
|
|
dtype = "?"
|
|
shape = entry.get("shape")
|
|
if shape is None:
|
|
dims: list[int] = []
|
|
elif isinstance(shape, list) and all(isinstance(dim, int) and not isinstance(dim, bool) for dim in shape):
|
|
dims = [int(dim) for dim in shape]
|
|
else:
|
|
raise ValueError(f"{file_name}: tensor {name!r} has an invalid shape")
|
|
offsets = entry.get("data_offsets")
|
|
nbytes = None
|
|
end = 0
|
|
if isinstance(offsets, list) and len(offsets) == 2 and all(isinstance(item, int) and not isinstance(item, bool) for item in offsets):
|
|
start, end = int(offsets[0]), int(offsets[1])
|
|
if start < 0 or end < start:
|
|
raise ValueError(f"{file_name}: tensor {name!r} has invalid data offsets")
|
|
nbytes = end - start
|
|
return TensorInfo(name, dtype, dims, nbytes, file_name), end
|
|
|
|
|
|
def load_shards(paths: list[Path]) -> tuple[list[TensorInfo], list[ShardInfo]]:
|
|
tensors: list[TensorInfo] = []
|
|
shards: list[ShardInfo] = []
|
|
owners: dict[str, str] = {}
|
|
for path in paths:
|
|
header, header_size = read_header(path)
|
|
meta = header.pop("__metadata__", None)
|
|
if not isinstance(meta, dict):
|
|
meta = {}
|
|
data_end = 0
|
|
for name, entry in header.items():
|
|
if name.startswith("__"):
|
|
continue
|
|
if name in owners:
|
|
raise ValueError(f"Tensor {name!r} appears in both {owners[name]} and {path.name}")
|
|
owners[name] = path.name
|
|
tensor, end = parse_tensor(name, entry, path.name)
|
|
data_end = max(data_end, end)
|
|
tensors.append(tensor)
|
|
shards.append(ShardInfo(path, header_size, data_end, meta))
|
|
return tensors, shards
|
|
|
|
|
|
def _natural_key(name: str) -> tuple:
|
|
return tuple((0, int(part)) if part.isdigit() else (1, part) for part in re.split(r"(\d+)", name))
|
|
|
|
|
|
def tensor_type(tensor: TensorInfo, quant_method: str, nvfp4: bool, comfy: bool) -> str:
|
|
leaf = tensor.name.rsplit(".", 1)[-1]
|
|
if nvfp4 and leaf == "weight_packed":
|
|
return "NVFP4"
|
|
if comfy and leaf == "weight" and tensor.dtype == "U8":
|
|
return "NVFP4"
|
|
if quant_method.lower() in {"gptq", "awq"} and leaf == "qweight":
|
|
return quant_method.upper()
|
|
return tensor.dtype
|
|
|
|
|
|
def _numel(shape: list[int]) -> int:
|
|
count = 1
|
|
for dim in shape:
|
|
count *= dim
|
|
return count
|
|
|
|
|
|
def parameter_count(tensors: list[TensorInfo], type_of) -> int | None:
|
|
"""Logical weights and biases.
|
|
|
|
Packed NVFP4 stores two FP4 values in each byte, so those element counts
|
|
are doubled. Scale tensors are not parameters. Other packed layouts are
|
|
left uncounted rather than reported as a partial total.
|
|
"""
|
|
total = 0
|
|
seen = False
|
|
for tensor in tensors:
|
|
leaf = tensor.name.rsplit(".", 1)[-1]
|
|
packed = leaf in {"weight_packed", "qweight"}
|
|
if packed and type_of(tensor) != "NVFP4":
|
|
return None
|
|
if leaf not in {"weight", "bias", "weight_packed"}:
|
|
continue
|
|
seen = True
|
|
count = _numel(tensor.shape)
|
|
if type_of(tensor) == "NVFP4":
|
|
count *= 2
|
|
total += count
|
|
return total if seen else None
|
|
|
|
|
|
def _is_model_dir(path: Path) -> bool:
|
|
return path.name in {"config.json", "adapter_config.json"} or path.name.endswith(".safetensors.index.json")
|
|
|
|
|
|
def _load(path: Path) -> Checkpoint:
|
|
if not path.exists():
|
|
raise SystemExit(f"Not found: {path}")
|
|
if path.is_file() and path.suffix == ".gguf":
|
|
raise SystemExit(f"{path} is a GGUF file. Use gguf_info.py for GGUF files.")
|
|
if path.is_file() and _is_model_dir(path):
|
|
root = path.parent
|
|
selected: list[Path] | None = None
|
|
elif path.is_file():
|
|
if path.suffix != ".safetensors":
|
|
raise SystemExit(f"{path} is not a model directory or a safetensors file")
|
|
root = path.parent
|
|
selected = [path]
|
|
elif path.is_dir():
|
|
root = path
|
|
selected = None
|
|
else:
|
|
raise SystemExit(f"Not a file or directory: {path}")
|
|
|
|
config_path = root / "config.json"
|
|
config_name = "config.json"
|
|
if not config_path.is_file() and (root / "adapter_config.json").is_file():
|
|
config_path = root / "adapter_config.json"
|
|
config_name = config_path.name
|
|
config: dict = {}
|
|
if config_path.is_file():
|
|
loaded = read_json(config_path)
|
|
if not isinstance(loaded, dict):
|
|
raise ValueError(f"{config_path} is not a JSON object")
|
|
config = loaded
|
|
|
|
generation_path = root / "generation_config.json"
|
|
generation: dict = {}
|
|
if generation_path.is_file():
|
|
loaded = read_json(generation_path)
|
|
if isinstance(loaded, dict):
|
|
generation = loaded
|
|
|
|
tokenizer_path = root / "tokenizer_config.json"
|
|
tokenizer: dict = {}
|
|
if tokenizer_path.is_file():
|
|
loaded = read_json(tokenizer_path)
|
|
if isinstance(loaded, dict):
|
|
tokenizer = loaded
|
|
|
|
quant, hf_quant = load_quant(config, root)
|
|
recipe = read_recipe_scheme(root / "recipe.yaml")
|
|
base_model = read_base_model(root / "README.md")
|
|
repo = config.get("_name_or_path")
|
|
if not base_model and isinstance(repo, str) and "/" in repo and not repo.startswith(("/", ".")):
|
|
base_model = repo
|
|
|
|
index_path = find_index(root)
|
|
index_payload: dict = {}
|
|
weight_map: dict = {}
|
|
indexed: list[Path] = []
|
|
if index_path is not None:
|
|
loaded = read_json(index_path)
|
|
if not isinstance(loaded, dict):
|
|
raise ValueError(f"{index_path} is not a JSON object")
|
|
index_payload = loaded
|
|
raw_map = loaded.get("weight_map")
|
|
weight_map = raw_map if isinstance(raw_map, dict) else {}
|
|
indexed = shards_from_index(index_path, loaded)
|
|
if selected is None:
|
|
shard_paths = indexed if indexed else discover_shards(root)
|
|
missing = [path.name for path in shard_paths if not path.is_file()]
|
|
if missing:
|
|
raise ValueError(f"Index entries are missing from {root}: {', '.join(missing)}")
|
|
else:
|
|
shard_paths = selected
|
|
|
|
if not config and not shard_paths:
|
|
raise SystemExit(f"No config.json or safetensors files in {root}")
|
|
|
|
tensors, shards = load_shards(shard_paths) if shard_paths else ([], [])
|
|
opened = {shard.path.name for shard in shards}
|
|
index_names = {path.name for path in indexed}
|
|
partial = bool(index_names) and opened != index_names and opened.issubset(index_names)
|
|
if index_names and opened == index_names:
|
|
missing = [name for name in weight_map if isinstance(name, str) and name not in {tensor.name for tensor in tensors}]
|
|
if missing:
|
|
raise ValueError(f"Index lists {missing[0]!r}, which is not in the safetensors headers")
|
|
|
|
comfy = any(tensor.name.endswith(".comfy_quant") for tensor in tensors) or any(
|
|
"_quantization_metadata" in shard.metadata for shard in shards
|
|
)
|
|
packed = any(tensor.name.endswith("weight_packed") for tensor in tensors)
|
|
recipe_nvfp4 = bool(recipe) and "NVFP4" in recipe.upper()
|
|
nvfp4 = quant_is_nvfp4(quant) or comfy or (recipe_nvfp4 and packed)
|
|
template, template_where, template_count = load_chat_template(root, tokenizer)
|
|
index_meta = index_payload.get("metadata") if isinstance(index_payload.get("metadata"), dict) else {}
|
|
|
|
return Checkpoint(
|
|
path=path,
|
|
root=root,
|
|
config=config,
|
|
config_name=config_name,
|
|
generation=generation,
|
|
tokenizer=tokenizer,
|
|
quant=quant,
|
|
hf_quant=hf_quant,
|
|
recipe=recipe,
|
|
base_model=base_model,
|
|
tensors=tensors,
|
|
shards=shards,
|
|
shard_total=len(indexed) if indexed else len(shards),
|
|
partial=partial,
|
|
index_meta=index_meta,
|
|
template=template,
|
|
template_where=template_where,
|
|
template_count=template_count,
|
|
nvfp4=nvfp4,
|
|
comfy=comfy,
|
|
)
|
|
|
|
|
|
def load_checkpoint(path: Path) -> Checkpoint:
|
|
try:
|
|
return _load(path)
|
|
except (OSError, ValueError, struct.error, UnicodeError) as exc:
|
|
raise SystemExit(str(exc)) from exc
|
|
|
|
|
|
def _type_of(checkpoint: Checkpoint):
|
|
method = str(checkpoint.quant.get("quant_method") or "")
|
|
|
|
def kind(tensor: TensorInfo) -> str:
|
|
return tensor_type(tensor, method, checkpoint.nvfp4, checkpoint.comfy)
|
|
|
|
return kind
|
|
|
|
|
|
def _quant_title(checkpoint: Checkpoint) -> str:
|
|
quant = checkpoint.quant
|
|
algo = quant.get("quant_algo")
|
|
if isinstance(algo, str) and algo:
|
|
return algo
|
|
if checkpoint.nvfp4:
|
|
return "NVFP4"
|
|
fmt = quant.get("format")
|
|
if isinstance(fmt, str) and fmt:
|
|
return fmt
|
|
method = quant.get("quant_method")
|
|
if isinstance(method, str) and method:
|
|
return method
|
|
if checkpoint.recipe:
|
|
return checkpoint.recipe
|
|
return ""
|
|
|
|
|
|
def print_identity(checkpoint: Checkpoint) -> None:
|
|
_line("file", checkpoint.path)
|
|
if checkpoint.shards:
|
|
size = sum(shard.path.stat().st_size for shard in checkpoint.shards)
|
|
_line("size", format_size(size))
|
|
model_type = _text(_present(checkpoint.config, ("model_type",)))
|
|
architectures = checkpoint.config.get("architectures")
|
|
if isinstance(architectures, list):
|
|
arch_text = ", ".join(str(item) for item in architectures if item not in (None, ""))
|
|
else:
|
|
arch_text = _text(architectures)
|
|
if model_type:
|
|
_line("architecture", model_type)
|
|
if arch_text and arch_text != model_type:
|
|
_line("class", arch_text)
|
|
elif arch_text:
|
|
_line("architecture", arch_text)
|
|
if checkpoint.path.is_dir() or checkpoint.config:
|
|
_line("name", checkpoint.root.name)
|
|
else:
|
|
_line("name", checkpoint.path.stem)
|
|
if checkpoint.base_model and checkpoint.base_model != checkpoint.root.name:
|
|
_line("base model", checkpoint.base_model)
|
|
dtype = _text(_present(checkpoint.config, ("dtype", "torch_dtype")))
|
|
if dtype:
|
|
_line("dtype", dtype)
|
|
|
|
|
|
def print_quantization(checkpoint: Checkpoint) -> None:
|
|
title = _quant_title(checkpoint)
|
|
if checkpoint.nvfp4 or checkpoint.quant:
|
|
if title:
|
|
_line("quantization", title)
|
|
elif checkpoint.recipe:
|
|
_line("recipe", checkpoint.recipe)
|
|
return
|
|
else:
|
|
return
|
|
method = _text(checkpoint.quant.get("quant_method"))
|
|
fmt = _text(checkpoint.quant.get("format"))
|
|
if method and method != title:
|
|
_line("quant method", method)
|
|
if fmt and fmt != title:
|
|
_line("quant format", fmt)
|
|
elif checkpoint.comfy and not fmt:
|
|
_line("quant format", "comfy")
|
|
schemes = _weight_schemes(checkpoint.quant)
|
|
bits = []
|
|
groups = []
|
|
act_bits = []
|
|
for weight_bits, group_size, activation_bits in schemes:
|
|
if isinstance(weight_bits, int) and not isinstance(weight_bits, bool) and weight_bits not in bits:
|
|
bits.append(weight_bits)
|
|
if isinstance(group_size, int) and not isinstance(group_size, bool) and group_size > 0 and group_size not in groups:
|
|
groups.append(group_size)
|
|
if isinstance(activation_bits, int) and not isinstance(activation_bits, bool) and activation_bits not in act_bits:
|
|
act_bits.append(activation_bits)
|
|
if bits:
|
|
_line("bits", ", ".join(str(item) for item in bits))
|
|
if groups:
|
|
_line("group size", ", ".join(str(item) for item in groups))
|
|
if act_bits:
|
|
_line("act bits", ", ".join(str(item) for item in act_bits))
|
|
status = _text(checkpoint.quant.get("quantization_status"))
|
|
if status:
|
|
_line("quant status", status)
|
|
version = checkpoint.quant.get("version")
|
|
if isinstance(version, (str, int, float)) and not isinstance(version, bool) and version != "":
|
|
_line("quant version", format_number(version) if isinstance(version, (int, float)) else version)
|
|
ignored = checkpoint.quant.get("ignore")
|
|
if isinstance(ignored, list) and ignored:
|
|
if all(isinstance(item, str) for item in ignored):
|
|
text = ", ".join(ignored)
|
|
if len(text) > 160:
|
|
text = text[:160] + "..."
|
|
_line("quant ignore", text)
|
|
else:
|
|
_line("quant ignore", _preview(ignored))
|
|
elif isinstance(ignored, str) and ignored:
|
|
_line("quant ignore", ignored)
|
|
layers = _comfy_layers(checkpoint)
|
|
if layers:
|
|
_line("quant layers", layers)
|
|
|
|
|
|
def _comfy_layers(checkpoint: Checkpoint) -> int:
|
|
total = 0
|
|
found = False
|
|
for shard in checkpoint.shards:
|
|
raw = shard.metadata.get("_quantization_metadata")
|
|
parsed = None
|
|
if isinstance(raw, str):
|
|
try:
|
|
parsed = json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
parsed = None
|
|
elif isinstance(raw, dict):
|
|
parsed = raw
|
|
if not isinstance(parsed, dict):
|
|
continue
|
|
layers = parsed.get("layers")
|
|
if isinstance(layers, dict):
|
|
total += len(layers)
|
|
found = True
|
|
return total if found else 0
|
|
|
|
|
|
def _submodule_text(block: dict) -> str:
|
|
parts = []
|
|
for label, keys in (
|
|
("layers", ("num_hidden_layers", "num_layers", "depth")),
|
|
("width", ("hidden_size", "d_model", "embed_dim")),
|
|
("heads", ("num_attention_heads", "num_heads")),
|
|
):
|
|
for key in keys:
|
|
value = block.get(key)
|
|
if isinstance(value, int) and not isinstance(value, bool) and value > 0:
|
|
parts.append(f"{value} {label}")
|
|
break
|
|
return ", ".join(parts)
|
|
|
|
|
|
def print_architecture(checkpoint: Checkpoint) -> None:
|
|
config = checkpoint.config
|
|
context = _positive_int(config, (
|
|
"max_position_embeddings",
|
|
"max_seq_len",
|
|
"n_ctx",
|
|
"n_positions",
|
|
"max_sequence_length",
|
|
"seq_length",
|
|
"model_max_length",
|
|
))
|
|
if context:
|
|
_line("context", context)
|
|
tok_length = _positive_int(checkpoint.tokenizer, ("model_max_length",))
|
|
if tok_length and tok_length != context:
|
|
_line("tok length", tok_length)
|
|
fields = (
|
|
("blocks", ("num_hidden_layers", "n_layer", "num_layers", "n_layers")),
|
|
("embedding", ("hidden_size", "n_embd", "d_model", "hidden_dim")),
|
|
("feed forward", ("intermediate_size", "n_inner", "ffn_dim", "intermediate_dim")),
|
|
("heads", ("num_attention_heads", "n_head", "num_heads")),
|
|
("kv heads", ("num_key_value_heads", "num_kv_heads", "n_head_kv")),
|
|
("experts", ("num_experts", "num_local_experts", "n_routed_experts", "moe_num_experts")),
|
|
("experts used", ("num_experts_per_tok", "num_experts_per_token")),
|
|
("expert ffn", ("moe_intermediate_size", "expert_intermediate_size")),
|
|
("vocab", ("vocab_size", "padded_vocab_size")),
|
|
)
|
|
values = {label: _positive_int(config, keys) for label, keys in fields}
|
|
head_dim = _positive_int(config, ("head_dim", "attention_head_dim", "head_size"))
|
|
if head_dim is None and values["embedding"] and values["heads"] and values["embedding"] % values["heads"] == 0:
|
|
if _present(config, ("qk_nope_head_dim", "qk_rope_head_dim")) is None:
|
|
head_dim = values["embedding"] // values["heads"]
|
|
for label, _keys in fields:
|
|
if values[label]:
|
|
_line(label, values[label])
|
|
if label == "kv heads" and head_dim:
|
|
_line("head dim", head_dim)
|
|
activation = _text(_present(config, ("hidden_act", "activation_function", "hidden_activation")))
|
|
if activation:
|
|
_line("activation", activation)
|
|
rope = _present(config, ("rope_theta",))
|
|
if isinstance(rope, (int, float)) and not isinstance(rope, bool) and rope != 0:
|
|
_line("rope theta", format_number(rope))
|
|
window = _positive_int(config, ("sliding_window", "sliding_window_size"))
|
|
if window and _present(config, ("use_sliding_window",)) is not False:
|
|
_line("sliding window", window)
|
|
layer_types = _present(config, ("layer_types", "sliding_window_pattern"))
|
|
if isinstance(layer_types, list) and layer_types:
|
|
if all(isinstance(item, bool) for item in layer_types):
|
|
sliding = sum(bool(item) for item in layer_types)
|
|
_line("swa layers", f"{sliding} sliding, {len(layer_types) - sliding} full")
|
|
elif all(isinstance(item, str) for item in layer_types):
|
|
counts = Counter(layer_types)
|
|
if not (len(counts) == 1 and next(iter(counts)) in {"full_attention", "full"}):
|
|
text = ", ".join(f"{count} {name}" for name, count in counts.most_common())
|
|
_line("layers", text)
|
|
for key, label in (("vision_config", "vision"), ("audio_config", "audio"), ("vision_tower", "vision")):
|
|
block = config.get(key)
|
|
if isinstance(block, dict):
|
|
text = _submodule_text(block)
|
|
if text:
|
|
_line(label, text)
|
|
elif isinstance(block, str) and block:
|
|
_line(label, block)
|
|
if _present(config, ("tie_word_embeddings", "tie_embeddings")) is True:
|
|
_line("tied embed", "yes")
|
|
|
|
|
|
def print_tokenizer(checkpoint: Checkpoint) -> None:
|
|
tokenizer = checkpoint.tokenizer
|
|
kind = _text(tokenizer.get("tokenizer_class"))
|
|
if not kind and (checkpoint.root / "tokenizer.model").is_file():
|
|
kind = "sentencepiece"
|
|
elif not kind and (checkpoint.root / "tokenizer.json").is_file():
|
|
kind = "tokenizer.json"
|
|
if kind:
|
|
_line("tokenizer", kind)
|
|
bos = _token_id(checkpoint.config, "bos_token_id")
|
|
if bos is None:
|
|
bos = _token_id(checkpoint.generation, "bos_token_id")
|
|
eos = _token_id(checkpoint.generation, "eos_token_id")
|
|
if eos is None:
|
|
eos = _token_id(checkpoint.config, "eos_token_id")
|
|
pad = _token_id(checkpoint.generation, "pad_token_id")
|
|
if pad is None:
|
|
pad = _token_id(checkpoint.config, "pad_token_id")
|
|
if bos is not None:
|
|
_line("bos", token_label(bos, tokenizer))
|
|
elif _token_text(tokenizer.get("bos_token")):
|
|
_line("bos", _token_text(tokenizer.get("bos_token")))
|
|
if eos is not None:
|
|
_line("eos", token_label(eos, tokenizer))
|
|
elif _token_text(tokenizer.get("eos_token")):
|
|
_line("eos", _token_text(tokenizer.get("eos_token")))
|
|
if pad is not None:
|
|
_line("pad", token_label(pad, tokenizer))
|
|
elif _token_text(tokenizer.get("pad_token")):
|
|
_line("pad", _token_text(tokenizer.get("pad_token")))
|
|
if checkpoint.template:
|
|
detail = f"{len(checkpoint.template)} chars"
|
|
if checkpoint.template_count > 1:
|
|
detail = f"{checkpoint.template_count} templates, {detail}"
|
|
if checkpoint.template_where:
|
|
detail = f"{detail}, {checkpoint.template_where}"
|
|
if chat_template_has_tools(checkpoint.template):
|
|
detail = f"{detail}, tools"
|
|
_line("chat template", f"yes ({detail})")
|
|
else:
|
|
_line("chat template", "no")
|
|
|
|
|
|
def print_tensor_summary(checkpoint: Checkpoint) -> None:
|
|
if checkpoint.partial:
|
|
_line("shards", f"{len(checkpoint.shards)} of {checkpoint.shard_total}")
|
|
elif len(checkpoint.shards) > 1:
|
|
_line("shards", str(len(checkpoint.shards)))
|
|
truncated = [
|
|
shard.path.name
|
|
for shard in checkpoint.shards
|
|
if shard.path.stat().st_size < 8 + shard.header_size + shard.data_end
|
|
]
|
|
if truncated:
|
|
shown = ", ".join(truncated[:4])
|
|
if len(truncated) > 4:
|
|
shown += f", ... ({len(truncated)})"
|
|
_line("truncated", shown)
|
|
kind = _type_of(checkpoint)
|
|
params = None if checkpoint.partial else parameter_count(checkpoint.tensors, kind)
|
|
if params:
|
|
_line("parameters", format_params(params))
|
|
_line("tensors", str(len(checkpoint.tensors)))
|
|
if not checkpoint.tensors:
|
|
return
|
|
counts: Counter[str] = Counter()
|
|
stored: dict[str, Counter[str]] = {}
|
|
known = 0
|
|
unknown = 0
|
|
for tensor in checkpoint.tensors:
|
|
name = kind(tensor)
|
|
counts[name] += 1
|
|
stored.setdefault(name, Counter())[tensor.dtype] += 1
|
|
if tensor.nbytes is None:
|
|
unknown += 1
|
|
else:
|
|
known += tensor.nbytes
|
|
width = max([10, *(len(name) for name in counts)])
|
|
for name, count in counts.most_common():
|
|
extra = ""
|
|
dtypes = stored[name]
|
|
if list(dtypes) != [name]:
|
|
if len(dtypes) == 1:
|
|
extra = f" (stored {next(iter(dtypes))})"
|
|
else:
|
|
detail = ", ".join(f"{dtype} x {number}" for dtype, number in dtypes.most_common())
|
|
extra = f" (stored {detail})"
|
|
print(f" {name:<{width}} {count}{extra}")
|
|
_line("tensor bytes", format_size(known))
|
|
if unknown:
|
|
_line("unmeasured", f"{unknown} tensors have no data length")
|
|
|
|
|
|
def print_summary(checkpoint: Checkpoint) -> None:
|
|
print_identity(checkpoint)
|
|
print_quantization(checkpoint)
|
|
print_architecture(checkpoint)
|
|
print_tokenizer(checkpoint)
|
|
print_tensor_summary(checkpoint)
|
|
|
|
|
|
def _metadata_tokenizer(tokenizer: dict) -> dict:
|
|
view = {}
|
|
for key, value in tokenizer.items():
|
|
if key == "added_tokens_decoder" and isinstance(value, dict):
|
|
view[key] = f"{len(value)} added tokens"
|
|
else:
|
|
view[key] = value
|
|
return view
|
|
|
|
|
|
def _metadata_safetensors(metadata: dict) -> dict:
|
|
view = {}
|
|
for key, value in metadata.items():
|
|
if key == "format" and value == "pt":
|
|
continue
|
|
if key == "_quantization_metadata":
|
|
parsed = None
|
|
if isinstance(value, str):
|
|
try:
|
|
parsed = json.loads(value)
|
|
except json.JSONDecodeError:
|
|
parsed = None
|
|
elif isinstance(value, dict):
|
|
parsed = value
|
|
layers = parsed.get("layers") if isinstance(parsed, dict) else None
|
|
if isinstance(parsed, dict) and isinstance(layers, dict):
|
|
version = parsed.get("format_version", "?")
|
|
view[key] = f"format {version}, {len(layers)} layers"
|
|
continue
|
|
view[key] = value
|
|
return view
|
|
|
|
|
|
def _print_mapping(title: str, mapping: dict) -> None:
|
|
if not mapping:
|
|
return
|
|
print()
|
|
print(title)
|
|
width = max(len(str(key)) for key in mapping)
|
|
for key, value in mapping.items():
|
|
print(f" {key:<{width}} {_preview(value)}")
|
|
|
|
|
|
def print_metadata(checkpoint: Checkpoint) -> None:
|
|
if checkpoint.config:
|
|
_print_mapping(checkpoint.config_name, checkpoint.config)
|
|
if checkpoint.hf_quant:
|
|
_print_mapping("hf_quant_config.json", checkpoint.hf_quant)
|
|
if checkpoint.generation:
|
|
_print_mapping("generation_config.json", checkpoint.generation)
|
|
if checkpoint.tokenizer:
|
|
_print_mapping("tokenizer_config.json", _metadata_tokenizer(checkpoint.tokenizer))
|
|
if checkpoint.template and not _template_parts(checkpoint.tokenizer.get("chat_template") if checkpoint.tokenizer else None):
|
|
_print_mapping(checkpoint.template_where or "chat_template", {"chat_template": checkpoint.template})
|
|
if checkpoint.index_meta:
|
|
_print_mapping("safetensors index", checkpoint.index_meta)
|
|
metas = []
|
|
for shard in checkpoint.shards:
|
|
cleaned = _metadata_safetensors(shard.metadata)
|
|
if cleaned and cleaned not in metas:
|
|
metas.append(cleaned)
|
|
title = "safetensors metadata" if len(checkpoint.shards) == 1 else f"safetensors metadata ({shard.path.name})"
|
|
if len(checkpoint.shards) > 1 and all(_metadata_safetensors(item.metadata) == cleaned for item in checkpoint.shards):
|
|
title = "safetensors metadata"
|
|
_print_mapping(title, cleaned)
|
|
if len(checkpoint.shards) > 1 and all(_metadata_safetensors(item.metadata) == cleaned for item in checkpoint.shards):
|
|
break
|
|
|
|
|
|
def _shape_text(tensor: TensorInfo) -> str:
|
|
return " x ".join(str(dim) for dim in tensor.shape) or "-"
|
|
|
|
|
|
def print_tensors(checkpoint: Checkpoint) -> None:
|
|
print()
|
|
if not checkpoint.tensors:
|
|
print("tensors")
|
|
print(" (none)")
|
|
return
|
|
kind = _type_of(checkpoint)
|
|
tensors = sorted(checkpoint.tensors, key=lambda tensor: _natural_key(tensor.name))
|
|
if any(kind(tensor) == "NVFP4" for tensor in tensors):
|
|
print("NVFP4 shapes are packed storage, 2 FP4 values per byte.")
|
|
name_width = max(len(tensor.name) for tensor in tensors)
|
|
type_width = max(8, max(len(kind(tensor)) for tensor in tensors))
|
|
shapes = [_shape_text(tensor) for tensor in tensors]
|
|
show_shard = len({tensor.file for tensor in tensors}) > 1
|
|
shape_width = max(len("SHAPE"), *(len(shape) for shape in shapes)) if show_shard else 0
|
|
shape_header = f"{'SHAPE':<{shape_width}}" if show_shard else "SHAPE"
|
|
header = f"{'NAME':<{name_width + 2}} {'TYPE':<{type_width}} {shape_header}"
|
|
if show_shard:
|
|
header += " SHARD"
|
|
print(header)
|
|
for tensor, shape in zip(tensors, shapes):
|
|
shown = f"{shape:<{shape_width}}" if show_shard else shape
|
|
row = f" {tensor.name:<{name_width}} {kind(tensor):<{type_width}} {shown}"
|
|
if show_shard:
|
|
row += f" {tensor.file}"
|
|
print(row)
|
|
|
|
|
|
def inspect(path: Path, metadata_dump: bool, tensor_dump: bool) -> None:
|
|
checkpoint = load_checkpoint(path)
|
|
print_summary(checkpoint)
|
|
if metadata_dump:
|
|
print_metadata(checkpoint)
|
|
if tensor_dump:
|
|
print_tensors(checkpoint)
|
|
|
|
|
|
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Print architecture, quantization, and tensor info from a Hugging Face checkpoint."
|
|
)
|
|
parser.add_argument("model", nargs="+", type=Path, help="A model directory or a .safetensors file")
|
|
parser.add_argument(
|
|
"--metadata",
|
|
action="store_true",
|
|
help="Print config, tokenizer, and safetensors header metadata. Long values are shortened.",
|
|
)
|
|
parser.add_argument("--tensors", action="store_true", help="Print every tensor name, type, and shape.")
|
|
return parser.parse_args(argv)
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> None:
|
|
args = parse_args(argv)
|
|
for index, path in enumerate(args.model):
|
|
if index:
|
|
print()
|
|
inspect(path, args.metadata, args.tensors)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
main()
|
|
except BrokenPipeError:
|
|
# head and less close the pipe early. Don't turn that into a traceback.
|
|
try:
|
|
sys.stdout.close()
|
|
except BrokenPipeError:
|
|
pass
|
|
try:
|
|
sys.stderr.close()
|
|
except BrokenPipeError:
|
|
pass
|
|
os._exit(0)
|