#!/usr/bin/env python3
"""Find hc-related and global keys in the checkpoint."""
import json
from pathlib import Path

CHECKPOINT_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4"
cdir = Path(CHECKPOINT_DIR)
idx = cdir / "model.safetensors.index.json"
with open(idx) as f:
    wmap = json.load(f).get("weight_map", {})

# Find all keys containing 'hc' or global (non-layers.) keys
hc_keys = sorted(k for k in wmap if 'hc' in k.lower())
print(f"HC keys ({len(hc_keys)}):")
for k in hc_keys[:20]:
    print(f"  {k}")
if len(hc_keys) > 20:
    print(f"  ... and {len(hc_keys)-20} more")

# Global keys (not starting with 'layers.')
global_keys = sorted(k for k in wmap if not k.startswith('layers.'))
print(f"\nGlobal keys ({len(global_keys)}):")
for k in global_keys:
    print(f"  {k}")

# Check: are there keys with 'hyper' or 'connection' in them?
hyper_keys = sorted(k for k in wmap if 'hyper' in k.lower() or 'connection' in k.lower())
print(f"\nHyper/connection keys ({len(hyper_keys)}):")
for k in hyper_keys[:10]:
    print(f"  {k}")

# Check layer 0 non-expert keys that contain 'hc' or 'hyper'
l0_hc = sorted(k for k in wmap if k.startswith('layers.0.') and ('hc' in k.lower() or 'hyper' in k.lower()))
print(f"\nLayer 0 hc/hyper keys: {l0_hc}")