import torch, json, time, re
import numpy as np
from transformers import AutoModelForImageTextToText, AutoTokenizer, AutoConfig
MODEL="/model"
cfg=AutoConfig.from_pretrained(MODEL); tc=cfg.text_config
L,E,K=tc.num_hidden_layers,tc.num_experts,tc.num_experts_per_tok
tok=AutoTokenizer.from_pretrained(MODEL,trust_remote_code=True)
print("cargando modelo...",flush=True); t0=time.time()
model=AutoModelForImageTextToText.from_pretrained(MODEL,device_map="auto",dtype=torch.bfloat16); model.eval()
print(f"cargado {time.time()-t0:.0f}s",flush=True)
counts=np.zeros((L,E),dtype=np.int64)
def mk(li):
    def h(m,i,o):
        lg=o[0] if isinstance(o,(tuple,list)) else o
        if not torch.is_tensor(lg): return
        lg=lg.reshape(-1,lg.shape[-1])
        if lg.shape[-1]!=E: return
        counts[li]+=torch.bincount(lg.float().topk(K,-1).indices.flatten(),minlength=E).cpu().numpy()
    return h
nh=0
for n,mod in model.named_modules():
    if re.search(r"layers\.(\d+)\.mlp\.gate$",n):
        mod.register_forward_hook(mk(int(re.search(r"layers\.(\d+)\.",n).group(1)))); nh+=1
print(f"hooks {nh}",flush=True)
tools=json.load(open("/tools_native.json"))
corpus=[json.loads(l) for l in open("/corpus_prune.jsonl")]
SYS7K=corpus[0]["messages"][0]["content"]
SHORT="Eres Valentina, asesora de ventas de Malva Label. Usa las herramientas para guardar datos y gestionar el carrito."
text=tok.apply_chat_template([{"role":"system","content":SYS7K},{"role":"user","content":"hola"}],tools=tools,add_generation_prompt=False,tokenize=False)
ids=tok(text,return_tensors="pt").input_ids.to(model.device)
with torch.no_grad(): model(input_ids=ids,use_cache=False)
np.save("/counts_sys7k.npy",counts); print(f"[1] sys7k: {int((counts>0).sum(1).mean())} exp/capa ({ids.shape[1]}tok)",flush=True)
counts[:]=0
t0=time.time();done=0;errs=0
for ci,conv in enumerate(corpus):
    try:
        msgs=[{"role":"system","content":SHORT}]+conv["messages"][1:]
        text=tok.apply_chat_template(msgs,tools=tools,add_generation_prompt=False,tokenize=False)
        ids=tok(text,return_tensors="pt",truncation=True,max_length=4096).input_ids.to(model.device)
        with torch.no_grad(): model(input_ids=ids,use_cache=False)
        done+=1
    except Exception: errs+=1
    if (ci+1)%50==0:
        np.save("/counts_conv.npy",counts)
        el=time.time()-t0; print(f"[2] conv {ci+1}/{len(corpus)} {el:.0f}s ETA {el/(ci+1)*(len(corpus)-ci-1)/60:.0f}min",flush=True)
np.save("/counts_conv.npy",counts)
print(f"[2] CONV DONE {done} ok {errs} err {time.time()-t0:.0f}s",flush=True)