import json
from safetensors import safe_open

idx = json.load(open("model.safetensors.index.json"))["weight_map"]
have = set()
for f in sorted(set(idx.values())):
    with safe_open(f, framework="pt") as st:
        have.update(st.keys())

print("indexed:", len(idx), "physical:", len(have))
print("missing from shards:", [k for k in idx if k not in have])
print("unindexed extras:", [k for k in have if k not in idx])

fc = idx.get("mtp.fc.weight")
if fc:
    with safe_open(fc, framework="pt") as st:
        t = st.get_slice("mtp.fc.weight")
        print("mtp.fc.weight:", t.get_shape(), t.get_dtype())  # expect [10240, 5120] BF16