Revision 5: negative-result model card, report rev. 5 (PDF/MD/figure), CITATION.cff, iterations 020-021 code, logs, sealed tests and checkpoints
7420c83 verified | # CUERDAS-021 common: sealed set, varied replacement pool, training loop, evaluation. | |
| import os, sys, json, hashlib, random, re, time | |
| import torch, torch.nn as nn, torch.nn.functional as F | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| import run018 as R | |
| INP = os.environ['INP']; OD = os.environ['OD'] | |
| raw = open(f'{INP}/sealed021.json').read(); H = hashlib.sha256(raw.encode()).hexdigest() | |
| assert H == 'a86a6a5ee375ae67a3e79f969c9553b0e5548ed3b0e80e5315862e6552ae72bc', H | |
| S = json.loads(raw); NEW = S['new_values']; EXCL = set(S['exclude_token_ids']) | |
| SETS = ['sealed_id', 'sealed_ood', 'sealed_1_new', 'sealed_2_new', 'sealed_4_new'] | |
| SMOKE = bool(os.environ.get('SMOKE')) | |
| if SMOKE: # dry run on dev messages only; sealed sets are never touched | |
| for s_ in SETS: S[s_] = R.FRESH['ood'][:6] | |
| STEPS = 4 if SMOKE else 3000; EVERY = 2 if SMOKE else 250 | |
| t1 = R.t1; tok = R.tok; Emb = R.Emb; D = R.D; FIELDS = R.FIELDS; V = R.V | |
| CAND = [torch.tensor(sorted({t1(w)[0] for w in V[j] + NEW[j]})) for j in range(4)] | |
| assert all(len(CAND[j]) == len(V[j]) + 8 for j in range(4)) | |
| def gidx(cand, j, w): return int((cand[j] == t1(w)[0]).nonzero()[0]) | |
| # ---- varied pool (training) and held-out validation pool; all disjoint from sealed tokens and from 014/dev/020 values ---- | |
| voc = tok.get_vocab(); BAN = EXCL | R.BANNED | |
| def ok_tok(t): return voc[t] not in BAN | |
| WORDS = sorted(t[1:] for t in voc if re.fullmatch(r'Ġ[a-z]{3,14}', t) and ok_tok(t)) | |
| NAMES = sorted(t[1:] for t in voc if re.fullmatch(r'Ġ[A-Z][a-z]{2,10}', t) and ok_tok(t)) | |
| SUBS = sorted(t for t in voc if re.fullmatch(r'[a-z]{2,5}', t)) | |
| g0 = random.Random(2101); g0.shuffle(WORDS); g0.shuffle(NAMES) | |
| VW, VN = WORDS[-300:], NAMES[-300:]; WORDS, NAMES = WORDS[:-300], NAMES[:-300] | |
| def sample_value(g, W_=None, N_=None): | |
| W_ = W_ or WORDS; N_ = N_ or NAMES; k = g.randrange(5) | |
| if k == 0: return g.choice(W_) | |
| if k == 1: return g.choice(N_) | |
| if k == 2: return str(g.randrange(10000)) | |
| if k == 3: return ' '.join(g.choice(W_) for _ in range(g.choice((2, 3)))) | |
| return g.choice(W_) + g.choice(SUBS) | |
| def sample_msgs(g, bs=12): | |
| out = [] | |
| for _ in range(bs): | |
| m = list(g.choice(R.FIT)) | |
| for j in range(4): | |
| if g.random() < .4: m[j] = sample_value(g) | |
| out.append(m) | |
| return out | |
| # restricted transfer validation: 50 held-out values per field (no numbers: first token not unique), options = field vocab + those | |
| gv = random.Random(2102); CANDV = []; VPOOL = [] | |
| for j in range(4): | |
| used = {t1(w)[0] for w in V[j]}; ws = [] | |
| while len(ws) < 50: | |
| k = gv.randrange(3); w = gv.choice(VW) if k == 0 else gv.choice(VN) if k == 1 else gv.choice(VW) + ' ' + gv.choice(VW) | |
| if t1(w)[0] not in used: used.add(t1(w)[0]); ws.append(w) | |
| VPOOL.append(ws); CANDV.append(torch.tensor(sorted(used))) | |
| VAL_TR = [[gv.choice(VPOOL[j]) for j in range(4)] for _ in range(96)] | |
| VALC = R.VAL[:96] | |
| if SMOKE: VALC, VAL_TR = VALC[:6], VAL_TR[:6] | |
| # ---- generic readout: rows = prefix (embeddings) + body tokens ---- | |
| def run_h(embs): | |
| 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 = R.lm.model(inputs_embeds=x, attention_mask=m).last_hidden_state[torch.arange(B), m.sum(1) - 1] | |
| return h, R.lm.lm_head(h) | |
| def k2_rows(cod, msgs, tpl): | |
| Sl = cod(torch.stack([torch.stack([R.lex(w) for w in m]) for m in msgs])) | |
| return [torch.cat([Sl[i], Emb[R.sufids(tpl, j)]]) for i in range(len(msgs)) for j in range(4)] | |
| def text_rows(msgs, tpl, fmt, prefix=None): | |
| rows = [] | |
| for m in msgs: | |
| for j in range(4): | |
| e = Emb[tok.encode(fmt.format(**dict(zip(FIELDS, m))) + tpl.format(field=FIELDS[j]), add_special_tokens=False)] | |
| rows.append(e if prefix is None else torch.cat([prefix, e])) | |
| return rows | |
| def per_ex(rowfn, msgs, cand, bs=24): | |
| out = [] | |
| for b in range(0, len(msgs), bs): | |
| mb = msgs[b:b+bs]; _, lg = run_h(rowfn(mb)); good = [True] * len(mb) | |
| for r_ in range(len(mb) * 4): | |
| i, j = divmod(r_, 4); good[i] &= int(lg[r_][cand[j]].argmax()) == gidx(cand, j, mb[i][j]) | |
| out += good | |
| return [float(x) for x in out] | |
| def val_score(rowfn_main): | |
| return sum(per_ex(rowfn_main, VALC, R.CAND)) / len(VALC) + sum(per_ex(rowfn_main, VAL_TR, CANDV)) / len(VAL_TR) | |
| def train(params, rowfn, teach, seed, tag, log, steps=None, lr=5e-4, every=None, set_mode=lambda t: None): | |
| """Same loss/schedule/selection as 018 distill2: CE + KL(teacher) + 0.5(1-cos) ; OneCycle; select on validation only.""" | |
| steps = steps or STEPS; every = every or EVERY | |
| g = random.Random(seed + 2100); opt = torch.optim.AdamW(params, lr=lr, weight_decay=1e-3) | |
| sch = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=lr, total_steps=steps, pct_start=.05) | |
| def snap(): return [p.detach().clone() for p in params] | |
| set_mode(False); v = val_score(lambda mb: rowfn(mb, R.SUF_MAIN)); set_mode(True); best = (v, snap(), 0); log({'run': tag, 'step': 0, 'val': v}) | |
| t0 = time.time() | |
| for s_ in range(1, steps + 1): | |
| msgs = sample_msgs(g); tpl = g.choice(R.SUF_TRAIN) | |
| with torch.no_grad(): ht, lt = R.readout_h(teach, msgs, tpl) | |
| hs, ls = run_h(rowfn(msgs, tpl)); tg = torch.tensor([t1(m[j])[0] for m in msgs for j in range(4)]) | |
| loss = F.cross_entropy(ls, tg) + F.kl_div(F.log_softmax(ls, -1), F.log_softmax(lt, -1), log_target=True, reduction='batchmean') + .5 * (1 - F.cosine_similarity(hs, ht, -1).mean()) | |
| opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(params, 1.0); opt.step(); sch.step() | |
| if s_ % every == 0: | |
| set_mode(False); v = val_score(lambda mb: rowfn(mb, R.SUF_MAIN)); set_mode(True) | |
| log({'run': tag, 'step': s_, 'val': v, 'loss': float(loss), 's': round(time.time() - t0)}) | |
| if v > best[0]: best = (v, snap(), s_) | |
| with torch.no_grad(): | |
| for p, b in zip(params, best[1]): p.copy_(b) | |
| set_mode(False); return best[2], best[0] | |
| def load_teacher(sd): | |
| t = R.Binder(sd); t.load_state_dict(torch.load(f'{INP}/binder4_1000_s{sd}.pt', map_location=R.DEV)); return t.eval().requires_grad_(False) | |
| def opener(name): | |
| f = open(f'{OD}/{name}', 'w') | |
| def log(d): f.write(json.dumps(d) + '\n'); f.flush(); print(json.dumps(d), flush=True) | |
| return log | |