| 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) |
|
|
