| |
| import sys; exec(open('iterate.py').read().split("CFG=")[0].replace('FMT="natural_labeled"','FMT="natural_labeled"').split("_,TL=eval_text")[0]) |
| FORMATS["rival_like"]=lambda m:f"The action is {A[m[0]]}. The target is {T[m[1]]}. The state is {S[m[2]]}. The window is {W[m[3]]}." |
| FORMATS["natural_labeled"]=FORMATS["natural_labeled"] |
| QS={"qa":[ids(f"\nQ: what is the {f}?\nA:") for f in FN],"cloze":[ids(f"\nThe {f} is") for f in FN],"cloze_colon":[ids(f"\n{f.capitalize()}:") for f in FN]} |
| random.seed(2); calS=random.sample(cov_cal,100)+random.sample(unc_cal,100) |
| sweep=[] |
| @torch.no_grad() |
| def perfield(fmt,ms): |
| pf=[0]*4; ok=0 |
| for i in range(0,len(ms),50): |
| b=ms[i:i+50]; L={} |
| for m in b: L.setdefault(len(ids(" "+FORMATS[fmt](m))),[]).append(m) |
| for sub in L.values(): |
| r=readout(E(torch.stack([ids(" "+FORMATS[fmt](m)) for m in sub]))); t=torch.tensor(sub); o=torch.ones(len(sub),dtype=torch.bool) |
| for f in range(4): c=r[f].argmax(1)==t[:,f]; pf[f]+=int(c.sum()); o&=c |
| ok+=int(o.sum()) |
| return ok/len(ms),[x/len(ms) for x in pf] |
| DEMO=(5,6,7,4) |
| PRES={"none":ids("Message:")} |
| for fmt in ["key_value","natural_labeled"]: |
| d=FORMATS[fmt](DEMO); PRES["1shot_"+fmt]=ids("Message: "+d+"".join(f"\nQ: what is the {f}?\nA: {FL[i][DEMO[i]]}" for i,f in enumerate(FN))+"\n\nMessage:") |
| Q[:]=QS["qa"] |
| for pn,pv in PRES.items(): |
| PRE=pv; qn="qa" |
| for fmt in ["key_value","natural_labeled","json"]: |
| if pn.startswith("1shot") and not pn.endswith(fmt) : continue |
| a,pf=perfield(fmt,calS); r=dict(pre=pn,query=qn,fmt=fmt,cal_acc=a,per_field=dict(zip(FN,pf)),avg_tokens=sum(len(ids(" "+FORMATS[fmt](m))) for m in calS)/len(calS)) |
| sweep.append(r); print(json.dumps(r),flush=True) |
| json.dump(sweep,open("sweep_013h.json","w"),indent=1) |
| good=[r for r in sweep if r["cal_acc"]>=0.97] |
| best=min(good,key=lambda r:r["avg_tokens"]) if good else max(sweep,key=lambda r:r["cal_acc"]) |
| print("BEST",best,flush=True); PRE=PRES[best["pre"]]; Q[:]=QS[best["query"]]; FMT=best["fmt"]; pos=lambda m:len(ids(" "+FORMATS[FMT](m))) |
| |
| okT_cal,_=eval_text(FMT,cal_all); okT_dev,_=eval_text(FMT,dev) |
| errs={} |
| for m,o in okT_cal.items(): |
| if not o: |
| for f in range(4): errs[f"{FN[f]}={FL[f][m[f]]}"]=errs.get(f"{FN[f]}={FL[f][m[f]]}",0)+1 |
| json.dump(dict(best=best,cal_full_acc=sum(okT_cal.values())/len(cal_all),dev_acc=sum(okT_dev.values())/len(dev),values_in_failed_msgs=errs),open("baseline_013h.json","w"),indent=1) |
| _,TL=eval_text(FMT,cov_train,keep_logits=True) |
| exec(open("defs_iter.py").read()) |
| RUNS=[("013h","fact4",200),("013i","add1",200),("013j","mlp4",160),("013k","mlp4",640),("013l","mlp1",640)] |
| ADS={} |
| for name,mode,steps in RUNS: |
| torch.manual_seed(0); random.seed(0) |
| ad=Fact(mode) if mode in("fact4","add1") else MLP(int(mode[-1])) |
| t0=time.time(); ad=train(ad,steps,True); ADS[name]=ad |
| r=evaluate(name,ad,cal_all,dev,okT_cal,okT_dev); r.update(mode=mode,steps=steps,train_s=round(time.time()-t0),params=sum(p.numel() for p in ad.parameters()),teacher=f'{FMT}/{best["pre"]}') |
| open("iteraciones.jsonl","a").write(json.dumps(r)+"\n"); torch.save(ad.state_dict(),f"adapter_{name}.pt"); print(name,json.dumps(r),flush=True) |
| |
| cands=[json.loads(l) for l in open("iteraciones.jsonl") if l.startswith('{"iter": "013') and json.loads(l).get("teacher")] |
| win=max(cands,key=lambda r:(r["promising_native"],r["native_acc_sysholdout"],-r["K"]))["iter"] |
| okT_f,_=eval_text(FMT,fresh); rf=evaluate(win+"_FRESH",ADS[win],cal_all,fresh,okT_cal,okT_f) |
| open("iteraciones.jsonl","a").write(json.dumps(rf)+"\n"); print("FRESH",json.dumps(rf),flush=True) |
| |
| ad=ADS[win]; AB={} |
| with torch.no_grad(): |
| t=torch.tensor(fresh); v=ad(t) |
| if v.shape[1]==4: |
| okn=torch.ones(len(fresh),dtype=torch.bool) |
| for f in range(4): okn&=(v[:,f]@E.weight[cand[f]].T).argmax(1)==t[:,f] |
| AB["nn_decode_no_llm_acc"]=float(okn.float().mean()) |
| AB["holdout_pairs_in_train"]=sum(1 for m in cov_train if (m[0],m[1]) in set(H_AT) or (m[2],m[3]) in set(H_SW)) |
| AB["demo_msg_in_test"]=DEMO in set(fresh+dev) |
| from transformers import AutoConfig |
| real=lm; lm=AutoModelForCausalLM.from_config(AutoConfig.from_pretrained(M)).eval(); E_real=E |
| okr,_=eval_ad(ad,fresh); AB["random_llm_native_acc"]=sum(okr.values())/len(fresh) |
| okrt,_=eval_text(FMT,fresh); AB["random_llm_text_acc"]=sum(okrt.values())/len(fresh) |
| lm=real; json.dump(AB,open("ablaciones_013.json","w"),indent=1); print("ABL",AB,flush=True) |
| |
| ad=ADS[win]; random.seed(3); ts_=random.sample([m for m in fresh if covered(m)],57)+random.sample([m for m in fresh if not covered(m)],43) |
| @torch.no_grad() |
| def one(m,route): |
| t0=time.perf_counter(); g=0.0 |
| if route=="text" or (route=="hyb" and not covered(m)): |
| if route=="hyb": g=time.perf_counter()-t0 |
| e=E(ids(" "+FORMATS[FMT](m)))[None] |
| else: |
| e=ad(torch.tensor([m])); g=time.perf_counter()-t0 |
| out=lm(inputs_embeds=torch.cat([E(PRE)[None],e],1),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*1e3,(time.perf_counter()-t0)*1e3,g*1e3 |
| for m in ts_[:10]: one(m,"text"); one(m,"native") |
| TIM={}; q=lambda v:sorted(v)[int(0.9*len(v))] |
| data={k:[] for k in ("text","native","hyb")} |
| for rep in range(3): |
| for m in ts_: |
| for k in random.sample(list(data),3): data[k].append(one(m,k)) |
| for k,v in data.items(): |
| tt=[x[0] for x in v]; tl=[x[1] for x in v]; tg=[x[2] for x in v] |
| TIM[k]=dict(ttft_med_ms=st.median(tt),ttft_p90_ms=q(tt),total_med_ms=st.median(tl),total_p90_ms=q(tl),gate_adapter_mean_ms=st.mean(tg),n=len(v)) |
| import resource |
| json.dump(dict(winner=win,timing_cpu=TIM,positions=dict(text=sum(map(pos,ts_))/len(ts_),native=ad.K,hyb=sum((ad.K if covered(m) else pos(m)) for m in ts_)/len(ts_)),peak_rss_MB=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024),open("timing_013.json","w"),indent=1) |
| print("TIM",json.dumps(TIM),"DONE",flush=True) |
|
|
