"""Structural validation for a modelopt NVFP4 checkpoint.

Checks that every tensor in model.safetensors.index.json exists in the
shards (and vice versa) and prints the dtype layout + quant config.

Usage:
  python validate_nvfp4.py [MODEL_DIR]   (default: current directory)
"""
import collections
import json
import os
import struct
import sys

DD = sys.argv[1] if len(sys.argv) > 1 else "."
idx = json.load(open(os.path.join(DD, "model.safetensors.index.json")))
wm = idx["weight_map"]

# headers of each shard (no tensor data is read)
have = {}
for fname in set(wm.values()):
    with open(os.path.join(DD, fname), "rb") as f:
        n = struct.unpack("<Q", f.read(8))[0]
        header = json.loads(f.read(n))
    for t in header:
        if t != "__metadata__":
            have[t] = (fname, header[t]["dtype"], header[t]["shape"])

missing = [t for t in wm if t not in have]
extra = [t for t in have if t not in wm]
print("tensors in index:", len(wm))
print("tensors in shards:", len(have))
print("missing:", len(missing), missing[:5])
print("extra:", len(extra), extra[:5])

dtypes = collections.Counter(v[1] for v in have.values())
print("dtypes:", dict(dtypes))
# packed NVFP4 weights show up as U8 with per-block scale tensors
suff = collections.Counter()
for t in have:
    s = t.split(".")[-1]
    suff[s] += 1
print("common suffixes:", suff.most_common(12))
q = json.load(open(os.path.join(DD, "hf_quant_config.json")))
print("quant_algo:", q["quantization"]["quant_algo"],
      "| group_size:", q["quantization"]["group_size"])

assert not missing and not extra, "shard/index mismatch!"
assert q["quantization"]["quant_algo"] == "NVFP4", "not an NVFP4 checkpoint!"
print("OK - valid NVFP4 checkpoint")