| |
| import torch, time, random, json, itertools, os, copy, statistics as st |
| import torch.nn as nn, torch.nn.functional as F |
| from transformers import AutoTokenizer, AutoModelForCausalLM |
| torch.manual_seed(0); random.seed(0); torch.set_num_threads(8) |
| M=os.environ.get("MODEL","HuggingFaceTB/SmolLM2-135M") |
| tok=AutoTokenizer.from_pretrained(M); lm=AutoModelForCausalLM.from_pretrained(M,dtype=torch.float32).eval() |
| for p in lm.parameters(): p.requires_grad=False |
| E=lm.get_input_embeddings(); D=E.weight.shape[1] |
| A=["open","close","move","copy","delete","send","check"] |
| T=["file","door","box","email","report","car","light"] |
| S=["new","old","red","blue","hot","cold","big","small"] |
| W=["today","monday","friday","night","morning"] |
| FN=["action","target","state","window"]; FL=[A,T,S,W] |
| msgs=list(itertools.product(range(7),range(7),range(8),range(5))) |
| at=list(itertools.product(range(7),range(7))); sw=list(itertools.product(range(8),range(5))) |
| random.shuffle(at); random.shuffle(sw) |
| H_AT=at[:9]; H_SW=sw[:12] |
| CAL_AT=set(H_AT[:3]); TST_AT=set(H_AT[3:]); CAL_SW=set(H_SW[:4]); TST_SW=set(H_SW[4:]) |
| def covered(m): return (m[0],m[1]) not in set(H_AT) and (m[2],m[3]) not in set(H_SW) |
| cov=[m for m in msgs if covered(m)]; random.shuffle(cov) |
| cov_train=cov[:672]; cov_cal=cov[672:896]; cov_test=cov[896:] |
| def ugroup(m): |
| tags={("c" if (m[0],m[1]) in CAL_AT else "t") for _ in [0] if (m[0],m[1]) in set(H_AT)} |
| if (m[2],m[3]) in set(H_SW): tags.add("c" if (m[2],m[3]) in CAL_SW else "t") |
| return tags |
| unc=[m for m in msgs if not covered(m)] |
| unc_cal=[m for m in unc if ugroup(m)=={"c"}]; unc_test=[m for m in unc if ugroup(m)=={"t"}] |
| SPL=dict(cov_train=len(cov_train),cov_cal=len(cov_cal),cov_test=len(cov_test),unc_cal=len(unc_cal),unc_test=len(unc_test),unc_discarded_mixed=len(unc)-len(unc_cal)-len(unc_test)) |
| print(SPL,flush=True) |
| FORMATS={ |
| "semicolon": lambda m:f"{A[m[0]]};{T[m[1]]};{S[m[2]]};{W[m[3]]}", |
| "bare_words": lambda m:f"{A[m[0]]} {T[m[1]]} {S[m[2]]} {W[m[3]]}", |
| "key_value": lambda m:f"action={A[m[0]]} target={T[m[1]]} state={S[m[2]]} window={W[m[3]]}", |
| "json": lambda m:json.dumps(dict(action=A[m[0]],target=T[m[1]],state=S[m[2]],window=W[m[3]])), |
| "natural": lambda m:f"{A[m[0]]} the {S[m[2]]} {T[m[1]]} {W[m[3]]}.", |
| "natural_labeled": lambda m:f"Action: {A[m[0]]}. Target: {T[m[1]]}. State: {S[m[2]]}. Window: {W[m[3]]}.", |
| } |
| def ids(s): return tok(s,add_special_tokens=False,return_tensors="pt").input_ids[0] |
| cand=[torch.tensor([ids(" "+w)[0].item() for w in L]) for L in FL] |
| for c in cand: assert len(set(c.tolist()))==len(c) |
| Q=[ids(f"\nQ: what is the {f}?\nA:") for f in FN]; PRE=ids("Message:") |
| QE=None |
| def readout(msg_emb): |
| B=msg_emb.shape[0] |
| if torch.is_grad_enabled(): |
| return [lm(inputs_embeds=torch.cat([E(PRE).expand(B,-1,-1),msg_emb,E(Q[f]).expand(B,-1,-1)],1)).logits[:,-1][:,cand[f]] for f in range(4)]; x=torch.cat([E(PRE).expand(B,-1,-1),msg_emb],1) |
| x=torch.cat([E(PRE).expand(B,-1,-1),msg_emb],1) |
| out=lm(inputs_embeds=x,use_cache=True); res=[] |
| for f in range(4): |
| c=copy.deepcopy(out.past_key_values) |
| res.append(lm(inputs_embeds=E(Q[f]).expand(B,-1,-1),past_key_values=c).logits[:,-1][:,cand[f]]) |
| return res |
| @torch.no_grad() |
| def eval_text(fmt,ms,keep_logits=False): |
| from collections import defaultdict |
| g=defaultdict(list); ok={}; lg={} |
| for m in ms: g[len(ids(" "+FORMATS[fmt](m)))].append(m) |
| for L,sub in g.items(): |
| for i in range(0,len(sub),64): |
| b=sub[i:i+64]; r=readout(E(torch.stack([ids(" "+FORMATS[fmt](m)) for m in b]))); t=torch.tensor(b) |
| o=torch.ones(len(b),dtype=torch.bool) |
| for f in range(4): o&=r[f].argmax(1)==t[:,f] |
| for j,m in enumerate(b): |
| ok[m]=bool(o[j]); |
| if keep_logits: lg[m]=[r[f][j] for f in range(4)] |
| return ok,lg |
| R={"model":M,"device":"CPU del box (8 hilos, compartido, sin GPU)","splits":SPL} |
| |
| random.seed(2); calS=random.sample(cov_cal,100)+random.sample(unc_cal,min(100,len(unc_cal))) |
| sweep={} |
| for fmt in FORMATS: |
| t0=time.time(); ok,_=eval_text(fmt,calS) |
| sweep[fmt]=dict(cal_acc=sum(ok.values())/len(ok),avg_tokens=sum(len(ids(" "+FORMATS[fmt](m))) for m in calS)/len(calS),secs=round(time.time()-t0,1)) |
| print(fmt,sweep[fmt],flush=True) |
| R["format_sweep_cal"]=sweep |
| best=max(sweep,key=lambda f:(round(sweep[f]["cal_acc"],2),-sweep[f]["avg_tokens"])) |
| R["teacher_format"]=best; FMT=best; print("teacher",best,flush=True) |
| pos=lambda m:len(ids(" "+FORMATS[FMT](m))) |
| |
| _,TL=eval_text(FMT,cov_train,keep_logits=True) |
| class Adapter(nn.Module): |
| def __init__(s,K): |
| super().__init__(); s.K=K |
| s.emb=nn.ModuleList([nn.Embedding(len(L),64) for L in FL]) |
| s.mlp=nn.Sequential(nn.Linear(256,512),nn.GELU(),nn.Linear(512,K*D)) |
| s.scale=nn.Parameter(torch.tensor(float(E.weight.norm(dim=1).mean()))) |
| def forward(s,mb): |
| h=torch.cat([s.emb[i](mb[:,i]) for i in range(4)],-1) |
| return F.normalize(s.mlp(h).view(-1,s.K,D),dim=-1)*s.scale |
| def train(K,epochs=int(os.environ.get("EPOCHS",8)),bs=32,Tk=2.0): |
| ad=Adapter(K); opt=torch.optim.AdamW(ad.parameters(),3e-3); data=list(cov_train) |
| for ep in range(epochs): |
| random.shuffle(data); tot=0 |
| for i in range(0,len(data),bs): |
| b=data[i:i+bs]; mb=torch.tensor(b); r=readout(ad(mb)); loss=0 |
| for f in range(4): |
| tl=torch.stack([TL[m][f] for m in b]) |
| loss=loss+F.cross_entropy(r[f],mb[:,f])+Tk*Tk*F.kl_div(F.log_softmax(r[f]/Tk,1),F.softmax(tl/Tk,1),reduction="batchmean") |
| opt.zero_grad(); loss.backward(); opt.step(); tot+=loss.item() |
| print(f"K={K} ep{ep} loss {tot:.2f}",flush=True) |
| return ad.eval() |
| @torch.no_grad() |
| def eval_ad(ad,ms): |
| ok={}; conf={} |
| for i in range(0,len(ms),64): |
| b=ms[i:i+64]; t=torch.tensor(b); r=readout(ad(t)) |
| o=torch.ones(len(b),dtype=torch.bool); c=torch.ones(len(b)) |
| for f in range(4): o&=r[f].argmax(1)==t[:,f]; c=torch.minimum(c,F.softmax(r[f],1).max(1).values) |
| for j,m in enumerate(b): ok[m]=bool(o[j]); conf[m]=float(c[j]) |
| return ok,conf |
| cal_all=cov_cal+unc_cal; test_all=cov_test+unc_test |
| tok_cal,_=eval_text(FMT,cal_all); tok_test,_=eval_text(FMT,test_all) |
| R["baseline_test"]=dict(acc=sum(tok_test.values())/len(test_all),pos_avg=sum(map(pos,test_all))/len(test_all),acc_cov=sum(tok_test[m] for m in cov_test)/len(cov_test),acc_unc=sum(tok_test[m] for m in unc_test)/len(unc_test)) |
| print(R["baseline_test"],flush=True) |
| ADS={} |
| for K in (1,2): |
| ts=time.time(); ad=train(K); ADS[K]=ad |
| okc,cc=eval_ad(ad,cal_all); okt,ct=eval_ad(ad,test_all) |
| |
| def hyb(okA,conf,okT,ms,thr): |
| nat=[m for m in ms if covered(m) and conf[m]>=thr]; s=set(nat) |
| acc=sum((okA[m] if m in s else okT[m]) for m in ms)/len(ms) |
| p=sum((K if m in s else pos(m)) for m in ms)/len(ms); return acc,p,len(nat)/len(ms) |
| base_cal=sum(tok_cal.values())/len(cal_all); thr_ok=None |
| for thr in [i/100 for i in range(0,100,2)]: |
| a,_,_=hyb(okc,cc,tok_cal,cal_all,thr) |
| if base_cal-a<=0.01: thr_ok=thr; break |
| if thr_ok is None: thr_ok=1.01 |
| a,p,fr=hyb(okt,ct,tok_test,test_all,thr_ok) |
| |
| a0,p0,fr0=hyb(okt,ct,tok_test,test_all,0.0) |
| R[f"K{K}"]=dict(params=sum(x.numel() for x in ad.parameters()),train_s=round(time.time()-ts), |
| adapter_acc_cov_test=sum(okt[m] for m in cov_test)/len(cov_test),adapter_acc_unc_test=sum(okt[m] for m in unc_test)/len(unc_test), |
| gate_threshold_from_cal=thr_ok,hybrid_test=dict(acc=a,pos_avg=p,native_frac=fr),coverage_only_test=dict(acc=a0,pos_avg=p0,native_frac=fr0)) |
| R[f"K{K}"]["_thr"]=thr_ok; print(R[f"K{K}"],flush=True) |
| json.dump(R,open("resultados-013b.json","w"),indent=1) |
| |
| random.seed(3); ts_=random.sample(cov_test,60)+random.sample(unc_test,40) |
| @torch.no_grad() |
| def one(m,route,K): |
| t0=time.perf_counter(); g=0.0 |
| if route=="hyb": |
| tg=time.perf_counter(); ad=ADS[K]; nat=False |
| if covered(m): |
| v=ad(torch.tensor([m])); r=readout(v) |
| c=min(float(F.softmax(x,1).max()) for x in r); nat=c>=R[f"K{K}"]["_thr"] |
| g=time.perf_counter()-tg |
| if nat: return time.perf_counter()-t0,g,g |
| tf=time.perf_counter(); e=E(ids(" "+FORMATS[FMT](m)))[None] |
| x=torch.cat([E(PRE)[None],e],1); out=lm(inputs_embeds=x,use_cache=True) |
| ttft=None |
| for f in range(4): |
| c=copy.deepcopy(out.past_key_values); lm(inputs_embeds=E(Q[f])[None],past_key_values=c) |
| if ttft is None: ttft=time.perf_counter()-t0 |
| return ttft,time.perf_counter()-t0,g |
| for m in ts_[:10]: one(m,"text",1); one(m,"hyb",1) |
| TIM={} |
| for name,route,K in (("texto",'text',1),("hibrido_K1",'hyb',1),("hibrido_K2",'hyb',2)): |
| tt=[];tl=[];tg=[];c0=os.times() |
| for rep in range(3): |
| for m in ts_: a,b,g=one(m,route,K); tt.append(a*1e3);tl.append(b*1e3);tg.append(g*1e3) |
| c1=os.times(); q=lambda v:sorted(v)[int(0.9*len(v))] |
| TIM[name]=dict(ttft_med_ms=st.median(tt),ttft_p90_ms=q(tt),total_med_ms=st.median(tl),total_p90_ms=q(tl),total_mean_ms=st.mean(tl),gate_adapter_mean_ms=st.mean(tg),cpu_s_per_msg=((c1.user-c0.user)+(c1.system-c0.system))/len(tl)) |
| print(name,TIM[name],flush=True) |
| R["timing_cpu"]=TIM |
| import resource; R["peak_rss_MB"]=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024 |
| json.dump(R,open("resultados-013b.json","w"),indent=1); print("DONE",flush=True) |
|
|
