| |
| 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 |
| 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 |
| @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 |
|
|
