"""Do the CHECKPOINT's tensor shapes agree with num_global_key_value_heads=4?

Patching a config to satisfy a constructor is only safe if the weights already have
the shape the patched value implies. If the abliteration reshaped attention, the
patch would silence the error and produce a quietly wrong quant -- worse than the
crash, because it ships.

For a full-attention layer, k_proj/v_proj out-features == num_kv_heads * head_dim.
zerofata's canonical v2: num_global_key_value_heads=4, global_head_dim=512
  => expected out-features 4 * 512 = 2048 on the GLOBAL (full-attention) layers.
"""
import json, sys
from safetensors import safe_open
from pathlib import Path

def probe(label, root, n_global_kv, global_head_dim, n_kv, head_dim):
    root = Path(root)
    idx = json.load(open(root / "model.safetensors.index.json"))["weight_map"]
    types = json.load(open(root / "config.json"))["text_config"]["layer_types"]
    full = [i for i, t in enumerate(types) if t == "full_attention"][:2]
    slide = [i for i, t in enumerate(types) if t == "sliding_attention"][:2]
    print(f"  -- {label}")
    print(f"     expected FULL   k/v out-features = {n_global_kv} x {global_head_dim} = {n_global_kv*global_head_dim}"
          if n_global_kv and global_head_dim else "     expected FULL   = (config lacks the fields)")
    print(f"     expected SLIDE  k/v out-features = {n_kv} x {head_dim} = {n_kv*head_dim}")
    for tag, idxs in (("full ", full), ("slide", slide)):
        for li in idxs:
            for proj in ("k_proj", "v_proj"):
                key = f"model.language_model.layers.{li}.self_attn.{proj}.weight"
                if key not in idx:
                    key = f"language_model.model.layers.{li}.self_attn.{proj}.weight"
                if key not in idx:
                    cand = [k for k in idx if f"layers.{li}.self_attn.{proj}" in k]
                    key = cand[0] if cand else None
                if not key:
                    print(f"     {tag} L{li} {proj}: KEY NOT FOUND"); continue
                with safe_open(root / idx[key], framework="pt") as f:
                    shape = f.get_slice(key).get_shape()
                print(f"     {tag} L{li} {proj}: shape {shape}  out-features={shape[0]}")

cfg = json.load(open("/tank/aimodels/G4-MeroMero-v2-31B-heretic-bf16/config.json"))["text_config"]
good = json.load(open("/tank/aimodels/meromero-v2-nvfp4-work/src/config.json"))["text_config"]
print(f"  canonical (zerofata): num_global_key_value_heads={good.get('num_global_key_value_heads')} "
      f"global_head_dim={good.get('global_head_dim')} num_key_value_heads={good.get('num_key_value_heads')} head_dim={good.get('head_dim')}")
probe("zerofata v2 (canonical)", "/tank/aimodels/meromero-v2-nvfp4-work/src",
      good.get("num_global_key_value_heads"), good.get("global_head_dim"),
      good.get("num_key_value_heads"), good.get("head_dim"))
probe("DogOnKeyboard v2 (to patch)", "/tank/aimodels/G4-MeroMero-v2-31B-heretic-bf16",
      good.get("num_global_key_value_heads"), good.get("global_head_dim"),
      cfg.get("num_key_value_heads"), cfg.get("head_dim"))