| import torch, json, time, re, glob, traceback |
| import numpy as np |
| from transformers import AutoModelForImageTextToText, AutoProcessor, AutoConfig |
| from PIL import Image |
| 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 |
| IMG_TOK=getattr(cfg,"image_token_id",None) |
| print(f"L={L} E={E} K={K} img_token={IMG_TOK}",flush=True) |
| proc=AutoProcessor.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_vis=np.zeros((L,E),dtype=np.int64) |
| counts_all=np.zeros((L,E),dtype=np.int64) |
| state={"imgmask":None} |
| 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 |
| idx=lg.float().topk(K,-1).indices |
| counts_all[li]+=torch.bincount(idx.flatten(),minlength=E).cpu().numpy() |
| mm=state["imgmask"] |
| if mm is not None and mm.shape[0]==idx.shape[0]: |
| vis=idx[mm.to(idx.device)] |
| if vis.numel(): counts_vis[li]+=torch.bincount(vis.flatten(),minlength=E).cpu().numpy() |
| return h |
| 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)))) |
| imgs=sorted(glob.glob("/pimg/*.jpg"))[:55] |
| print(f"{len(imgs)} imagenes",flush=True) |
| done=0;errs=0;t0=time.time() |
| for ii,ip in enumerate(imgs): |
| try: |
| image=Image.open(ip).convert("RGB") |
| msgs=[{"role":"user","content":[{"type":"image"},{"type":"text","text":"Describe esta imagen para un negocio de lenceria."}]}] |
| text=proc.apply_chat_template(msgs,add_generation_prompt=True,tokenize=False) |
| inputs=proc(text=[text],images=[image],return_tensors="pt").to(model.device) |
| ids=inputs["input_ids"][0] |
| state["imgmask"]=(ids==IMG_TOK) if IMG_TOK is not None else None |
| with torch.no_grad(): model(**inputs,use_cache=False) |
| nvis=int(state["imgmask"].sum()) if state["imgmask"] is not None else 0 |
| state["imgmask"]=None |
| done+=1 |
| if ii<3: print(f" img{ii}: {ids.shape[0]} tok, {nvis} visuales",flush=True) |
| except Exception as e: |
| errs+=1 |
| if errs<=3: traceback.print_exc() |
| if (ii+1)%10==0: |
| np.save("/counts_vision.npy",counts_vis); np.save("/counts_visall.npy",counts_all) |
| print(f"img {ii+1}/{len(imgs)} {time.time()-t0:.0f}s",flush=True) |
| np.save("/counts_vision.npy",counts_vis); np.save("/counts_visall.npy",counts_all) |
| nv=(counts_vis>0).sum(1) |
| print(f"VISION DONE {done} ok {errs} err. expertos por tokens VISUALES: min={nv.min()} max={nv.max()} mean={nv.mean():.0f} de {E}",flush=True) |
|
|
