# CUERDAS-015: transferencia (valores/consultas nuevas), vocabulario aleatorizado, K=1 vs K=4. CPU. import json,random,time,re,sys,statistics as st,itertools import torch,torch.nn as nn,torch.nn.functional as F from transformers import AutoModelForCausalLM,AutoTokenizer,AutoConfig torch.set_num_threads(8) import os SRC=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') DEV='cuda' if torch.cuda.is_available() else 'cpu' tok=AutoTokenizer.from_pretrained(MP); lm=AutoModelForCausalLM.from_pretrained(MP,torch_dtype=torch.float32).eval().requires_grad_(False).to(DEV) if DEV=='cuda': torch.set_default_device('cuda') Emb=lm.get_input_embeddings().weight; D=Emb.shape[1] meta=json.load(open(f'{SRC}/data/expanded.json')); V=meta['vocabulary']; FIELDS=['action','target','status','window'] rows=meta['rows']; SPL=meta['splits'] TPL="The action is {action}. The target is {target}. The status is {status}. The window is {window}." SUF_MAIN="\nAccording to the record, the {field} is" SUF_TRAIN=[SUF_MAIN,"\nThe {field} is","\nIn this record, the {field} is","\nQuestion: what is the {field}? Answer:"] SUF_HELD={"recorded":"\nThe recorded {field} is","value":"\nThe {field} value is","past":"\nThe {field} was"} def t1(w): return tok.encode(' '+w,add_special_tokens=False) # ---- vocabularios nuevos (nunca en entrenamiento) ---- NEW={'action':['inspect','repair','cancel','approve','archive','scan','print','send'], 'target':['sensor','battery','printer','server','truck','door','screen','cable'], 'status':['broken','stable','late','empty','full','warm','missing','urgent'], 'window':['night','evening','weekend','noon','spring','summer','winter','midnight']} PHRASE={'action':['shut down','back up','sign off','log out'],'target':['air filter','fuel tank','water heater','hard drive'], 'status':['out of stock','on hold','under review','half full'],'window':['next week','last night','this month','later today']} # limpiar: primer token único dentro de (vocab original ∪ nuevos ∪ frases) por campo for j,f in enumerate(FIELDS): used={t1(w)[0] for w in V[j]}; keep=[];keepp=[] for w in NEW[f]: 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]) for w in PHRASE[f]: if len(t1(w))>1 and t1(w)[0] not in used: keepp.append(w); used.add(t1(w)[0]) NEW[f]=keep; PHRASE[f]=keepp BANNED={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} # pool de palabras aleatorias para entrenamiento (meta-aprendizaje de binding) POOL=[] for t,i in tok.get_vocab().items(): if t.startswith('Ġ') and re.fullmatch(r'[a-z]{3,9}',t[1:]) and i not in BANNED: POOL.append(t[1:]) random.Random(0).shuffle(POOL); POOL=POOL[:3000] # ---- splits: los de 014 + test FRESCO nuevo del universo no usado ---- used_vals={tuple(r['values']) for r in rows} fit=[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} rng=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) fresh_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] def W(v): return [V[j][v[j]] for j in range(4)] FIT=[W(v) for v in fit]; VAL=[W(tuple(rows[i]['values'])) for i in SPL['validation']] TEST014={'id':[W(tuple(rows[i]['values'])) for i in SPL['id_test']],'ood':[W(tuple(rows[i]['values'])) for i in SPL['ood_test']]} FRESH={'id':[W(v) for v in fresh_id],'ood':[W(v) for v in fresh_ood]} r2=random.Random(77) def novel_set(nnew,n=96,phr=False): # nnew campos con valor nuevo (los demás del vocab), tomados de FRESH out=[] for k in range(n): base=list(FRESH['ood'][k%len(FRESH['ood'])]); js=r2.sample(range(4),nnew) for j in js: base[j]=r2.choice(PHRASE[FIELDS[j]] if phr else NEW[FIELDS[j]]) out.append(base) return out TRANSFER={'1_new':novel_set(1),'2_new':novel_set(2),'4_new':novel_set(4),'1_phrase':novel_set(1,phr=True)} CAND=[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) CAND=[torch.tensor(c) for c in CAND] def lex(w): return Emb[t1(w)].mean(0) def gold_idx(j,w): return int((CAND[j]==t1(w)[0]).nonzero()[0]) # ---- codecs ---- class Binder(nn.Module): # K=4: ranura por campo, binding de cualquier palabra a un rol def __init__(s,seed): super().__init__(); torch.manual_seed(seed); s.K=4 anc=torch.stack([Emb[t1(f)[0]] for f in FIELDS]); s.roles=nn.Parameter(anc.clone()) 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) s.register_buffer('sc',Emb.std().clone()) def forward(s,L): c=L+s.roles[None]; return c+s.sc*s.net(c) # L: [B,4,D] class Pack1(nn.Module): # K=1: empaqueta los 4 campos en UNA posición def __init__(s,seed): super().__init__(); torch.manual_seed(seed); s.K=1 s.roles=nn.Parameter(torch.randn(4,D)*.02) 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()) s.mix=nn.Parameter(torch.full((4,),.25)) def forward(s,L): c=L+s.roles[None]; base=(s.mix[None,:,None]*c).sum(1) return (base+s.sc*s.net(c.flatten(1)))[:,None] class Their014(nn.Module): # su RoleAdapter, con léxico extensible (mismo cálculo que su forward) def __init__(s,seed): super().__init__(); ck=torch.load(f'{SRC}/checkpoints/expanded_role_ce_k4_seed{seed}/best.pt',map_location=DEV,weights_only=False)['adapter_state'] s.K=4; s.roles=nn.Parameter(ck['roles']); s.sc=ck['scale'] 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.')}) def forward(s,L): c=L+s.roles[None]; return c+s.sc*s.net(c) # ---- lectura por el LLM congelado ---- def sufids(tpl,j): return tok.encode(tpl.format(field=FIELDS[j]),add_special_tokens=False) def run_rows(embs,tails): # listas de tensores [L,D] -> logits última posición (padding a la derecha) L=max(e.shape[0] for e in embs); B=len(embs) x=torch.zeros(B,L,D); m=torch.zeros(B,L,dtype=torch.long) for i,e in enumerate(embs): x[i,:e.shape[0]]=e; m[i,:e.shape[0]]=1 h=lm.model(inputs_embeds=x,attention_mask=m).last_hidden_state[torch.arange(B),m.sum(1)-1] return lm.lm_head(h) def codec_logits(cod,msgs,tpl,js=range(4)): S=cod(torch.stack([torch.stack([lex(w) for w in m]) for m in msgs])) embs=[];meta_=[] for i,m in enumerate(msgs): for j in js: embs.append(torch.cat([S[i],Emb[sufids(tpl,j)]])); meta_.append((i,j)) return run_rows(embs,None),meta_ def text_logits(msgs,tpl): embs=[];meta_=[] for i,m in enumerate(msgs): 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)) return run_rows(embs,None),meta_ @torch.no_grad() def acc(route,msgs,tpl=SUF_MAIN,cod=None,bs=24): ok=[] for b in range(0,len(msgs),bs): mb=msgs[b:b+bs]; lg,mt=(text_logits(mb,tpl) if route=='text' else codec_logits(cod,mb,tpl)) good=torch.ones(len(mb),dtype=torch.bool) for r,(i,j) in enumerate(mt): good[i]&=int(lg[r][CAND[j]].argmax())==gold_idx(j,mb[i][j]) ok+=good.tolist() return sum(ok)/len(ok) @torch.no_grad() def nolm(cod,msgs): # anti-lookup: decodificar valor desde las ranuras por coseno, sin LLM S=cod(torch.stack([torch.stack([lex(w) for w in m]) for m in msgs])); ok=0 for i,m in enumerate(msgs): g=True for j in range(4): cands=F.normalize(Emb[CAND[j]],dim=-1); s=S[i,j] if S.shape[1]==4 else S[i].mean(0) g&=int((F.normalize(s,dim=-1)@cands.T).argmax())==gold_idx(j,m[j]) ok+=g return ok/len(msgs) def train(cod,steps,seed,p_rand=.5,p_phr=.15,bs=12): g=random.Random(seed); opt=torch.optim.AdamW(cod.parameters(),lr=2e-3,weight_decay=1e-3) for s in range(steps): msgs=[] for _ in range(bs): m=list(g.choice(FIT)) for j in range(4): u=g.random() if ubje',L,s.V,s.U)+s.roles[None] base=torch.einsum('kj,bjd->bkd',s.a,c) return base+s.sc*s.net(c.flatten(1)).view(-1,s.K,D) def train_kd(cod,steps,seed,p_rand=.3,bs=12,lr=1e-3): g=random.Random(seed); opt=torch.optim.AdamW(cod.parameters(),lr=lr,weight_decay=1e-3) sch=torch.optim.lr_scheduler.OneCycleLR(opt,max_lr=lr,total_steps=steps,pct_start=.05) for s_ in range(steps): msgs=[] for _ in range(bs): m=list(g.choice(FIT)) for j in range(4): if g.random()