File size: 5,012 Bytes
5851d4e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | # Cuerdas 013b: LLM real congelado, splits disjuntos, sweep de formato, destilación, gate calibrado. CPU.
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] # holdout sistemático: 40x28 = 1120 cubiertas
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): # prefill compartido (KV cache) + 4 preguntas -> lista de logits [B,ncand]
B=msg_emb.shape[0]
if torch.is_grad_enabled(): # entrenamiento: sin cache (deepcopy no admite grafo); misma función
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
|