| |
| 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) |
| |
| 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']} |
| |
| 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=[] |
| 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] |
| |
| 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): |
| 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)] |
| 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]) |
| |
| class Binder(nn.Module): |
| 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) |
| class Pack1(nn.Module): |
| 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): |
| 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) |
| |
| def sufids(tpl,j): return tok.encode(tpl.format(field=FIELDS[j]),add_special_tokens=False) |
| def run_rows(embs,tails): |
| 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): |
| 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 u<p_phr: m[j]=g.choice(POOL)+' '+g.choice(POOL) |
| elif u<p_phr+p_rand: m[j]=g.choice(POOL) |
| msgs.append(m) |
| tpl=g.choice(SUF_TRAIN); lg,mt=codec_logits(cod,msgs,tpl) |
| gold=torch.tensor([t1(msgs[i][j])[0] for i,j in mt]); loss=F.cross_entropy(lg,gold) |
| opt.zero_grad(); loss.backward(); opt.step() |
| return cod.eval() |
| OUT=OD+'/results016.jsonl' |
| def log(d): open(OUT,'a').write(json.dumps(d)+"\n"); print(json.dumps(d),flush=True) |
| def evaluate_all(name,cod=None,quick=False): |
| route='text' if cod is None else 'codec'; r={'name':name} |
| r['val']=acc(route,VAL[:96],cod=cod) |
| if not quick: |
| for k,v in FRESH.items(): r['fresh_'+k]=acc(route,v,cod=cod) |
| for k,v in TRANSFER.items(): r['transfer_'+k]=acc(route,v,cod=cod) |
| sub=FRESH['id'][:48]+FRESH['ood'][:48] |
| for k,t in SUF_HELD.items(): r['query_'+k]=acc(route,sub,t,cod=cod) |
| if cod is not None: r['noLLM_fresh_ood']=nolm(cod,FRESH['ood'][:200]); r['noLLM_transfer_4new']=nolm(cod,TRANSFER['4_new']) |
| log(r); return r |
|
|
| def _branch(cache): |
| if hasattr(cache,'to_legacy_cache'): cache=cache.to_legacy_cache() |
| return tuple((k.repeat_interleave(4,0),v.repeat_interleave(4,0)) for k,v in cache) |
| @torch.no_grad() |
| def one_event(route,m,cod,dev): |
| tails=[sufids(SUF_MAIN,j) for j in range(4)]; sn=0 |
| while all(t[sn]==tails[0][sn] for t in tails): sn+=1 |
| if route=='text': |
| pre=Emb[tok.encode(TPL.format(**dict(zip(FIELDS,m))),add_special_tokens=False)+tails[0][:sn]][None] |
| else: |
| pre=torch.cat([cod(torch.stack([lex(w) for w in m])[None].to(dev))[0],Emb[tails[0][:sn]]])[None] |
| o=lm.model(inputs_embeds=pre,use_cache=True) |
| 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] |
| out=[int(lm.lm_head(h)[j][CAND[j].to(dev)].argmax()) for j in range(4)] |
| return pre.shape[1]+4*len(tails[0][sn:]) |
| def latency(cod,dev,n=40,reps=3): |
| sync=(lambda: torch.cuda.synchronize()) if dev!='cpu' else (lambda:None) |
| ev=FRESH['ood'][:n]; T={'text':[],'codec':[]}; P={} |
| for m in ev[:5]: one_event('text',m,cod,dev); one_event('codec',m,cod,dev) |
| rr=random.Random(3) |
| for _ in range(reps): |
| for m in ev: |
| for r in rr.sample(['text','codec'],2): |
| sync(); t0=time.perf_counter(); P[r]=one_event(r,m,cod,dev); sync(); T[r].append(1e3*(time.perf_counter()-t0)) |
| q=lambda x:sorted(x)[int(.9*len(x))] |
| return {k:dict(median_ms=st.median(v),p90_ms=q(v),positions=P[k]) for k,v in T.items()} |
|
|
| class PackK(nn.Module): |
| def __init__(s,K,seed,r=64): |
| super().__init__(); torch.manual_seed(seed); s.K=K |
| 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) |
| s.a=nn.Parameter(torch.full((K,4),1.0/4*K if K<4 else 1.0)) |
| 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) |
| s.register_buffer('sc',Emb.std().clone()) |
| def forward(s,L): |
| c=L+torch.einsum('bjd,jdr,jer->bje',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()<p_rand: m[j]=g.choice(POOL) |
| msgs.append(m) |
| tpl=g.choice(SUF_TRAIN) |
| S=cod(torch.stack([torch.stack([lex(w) for w in m]) for m in msgs])) |
| embs=[];hs=[];tg=[] |
| for i,m in enumerate(msgs): |
| for j in range(4): embs.append(torch.cat([S[i],Emb[sufids(tpl,j)]])); tg.append(t1(m[j])[0]) |
| with torch.no_grad(): |
| 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)] |
| 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) |
| for i,e in enumerate(te): x[i,:e.shape[0]]=e; mk[i,:e.shape[0]]=1 |
| 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) |
| 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) |
| for i,e in enumerate(embs): x[i,:e.shape[0]]=e; mk[i,:e.shape[0]]=1 |
| 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) |
| 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()) |
| opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(cod.parameters(),1.0); opt.step(); sch.step() |
| if s_%500==0: print(json.dumps({'step':s_,'loss':float(loss)}),flush=True) |
| return cod.eval() |
| @torch.no_grad() |
| def probe(cod): |
| def slots(ms): return cod(torch.stack([torch.stack([lex(w) for w in m]) for m in ms])).flatten(1) |
| Xtr=torch.cat([slots(FIT[i:i+100]) for i in range(0,len(FIT),100)]); res={} |
| for name,ms in (('fresh_ood',FRESH['ood']),('transfer_4_new',TRANSFER['4_new'])): |
| Xte=slots(ms); ok=torch.ones(len(ms),dtype=torch.bool,device='cpu') |
| for j in range(4): |
| 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) |
| P=F.normalize(Xte@A,dim=-1)@F.normalize(Emb[CAND[j]],dim=-1).T |
| ok&=P.argmax(1).cpu()==torch.tensor([gold_idx(j,m[j]) for m in ms],device='cpu') |
| res['probe_'+name]=float(ok.float().mean()) |
| return res |
| if __name__=='__main__': |
| tag=os.environ.get('TAG','135m'); INP=os.environ.get('INP','') |
| for stage in sys.argv[1:]: |
| kind,K,steps,seeds=stage.split(':'); K=int(K); steps=int(steps) |
| for sd in map(int,seeds.split(',')): |
| if kind=='load': |
| cod=PackK(K,sd); cod.load_state_dict(torch.load(f'{INP}/pack{K}_{steps}_{tag}_s{sd}.pt',map_location=DEV)); cod.eval() |
| log({'name':f'pack{K}_{steps}_{tag}_s{sd}',**probe(cod)}) |
| log({'latency':f'pack{K}_{tag}_s{sd}','gpu':torch.cuda.get_device_name() if DEV=='cuda' else 'cpu',**latency(cod,DEV)}) |
| else: |
| t0=time.time(); cod=train_kd(PackK(K,sd),steps,sd); torch.save(cod.state_dict(),f'{OD}/pack{K}_{steps}_{tag}_s{sd}.pt') |
| r=evaluate_all(f'pack{K}_{steps}_{tag}_s{sd}',cod); log({'name':r['name'],'train_s':round(time.time()-t0)}) |
| log({'name':r['name'],**probe(cod)}) |
|
|
