| |
| 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,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+'/results015.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 |
| if __name__=='__main__' and False: |
| stage=sys.argv[1] |
| if stage=='base': |
| 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()}}) |
| evaluate_all('texto') |
| for sd in (17,29,43): evaluate_all(f'014_seed{sd}',Their014(sd)) |
| else: |
| kind,steps=stage.split(':'); steps=int(steps) |
| for sd in (17,29,43): |
| t0=time.time(); cod=(Binder if kind=='binder4' else Pack1)(sd); cod=train(cod,steps,sd) |
| torch.save(cod.state_dict(),f'{OD}/{kind}_{steps}_s{sd}.pt') |
| r=evaluate_all(f'{kind}_{steps}_s{sd}',cod); log({'name':r['name'],'train_s':round(time.time()-t0)}) |
| 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()} |
|
|
| if __name__=='__main__': |
| for stage in sys.argv[1:]: |
| if stage=='base': |
| 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'}) |
| evaluate_all('texto') |
| for sd in (17,29,43): evaluate_all(f'014_seed{sd}',Their014(sd)) |
| elif stage.startswith('lat'): |
| _,kind,steps=stage.split(':'); cods={'014':Their014(17)} |
| 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()} |
| 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)}) |
| else: |
| kind,steps=stage.split(':'); steps=int(steps) |
| for sd in (17,29,43): |
| t0=time.time(); cod=(Binder if kind=='binder4' else Pack1)(sd); cod=train(cod,steps,sd) |
| torch.save(cod.state_dict(),f'{OD}/{kind}_{steps}_s{sd}.pt') |
| r=evaluate_all(f'{kind}_{steps}_s{sd}_{DEV}',cod); log({'name':r['name'],'train_s':round(time.time()-t0)}) |
|
|
