"""Read-only structural probe of an R49 candidate carrier.

Answers, by measurement rather than by reading the config:
  * does transformers on this box load the checkpoint at all,
  * which module paths are nn.Linear (the only LoRA-attachable leaves),
  * how the parameter budget splits across text body / vision tower / MTP /
    embeddings, so "0.8B carrier" can be reported honestly,
  * which attention implementations the class accepts.

Loads on CPU in bf16. No training, no GPU, nothing written but stdout.
"""
import json, sys, collections, re
import torch
from transformers import AutoConfig, AutoModelForCausalLM

path = sys.argv[1]
print(f"== {path}")
cfg = AutoConfig.from_pretrained(path, trust_remote_code=False)
print("   config class     :", type(cfg).__name__)
print("   architectures    :", getattr(cfg, "architectures", None))
tc = getattr(cfg, "text_config", None)
if tc is not None:
    lt = getattr(tc, "layer_types", None) or []
    print("   text layers      :", getattr(tc, "num_hidden_layers", "?"),
          "| full_attention:", lt.count("full_attention"),
          "| linear_attention:", lt.count("linear_attention"))
    print("   hidden/inter     :", getattr(tc, "hidden_size", "?"), "/", getattr(tc, "intermediate_size", "?"))
    print("   vocab            :", getattr(tc, "vocab_size", "?"), "| tied:", getattr(tc, "tie_word_embeddings", "?"))

try:
    model = AutoModelForCausalLM.from_pretrained(
        path, dtype=torch.bfloat16, device_map="cpu",
        attn_implementation="sdpa",
    )
except Exception as e:
    print("   LOAD FAILED:", type(e).__name__, str(e)[:400])
    raise SystemExit(1)
print("   model class      :", type(model).__name__)
print("   attn impl        :", getattr(model.config, "_attn_implementation", "?"))

# --- parameter budget ------------------------------------------------------
buckets = collections.Counter()
def bucket(name):
    if ".visual." in name or name.startswith("visual."):        return "vision_tower"
    if name.startswith("mtp.") or ".mtp." in name:              return "mtp_head"
    if "embed_tokens" in name or name.endswith("lm_head.weight"): return "embeddings"
    if "linear_attn" in name:                                   return "text_linear_attn"
    if "self_attn" in name:                                     return "text_full_attn"
    if ".mlp." in name:                                         return "text_mlp"
    return "text_other"
for n, p in model.named_parameters():
    buckets[bucket(n)] += p.numel()
total = sum(buckets.values())
print(f"   TOTAL params     : {total/1e9:.3f} B")
for k, v in sorted(buckets.items(), key=lambda kv: -kv[1]):
    print(f"     {k:<18} {v/1e6:9.1f} M  ({100*v/total:5.1f}%)")

# --- LoRA-attachable leaves ------------------------------------------------
lin = collections.defaultdict(list)
for name, mod in model.named_modules():
    if isinstance(mod, torch.nn.Linear):
        lin[bucket(name + ".weight")].append(name)
print("   nn.Linear leaves by region:")
for region in sorted(lin):
    names = lin[region]
    tmpl = sorted({re.sub(r"\.\d+\.", ".N.", n) for n in names})
    print(f"     {region:<18} {len(names):4d} modules, {len(tmpl)} distinct shapes")
    for t in tmpl:
        print(f"        {t}")