File size: 25,075 Bytes
5851d4e | 1 2 3 4 5 6 7 8 9 10 11 | import os,subprocess,sys,glob
subprocess.run([sys.executable,'-m','pip','install','-q','transformers==4.46.3'],check=True)
inp=os.path.dirname(glob.glob('/kaggle/input/**/k2_s17.pt',recursive=True)[0])
S='/kaggle/working/src'; os.makedirs(S+'/data',exist_ok=True); os.system(f'cp {inp}/expanded.json {S}/data/')
from transformers import AutoTokenizer,AutoModelForCausalLM
mp='/kaggle/working/m135'; R='12fd25f77366fa6b3b4b768ec3050bf629380bac'
AutoTokenizer.from_pretrained('HuggingFaceTB/SmolLM2-135M-Instruct',revision=R).save_pretrained(mp); AutoModelForCausalLM.from_pretrained('HuggingFaceTB/SmolLM2-135M-Instruct',revision=R).save_pretrained(mp)
open('/kaggle/working/run018.py','w').write('# CUERDAS-015: transferencia (valores/consultas nuevas), vocabulario aleatorizado, K=1 vs K=4. CPU.\nimport json,random,time,re,sys,statistics as st,itertools\nimport torch,torch.nn as nn,torch.nn.functional as F\nfrom transformers import AutoModelForCausalLM,AutoTokenizer,AutoConfig\ntorch.set_num_threads(8)\nimport os\nSRC=os.environ.get(\'SRC\',\'/workspace/cuerdas014/x/cuerdas014\'); MP=os.environ.get(\'MP\',\'/workspace/cuerdas014/x/models/smollm2-135m-instruct\'); OD=os.environ.get(\'OD\',\'/workspace/cuerdas015\')\nDEV=\'cuda\' if torch.cuda.is_available() else \'cpu\'\ntok=AutoTokenizer.from_pretrained(MP); lm=AutoModelForCausalLM.from_pretrained(MP,torch_dtype=torch.float32).eval().requires_grad_(False).to(DEV)\nif DEV==\'cuda\': torch.set_default_device(\'cuda\')\nEmb=lm.get_input_embeddings().weight; D=Emb.shape[1]\nmeta=json.load(open(f\'{SRC}/data/expanded.json\')); V=meta[\'vocabulary\']; FIELDS=[\'action\',\'target\',\'status\',\'window\']\nrows=meta[\'rows\']; SPL=meta[\'splits\']\nTPL="The action is {action}. The target is {target}. The status is {status}. The window is {window}."\nSUF_MAIN="\\nAccording to the record, the {field} is"\nSUF_TRAIN=[SUF_MAIN,"\\nThe {field} is","\\nIn this record, the {field} is","\\nQuestion: what is the {field}? Answer:"]\nSUF_HELD={"recorded":"\\nThe recorded {field} is","value":"\\nThe {field} value is","past":"\\nThe {field} was"}\ndef t1(w): return tok.encode(\' \'+w,add_special_tokens=False)\n# ---- vocabularios nuevos (nunca en entrenamiento) ----\nNEW={\'action\':[\'inspect\',\'repair\',\'cancel\',\'approve\',\'archive\',\'scan\',\'print\',\'send\'],\n \'target\':[\'sensor\',\'battery\',\'printer\',\'server\',\'truck\',\'door\',\'screen\',\'cable\'],\n \'status\':[\'broken\',\'stable\',\'late\',\'empty\',\'full\',\'warm\',\'missing\',\'urgent\'],\n \'window\':[\'night\',\'evening\',\'weekend\',\'noon\',\'spring\',\'summer\',\'winter\',\'midnight\']}\nPHRASE={\'action\':[\'shut down\',\'back up\',\'sign off\',\'log out\'],\'target\':[\'air filter\',\'fuel tank\',\'water heater\',\'hard drive\'],\n \'status\':[\'out of stock\',\'on hold\',\'under review\',\'half full\'],\'window\':[\'next week\',\'last night\',\'this month\',\'later today\']}\n# limpiar: primer token único dentro de (vocab original ∪ nuevos ∪ frases) por campo\nfor j,f in enumerate(FIELDS):\n used={t1(w)[0] for w in V[j]}; keep=[];keepp=[]\n for w in NEW[f]:\n if len(t1(w))==1 and t1(w)[0] not in used and w not in V[j]: keep.append(w); used.add(t1(w)[0])\n for w in PHRASE[f]:\n if len(t1(w))>1 and t1(w)[0] not in used: keepp.append(w); used.add(t1(w)[0])\n NEW[f]=keep; PHRASE[f]=keepp\nBANNED={t1(w)[0] for f in FIELDS for w in NEW[f]+PHRASE[f]}|{t1(w)[0] for vs in V for w in vs}\n# pool de palabras aleatorias para entrenamiento (meta-aprendizaje de binding)\nPOOL=[]\nfor t,i in tok.get_vocab().items():\n if t.startswith(\'Ġ\') and re.fullmatch(r\'[a-z]{3,9}\',t[1:]) and i not in BANNED: POOL.append(t[1:])\nrandom.Random(0).shuffle(POOL); POOL_VAL=POOL[3000:3200]; POOL=POOL[:3000]\n# ---- splits: los de 014 + test FRESCO nuevo del universo no usado ----\nused_vals={tuple(r[\'values\']) for r in rows}\nfit=[tuple(rows[i][\'values\']) for i in SPL[\'fit\']]; seen_at={v[:2] for v in fit}; seen_sw={v[2:] for v in fit}\nrng=random.Random(15015); uni=[v for v in itertools.product(*[range(len(x)) for x in V]) if v not in used_vals]; rng.shuffle(uni)\nfresh_id=[v for v in uni if v[:2] in seen_at and v[2:] in seen_sw][:200]; fresh_ood=[v for v in uni if not(v[:2] in seen_at and v[2:] in seen_sw)][:400]\ndef W(v): return [V[j][v[j]] for j in range(4)]\nFIT=[W(v) for v in fit]; VAL=[W(tuple(rows[i][\'values\'])) for i in SPL[\'validation\']]\nTEST014={\'id\':[W(tuple(rows[i][\'values\'])) for i in SPL[\'id_test\']],\'ood\':[W(tuple(rows[i][\'values\'])) for i in SPL[\'ood_test\']]}\nFRESH={\'id\':[W(v) for v in fresh_id],\'ood\':[W(v) for v in fresh_ood]}\nr2=random.Random(77)\ndef novel_set(nnew,n=96,phr=False): # nnew campos con valor nuevo (los demás del vocab), tomados de FRESH\n out=[]\n for k in range(n):\n base=list(FRESH[\'ood\'][k%len(FRESH[\'ood\'])]); js=r2.sample(range(4),nnew)\n for j in js: base[j]=r2.choice(PHRASE[FIELDS[j]] if phr else NEW[FIELDS[j]])\n out.append(base)\n return out\nTRANSFER={\'1_new\':novel_set(1),\'2_new\':novel_set(2),\'4_new\':novel_set(4),\'1_phrase\':novel_set(1,phr=True)}\nCAND=[sorted({t1(w)[0] for w in V[j]+NEW[f]+PHRASE[f]}) for j,f in enumerate(FIELDS)] # opciones válidas: abiertas (vocab+nuevos)\nCAND=[torch.tensor(c) for c in CAND]\ndef lex(w): return Emb[t1(w)].mean(0)\ndef gold_idx(j,w): return int((CAND[j]==t1(w)[0]).nonzero()[0])\n# ---- codecs ----\nclass Binder(nn.Module): # K=4: ranura por campo, binding de cualquier palabra a un rol\n def __init__(s,seed):\n super().__init__(); torch.manual_seed(seed); s.K=4\n anc=torch.stack([Emb[t1(f)[0]] for f in FIELDS]); s.roles=nn.Parameter(anc.clone())\n s.net=nn.Sequential(nn.LayerNorm(D),nn.Linear(D,128),nn.GELU(),nn.Linear(128,D)); nn.init.normal_(s.net[-1].weight,std=.002); nn.init.zeros_(s.net[-1].bias)\n s.register_buffer(\'sc\',Emb.std().clone())\n def forward(s,L): c=L+s.roles[None]; return c+s.sc*s.net(c) # L: [B,4,D]\nclass Pack1(nn.Module): # K=1: empaqueta los 4 campos en UNA posición\n def __init__(s,seed):\n super().__init__(); torch.manual_seed(seed); s.K=1\n s.roles=nn.Parameter(torch.randn(4,D)*.02)\n s.net=nn.Sequential(nn.LayerNorm(4*D),nn.Linear(4*D,512),nn.GELU(),nn.Linear(512,D)); s.register_buffer(\'sc\',Emb.std().clone())\n s.mix=nn.Parameter(torch.full((4,),.25))\n def forward(s,L):\n c=L+s.roles[None]; base=(s.mix[None,:,None]*c).sum(1)\n return (base+s.sc*s.net(c.flatten(1)))[:,None]\nclass Their014(nn.Module): # su RoleAdapter, con léxico extensible (mismo cálculo que su forward)\n def __init__(s,seed):\n super().__init__(); ck=torch.load(f\'{SRC}/checkpoints/expanded_role_ce_k4_seed{seed}/best.pt\',map_location=DEV,weights_only=False)[\'adapter_state\']\n s.K=4; s.roles=nn.Parameter(ck[\'roles\']); s.sc=ck[\'scale\']\n s.net=nn.Sequential(nn.LayerNorm(D),nn.Linear(D,128),nn.GELU(),nn.Linear(128,D)); s.net.load_state_dict({k[4:]:v for k,v in ck.items() if k.startswith(\'net.\')})\n def forward(s,L): c=L+s.roles[None]; return c+s.sc*s.net(c)\n# ---- lectura por el LLM congelado ----\ndef sufids(tpl,j): return tok.encode(tpl.format(field=FIELDS[j]),add_special_tokens=False)\ndef run_rows(embs,tails): # listas de tensores [L,D] -> logits última posición (padding a la derecha)\n L=max(e.shape[0] for e in embs); B=len(embs)\n x=torch.zeros(B,L,D); m=torch.zeros(B,L,dtype=torch.long)\n for i,e in enumerate(embs): x[i,:e.shape[0]]=e; m[i,:e.shape[0]]=1\n h=lm.model(inputs_embeds=x,attention_mask=m).last_hidden_state[torch.arange(B),m.sum(1)-1]\n return lm.lm_head(h)\ndef codec_logits(cod,msgs,tpl,js=range(4)):\n S=cod(torch.stack([torch.stack([lex(w) for w in m]) for m in msgs]))\n embs=[];meta_=[]\n for i,m in enumerate(msgs):\n for j in js: embs.append(torch.cat([S[i],Emb[sufids(tpl,j)]])); meta_.append((i,j))\n return run_rows(embs,None),meta_\ndef text_logits(msgs,tpl):\n embs=[];meta_=[]\n for i,m in enumerate(msgs):\n for j in range(4): embs.append(Emb[tok.encode(TPL.format(**dict(zip(FIELDS,m)))+tpl.format(field=FIELDS[j]),add_special_tokens=False)]); meta_.append((i,j))\n return run_rows(embs,None),meta_\n@torch.no_grad()\ndef acc(route,msgs,tpl=SUF_MAIN,cod=None,bs=24):\n ok=[]\n for b in range(0,len(msgs),bs):\n mb=msgs[b:b+bs]; lg,mt=(text_logits(mb,tpl) if route==\'text\' else codec_logits(cod,mb,tpl))\n good=torch.ones(len(mb),dtype=torch.bool)\n for r,(i,j) in enumerate(mt): good[i]&=int(lg[r][CAND[j]].argmax())==gold_idx(j,mb[i][j])\n ok+=good.tolist()\n return sum(ok)/len(ok)\n@torch.no_grad()\ndef nolm(cod,msgs): # anti-lookup: decodificar valor desde las ranuras por coseno, sin LLM\n S=cod(torch.stack([torch.stack([lex(w) for w in m]) for m in msgs])); ok=0\n for i,m in enumerate(msgs):\n g=True\n for j in range(4):\n cands=F.normalize(Emb[CAND[j]],dim=-1); s=S[i,j] if S.shape[1]==4 else S[i].mean(0)\n g&=int((F.normalize(s,dim=-1)@cands.T).argmax())==gold_idx(j,m[j])\n ok+=g\n return ok/len(msgs)\ndef train(cod,steps,seed,p_rand=.5,p_phr=.15,bs=12):\n g=random.Random(seed); opt=torch.optim.AdamW(cod.parameters(),lr=2e-3,weight_decay=1e-3)\n for s in range(steps):\n msgs=[]\n for _ in range(bs):\n m=list(g.choice(FIT))\n for j in range(4):\n u=g.random()\n if u<p_phr: m[j]=g.choice(POOL)+\' \'+g.choice(POOL)\n elif u<p_phr+p_rand: m[j]=g.choice(POOL)\n msgs.append(m)\n tpl=g.choice(SUF_TRAIN); lg,mt=codec_logits(cod,msgs,tpl)\n gold=torch.tensor([t1(msgs[i][j])[0] for i,j in mt]); loss=F.cross_entropy(lg,gold) # CE sobre TODO el vocabulario: copiar la palabra ligada al rol\n opt.zero_grad(); loss.backward(); opt.step()\n return cod.eval()\nOUT=OD+\'/results018.jsonl\'\ndef log(d): open(OUT,\'a\').write(json.dumps(d)+"\\n"); print(json.dumps(d),flush=True)\ndef evaluate_all(name,cod=None,quick=False):\n route=\'text\' if cod is None else \'codec\'; r={\'name\':name}\n r[\'val\']=acc(route,VAL[:96],cod=cod)\n if not quick:\n for k,v in FRESH.items(): r[\'fresh_\'+k]=acc(route,v,cod=cod)\n for k,v in TRANSFER.items(): r[\'transfer_\'+k]=acc(route,v,cod=cod)\n sub=FRESH[\'id\'][:48]+FRESH[\'ood\'][:48]\n for k,t in SUF_HELD.items(): r[\'query_\'+k]=acc(route,sub,t,cod=cod)\n if cod is not None: r[\'noLLM_fresh_ood\']=nolm(cod,FRESH[\'ood\'][:200]); r[\'noLLM_transfer_4new\']=nolm(cod,TRANSFER[\'4_new\'])\n log(r); return r\n\ndef _branch(cache):\n if hasattr(cache,\'to_legacy_cache\'): cache=cache.to_legacy_cache()\n return tuple((k.repeat_interleave(4,0),v.repeat_interleave(4,0)) for k,v in cache)\n@torch.no_grad()\ndef one_event(route,m,cod,dev):\n tails=[sufids(SUF_MAIN,j) for j in range(4)]; sn=0\n while all(t[sn]==tails[0][sn] for t in tails): sn+=1\n if route==\'text\':\n pre=Emb[tok.encode(TPL.format(**dict(zip(FIELDS,m))),add_special_tokens=False)+tails[0][:sn]][None]\n else:\n pre=torch.cat([cod(torch.stack([lex(w) for w in m])[None].to(dev))[0],Emb[tails[0][:sn]]])[None]\n o=lm.model(inputs_embeds=pre,use_cache=True)\n h=lm.model(input_ids=torch.tensor([t[sn:] for t in tails],device=dev),past_key_values=_branch(o.past_key_values),use_cache=True).last_hidden_state[:,-1]\n out=[int(lm.lm_head(h)[j][CAND[j].to(dev)].argmax()) for j in range(4)]\n return pre.shape[1]+4*len(tails[0][sn:])\ndef latency(cod,dev,n=40,reps=3):\n sync=(lambda: torch.cuda.synchronize()) if dev!=\'cpu\' else (lambda:None)\n ev=FRESH[\'ood\'][:n]; T={\'text\':[],\'codec\':[]}; P={}\n for m in ev[:5]: one_event(\'text\',m,cod,dev); one_event(\'codec\',m,cod,dev)\n rr=random.Random(3)\n for _ in range(reps):\n for m in ev:\n for r in rr.sample([\'text\',\'codec\'],2):\n sync(); t0=time.perf_counter(); P[r]=one_event(r,m,cod,dev); sync(); T[r].append(1e3*(time.perf_counter()-t0))\n q=lambda x:sorted(x)[int(.9*len(x))]\n return {k:dict(median_ms=st.median(v),p90_ms=q(v),positions=P[k]) for k,v in T.items()}\n\nclass PackK(nn.Module): # superposición aprendida: ranura_k = sum_j a_kj (e_j + U_j V_j^T e_j + rol_j) + MLP(concat)\n def __init__(s,K,seed,r=64):\n super().__init__(); torch.manual_seed(seed); s.K=K\n s.roles=nn.Parameter(torch.randn(4,D)*.02); s.U=nn.Parameter(torch.randn(4,D,r)*.01); s.V=nn.Parameter(torch.randn(4,D,r)*.01)\n s.a=nn.Parameter(torch.full((K,4),1.0/4*K if K<4 else 1.0))\n s.net=nn.Sequential(nn.LayerNorm(4*D),nn.Linear(4*D,1024),nn.GELU(),nn.Linear(1024,K*D)); nn.init.normal_(s.net[-1].weight,std=.002); nn.init.zeros_(s.net[-1].bias)\n s.register_buffer(\'sc\',Emb.std().clone())\n def forward(s,L):\n c=L+torch.einsum(\'bjd,jdr,jer->bje\',L,s.V,s.U)+s.roles[None]\n base=torch.einsum(\'kj,bjd->bkd\',s.a,c)\n return base+s.sc*s.net(c.flatten(1)).view(-1,s.K,D)\ndef train_kd(cod,steps,seed,p_rand=.3,bs=12,lr=1e-3):\n g=random.Random(seed); opt=torch.optim.AdamW(cod.parameters(),lr=lr,weight_decay=1e-3)\n sch=torch.optim.lr_scheduler.OneCycleLR(opt,max_lr=lr,total_steps=steps,pct_start=.05)\n for s_ in range(steps):\n msgs=[]\n for _ in range(bs):\n m=list(g.choice(FIT))\n for j in range(4):\n if g.random()<p_rand: m[j]=g.choice(POOL)\n msgs.append(m)\n tpl=g.choice(SUF_TRAIN)\n S=cod(torch.stack([torch.stack([lex(w) for w in m]) for m in msgs]))\n embs=[];hs=[];tg=[]\n for i,m in enumerate(msgs):\n for j in range(4): embs.append(torch.cat([S[i],Emb[sufids(tpl,j)]])); tg.append(t1(m[j])[0])\n with torch.no_grad():\n te=[Emb[tok.encode(TPL.format(**dict(zip(FIELDS,m)))+tpl.format(field=FIELDS[j]),add_special_tokens=False)] for m in msgs for j in range(4)]\n Lt=max(e.shape[0] for e in te); x=torch.zeros(len(te),Lt,D); mk=torch.zeros(len(te),Lt,dtype=torch.long)\n for i,e in enumerate(te): x[i,:e.shape[0]]=e; mk[i,:e.shape[0]]=1\n ht=lm.model(inputs_embeds=x,attention_mask=mk).last_hidden_state[torch.arange(len(te)),mk.sum(1)-1]; lt=lm.lm_head(ht)\n L=max(e.shape[0] for e in embs); x=torch.zeros(len(embs),L,D); mk=torch.zeros(len(embs),L,dtype=torch.long)\n for i,e in enumerate(embs): x[i,:e.shape[0]]=e; mk[i,:e.shape[0]]=1\n hsd=lm.model(inputs_embeds=x,attention_mask=mk).last_hidden_state[torch.arange(len(embs)),mk.sum(1)-1]; ls=lm.lm_head(hsd)\n loss=F.cross_entropy(ls,torch.tensor(tg))+F.kl_div(F.log_softmax(ls,-1),F.log_softmax(lt,-1),log_target=True,reduction=\'batchmean\')+.5*(1-F.cosine_similarity(hsd,ht,-1).mean())\n opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(cod.parameters(),1.0); opt.step(); sch.step()\n if s_%500==0: print(json.dumps({\'step\':s_,\'loss\':float(loss)}),flush=True)\n return cod.eval()\n@torch.no_grad()\ndef probe(cod): # anti-lookup fuerte: sonda lineal (ridge, sin LLM) entrenada en slots de FIT, evaluada en FRESH ood y 4_new\n def slots(ms): return cod(torch.stack([torch.stack([lex(w) for w in m]) for m in ms])).flatten(1)\n Xtr=torch.cat([slots(FIT[i:i+100]) for i in range(0,len(FIT),100)]); res={}\n for name,ms in ((\'fresh_ood\',FRESH[\'ood\']),(\'transfer_4_new\',TRANSFER[\'4_new\'])):\n Xte=slots(ms); ok=torch.ones(len(ms),dtype=torch.bool,device=\'cpu\')\n for j in range(4): # regresión al embedding del valor, decodificación por coseno contra opciones (abierta a palabras nuevas)\n Y=torch.stack([lex(m[j]) for m in FIT]); A=torch.linalg.solve(Xtr.T@Xtr+10*torch.eye(Xtr.shape[1]),Xtr.T@Y)\n P=F.normalize(Xte@A,dim=-1)@F.normalize(Emb[CAND[j]],dim=-1).T\n ok&=P.argmax(1).cpu()==torch.tensor([gold_idx(j,m[j]) for m in ms],device=\'cpu\')\n res[\'probe_\'+name]=float(ok.float().mean())\n return res\n\nrv=random.Random(99)\nVAL_T=[[rv.choice(POOL_VAL) for j in range(4)] for _ in range(64)] # validación de transferencia: palabras nunca usadas en entrenamiento ni en test\nCAND_FULL=None\n@torch.no_grad()\ndef acc_open(cod,msgs,bs=16): # exactitud sobre TODO el vocabulario (palabras de validación no están en CAND)\n ok=[]\n for b in range(0,len(msgs),bs):\n mb=msgs[b:b+bs]; lg,mt=codec_logits(cod,mb,SUF_MAIN); good=torch.ones(len(mb),dtype=torch.bool)\n for r,(i,j) in enumerate(mt): good[i]&=int(lg[r].argmax())==t1(mb[i][j])[0]\n ok+=good.tolist()\n return sum(ok)/len(ok)\ndef readout_h(cod,msgs,tpl):\n S=cod(torch.stack([torch.stack([lex(w) for w in m]) for m in msgs])); embs=[]\n for i,m in enumerate(msgs):\n for j in range(4): embs.append(torch.cat([S[i],Emb[sufids(tpl,j)]]))\n L=max(e.shape[0] for e in embs); x=torch.zeros(len(embs),L,D); mk=torch.zeros(len(embs),L,dtype=torch.long)\n for i,e in enumerate(embs): x[i,:e.shape[0]]=e; mk[i,:e.shape[0]]=1\n h=lm.model(inputs_embeds=x,attention_mask=mk).last_hidden_state[torch.arange(len(embs)),mk.sum(1)-1]; return h,lm.lm_head(h)\ndef distill(stu,teach,steps,seed,tagname,bs=12,lr=1e-3,every=250):\n g=random.Random(seed); opt=torch.optim.AdamW(stu.parameters(),lr=lr,weight_decay=1e-3)\n sch=torch.optim.lr_scheduler.OneCycleLR(opt,max_lr=lr,total_steps=steps,pct_start=.05); best=(-1,None,0)\n for s_ in range(1,steps+1):\n msgs=[]\n for _ in range(bs):\n m=list(g.choice(FIT))\n for j in range(4):\n u=g.random()\n if u<.1: m[j]=g.choice(POOL)+\' \'+g.choice(POOL)\n elif u<.5: m[j]=g.choice(POOL)\n msgs.append(m)\n tpl=g.choice(SUF_TRAIN)\n with torch.no_grad(): ht,lt=readout_h(teach,msgs,tpl)\n hs,ls=readout_h(stu,msgs,tpl); tg=torch.tensor([t1(m[j])[0] for m in msgs for j in range(4)])\n loss=F.cross_entropy(ls,tg)+F.kl_div(F.log_softmax(ls,-1),F.log_softmax(lt,-1),log_target=True,reduction=\'batchmean\')+.5*(1-F.cosine_similarity(hs,ht,-1).mean())\n opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(stu.parameters(),1.0); opt.step(); sch.step()\n if s_%every==0:\n stu.eval(); v1=acc(\'codec\',VAL[:96],cod=stu); v2=acc_open(stu,VAL_T); stu.train()\n sc=v1+v2; log({\'run\':tagname,\'step\':s_,\'val_contract\':v1,\'val_transfer_open\':v2})\n if sc>best[0]: best=(sc,{k:v.detach().clone() for k,v in stu.state_dict().items()},s_)\n stu.load_state_dict(best[1]); stu.eval(); return stu,best[2]\ndef grow_down(src,K): # inicializar K pequeño desde K mayor (curriculum)\n dst=PackK(K,0); sd=src.state_dict(); dd=dst.state_dict()\n for k in (\'roles\',\'U\',\'V\',\'net.0.weight\',\'net.0.bias\',\'net.1.weight\',\'net.1.bias\',\'sc\'): dd[k]=sd[k].clone()\n Ks=src.K; dd[\'a\']=sd[\'a\'].view(K,Ks//K,4).sum(1) if Ks%K==0 else sd[\'a\'].mean(0,keepdim=True).expand(K,4).clone()\n W=sd[\'net.3.weight\'].view(Ks,D,-1); b=sd[\'net.3.bias\'].view(Ks,D)\n dd[\'net.3.weight\']=W.view(K,Ks//K,D,-1).sum(1).reshape(K*D,-1); dd[\'net.3.bias\']=b.view(K,Ks//K,D).sum(1).reshape(-1)\n dst.load_state_dict(dd); return dst\n@torch.no_grad()\ndef batched_latency(cod,bsz,reps=5):\n sync=(lambda: torch.cuda.synchronize()) if DEV==\'cuda\' else (lambda:None)\n ms=FRESH[\'ood\'][:bsz]; res={}\n for route in (\'text\',\'codec\'):\n for _ in range(2): (text_logits(ms,SUF_MAIN) if route==\'text\' else codec_logits(cod,ms,SUF_MAIN))\n T=[]\n for _ in range(reps):\n sync(); t0=time.perf_counter(); (text_logits(ms,SUF_MAIN) if route==\'text\' else codec_logits(cod,ms,SUF_MAIN)); sync(); T.append(1e3*(time.perf_counter()-t0))\n res[route]=dict(median_ms=st.median(T),max_ms=max(T))\n return res\n\n# validación de transferencia RESTRINGIDA: opciones = vocab del campo + 40 palabras reservadas por campo (disjuntas de train y test)\nCANDV=[];VPOOL=[]\nfor j in range(4):\n ws=POOL_VAL[j*50:(j+1)*50]; ids=sorted({t1(w)[0] for w in V[j]}|{t1(w)[0] for w in ws}); CANDV.append(torch.tensor(ids)); VPOOL.append(ws)\nrv2=random.Random(123); VAL_TR=[[rv2.choice(VPOOL[j]) for j in range(4)] for _ in range(96)]\n@torch.no_grad()\ndef acc_valT(cod,bs=24):\n ok=[]\n for b in range(0,len(VAL_TR),bs):\n mb=VAL_TR[b:b+bs]; lg,mt=codec_logits(cod,mb,SUF_MAIN); good=torch.ones(len(mb),dtype=torch.bool)\n for r,(i,j) in enumerate(mt): good[i]&=int(CANDV[j][lg[r][CANDV[j]].argmax()])==t1(mb[i][j])[0]\n ok+=good.tolist()\n return sum(ok)/len(ok)\ndef distill2(stu,teach,steps,seed,tagname,bs=12,lr=5e-4,every=250):\n g=random.Random(seed+1000); opt=torch.optim.AdamW(stu.parameters(),lr=lr,weight_decay=1e-3)\n sch=torch.optim.lr_scheduler.OneCycleLR(opt,max_lr=lr,total_steps=steps,pct_start=.05)\n stu.eval(); v1=acc(\'codec\',VAL[:96],cod=stu); v2=acc_valT(stu); stu.train()\n best=(v1+v2,{k:v.detach().clone() for k,v in stu.state_dict().items()},0); log({\'run\':tagname,\'step\':0,\'val_contract\':v1,\'val_transfer_restricted\':v2})\n for s_ in range(1,steps+1):\n msgs=[]\n for _ in range(bs):\n m=list(g.choice(FIT))\n for j in range(4):\n u=g.random()\n if u<.1: m[j]=g.choice(POOL)+\' \'+g.choice(POOL)\n elif u<.4: m[j]=g.choice(POOL)\n msgs.append(m)\n tpl=g.choice(SUF_TRAIN)\n with torch.no_grad(): ht,lt=readout_h(teach,msgs,tpl)\n hs,ls=readout_h(stu,msgs,tpl); tg=torch.tensor([t1(m[j])[0] for m in msgs for j in range(4)])\n loss=F.cross_entropy(ls,tg)+F.kl_div(F.log_softmax(ls,-1),F.log_softmax(lt,-1),log_target=True,reduction=\'batchmean\')+.5*(1-F.cosine_similarity(hs,ht,-1).mean())\n opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(stu.parameters(),1.0); opt.step(); sch.step()\n if s_%every==0:\n stu.eval(); v1=acc(\'codec\',VAL[:96],cod=stu); v2=acc_valT(stu); stu.train()\n log({\'run\':tagname,\'step\':s_,\'val_contract\':v1,\'val_transfer_restricted\':v2})\n if v1+v2>best[0]: best=(v1+v2,{k:v.detach().clone() for k,v in stu.state_dict().items()},s_)\n stu.load_state_dict(best[1]); stu.eval(); return stu,best[2]\n@torch.no_grad()\ndef batched_shared(cod,bsz,reps=5): # prefijo KV compartido en batch, ambas rutas\n sync=(lambda: torch.cuda.synchronize()) if DEV==\'cuda\' else (lambda:None)\n ms=FRESH[\'ood\'][:bsz]; tails=[sufids(SUF_MAIN,j) for j in range(4)]; sn=0\n while all(t[sn]==tails[0][sn] for t in tails): sn+=1\n TT=torch.tensor([t[sn:] for t in tails]).repeat(bsz,1)\n def run(route):\n if route==\'text\':\n ids=[tok.encode(TPL.format(**dict(zip(FIELDS,m))),add_special_tokens=False)+tails[0][:sn] for m in ms]\n assert len(set(map(len,ids)))==1; pre=Emb[torch.tensor(ids)]\n else:\n S=cod(torch.stack([torch.stack([lex(w) for w in m]) for m in ms])); pre=torch.cat([S,Emb[tails[0][:sn]][None].expand(bsz,-1,-1)],1)\n o=lm.model(inputs_embeds=pre,use_cache=True)\n h=lm.model(input_ids=TT,past_key_values=_branch(o.past_key_values),use_cache=True).last_hidden_state[:,-1]\n lg=lm.lm_head(h); return pre.shape[1]\n res={}\n for route in (\'text\',\'codec\'):\n for _ in range(2): run(route)\n T=[]\n for _ in range(reps): sync(); t0=time.perf_counter(); p=run(route); sync(); T.append(1e3*(time.perf_counter()-t0))\n res[route]=dict(median_ms=st.median(T),max_ms=max(T),prefix_positions=p)\n return res\nif __name__==\'__main__\':\n INP=os.environ.get(\'INP\',\'\'); seeds=[int(x) for x in os.environ.get(\'SEEDS\',\'17,29,43\').split(\',\')]\n for sd in seeds:\n t0=time.time(); teach=Binder(sd); teach.load_state_dict(torch.load(f\'{INP}/binder4_1000_s{sd}.pt\',map_location=DEV)); teach.eval().requires_grad_(False)\n k2=PackK(2,sd); k2.load_state_dict(torch.load(f\'{INP}/k2_s{sd}.pt\',map_location=DEV))\n log({\'name\':f\'k2_017_s{sd}\',\'val_transfer_restricted\':acc_valT(k2.eval())})\n k2,b=distill2(k2,teach,int(os.environ.get(\'S2\',3000)),sd,f\'k2long_s{sd}\'); torch.save(k2.state_dict(),f\'{OD}/k2long_s{sd}.pt\')\n r=evaluate_all(f\'k2long_s{sd}\',k2); log({\'name\':r[\'name\'],\'best_step\':b,\'train_s\':round(time.time()-t0),**probe(k2)})\n if sd==seeds[0]:\n for bz in (1,16,64): log({\'latency_shared\':\'k2long\',\'batch\':bz,\'gpu\':torch.cuda.get_device_name() if DEV==\'cuda\' else \'cpu\',**batched_shared(k2,bz)})\n')
subprocess.run([sys.executable,'/kaggle/working/run018.py'],env=dict(os.environ,SRC=S,MP=mp,OD='/kaggle/working',INP=inp))
os.system('rm -rf /kaggle/working/m135')
|