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/**/expanded.json',recursive=True)[0])
S='/kaggle/working/src'; os.makedirs(S+'/data',exist_ok=True)
os.system(f'cp {inp}/expanded.json {S}/data/')
for s in (17,29,43):
    os.makedirs(f'{S}/checkpoints/expanded_role_ce_k4_seed{s}',exist_ok=True); os.system(f'cp {inp}/ck{s}/best.pt {S}/checkpoints/expanded_role_ce_k4_seed{s}/ 2>/dev/null || (cd {S}/checkpoints/expanded_role_ce_k4_seed{s} && unzip -o -q {inp}/ck{s}.zip)')
os.system(f'ls -R {inp} | head -20')
from transformers import AutoTokenizer,AutoModelForCausalLM
R='12fd25f77366fa6b3b4b768ec3050bf629380bac'; MP='/kaggle/working/model'
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/run015.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=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,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+\'/results015.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\nif __name__==\'__main__\' and False:\n    stage=sys.argv[1]\n    if stage==\'base\':\n        log({\'info\':\'vocab\',**{f:NEW[f] for f in FIELDS},\'phrases\':PHRASE,\'pool\':len(POOL),\'fresh\':{k:len(v) for k,v in FRESH.items()}})\n        evaluate_all(\'texto\')\n        for sd in (17,29,43): evaluate_all(f\'014_seed{sd}\',Their014(sd))\n    else:\n        kind,steps=stage.split(\':\'); steps=int(steps)\n        for sd in (17,29,43):\n            t0=time.time(); cod=(Binder if kind==\'binder4\' else Pack1)(sd); cod=train(cod,steps,sd)\n            torch.save(cod.state_dict(),f\'{OD}/{kind}_{steps}_s{sd}.pt\')\n            r=evaluate_all(f\'{kind}_{steps}_s{sd}\',cod); log({\'name\':r[\'name\'],\'train_s\':round(time.time()-t0)})\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\nif __name__==\'__main__\':\n    for stage in sys.argv[1:]:\n        if stage==\'base\':\n            log({\'info\':\'vocab\',**{f:NEW[f] for f in FIELDS},\'phrases\':PHRASE,\'pool\':len(POOL),\'fresh\':{k:len(v) for k,v in FRESH.items()},\'device\':DEV,\'gpu\':torch.cuda.get_device_name() if DEV==\'cuda\' else \'cpu\'})\n            evaluate_all(\'texto\')\n            for sd in (17,29,43): evaluate_all(f\'014_seed{sd}\',Their014(sd))\n        elif stage.startswith(\'lat\'):\n            _,kind,steps=stage.split(\':\'); cods={\'014\':Their014(17)}\n            if kind!=\'014\': c=(Binder if kind==\'binder4\' else Pack1)(17); c.load_state_dict(torch.load(f\'{OD}/{kind}_{steps}_s17.pt\',map_location=DEV)); cods={kind:c.eval()}\n            for k,c in cods.items(): log({\'latency\':k,\'steps\':steps,\'device\':DEV,\'gpu\':torch.cuda.get_device_name() if DEV==\'cuda\' else \'cpu\',**latency(c,DEV)})\n        else:\n            kind,steps=stage.split(\':\'); steps=int(steps)\n            for sd in (17,29,43):\n                t0=time.time(); cod=(Binder if kind==\'binder4\' else Pack1)(sd); cod=train(cod,steps,sd)\n                torch.save(cod.state_dict(),f\'{OD}/{kind}_{steps}_s{sd}.pt\')\n                r=evaluate_all(f\'{kind}_{steps}_s{sd}_{DEV}\',cod); log({\'name\':r[\'name\'],\'train_s\':round(time.time()-t0)})\n')
env=dict(os.environ,SRC=S,MP=MP,OD='/kaggle/working')
subprocess.run([sys.executable,'/kaggle/working/run015.py','base','binder4:1000','pack1:1000','lat:014:0','lat:binder4:1000','lat:pack1:1000'],env=env)