Buckets:
| """E5 (extension): thinking in tokens, written by diffusion. Same S1 backbone + decision head; T note slots after | |
| the answer masks hold real text (the E3 notes, S1-tokenised, padded with <|endoftext|>; never cut: T >= the | |
| longest kept note). | |
| Train (one pass per item): items with a note get a random mask ratio t ~ U(0,1) over the note slots; | |
| loss = decision loss + MDLM CE on the masked note tokens. Items without a note keep all slots masked (decision | |
| loss only) -- the same start state as inference. | |
| Inference: K unmasking steps (most confident first, T/K slots per step) write the notes; a final pass decides. | |
| Passes per decision = K + 1 (compare with E3 at the same number of passes). | |
| """ | |
| import argparse | |
| import json | |
| import math | |
| import os | |
| import random | |
| import sys | |
| import time | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| from common import DATA, PROC, S1_PATH, Renderer, load_canonical, load_core, tokenizer_and_mapping # noqa: E402 | |
| from model import ThinkDecide, collate, item_logits, load_backbone # noqa: E402 | |
| from train import micro_batches # noqa: E402 | |
| from shared_yesno_supervised import supervised_terms # noqa: E402 | |
| PAD_NOTE = 151643 # <|endoftext|> | |
| def note_ids(tok, note, T): | |
| ids = tok.encode(note, add_special_tokens=False) | |
| assert len(ids) <= T, "note longer than T (%d > %d): raise T, never cut" % (len(ids), T) | |
| return ids + [PAD_NOTE] * (T - len(ids)) | |
| def forward_h(model, batch): | |
| """One pass; returns pre-norm hidden [B,L,d].""" | |
| ids, attn = batch["ids"], batch["attn"] | |
| x = model.bb.embed_tokens(ids) | |
| pos = torch.arange(ids.shape[1], device=ids.device)[None] | |
| return model.layers(x, attn[:, None, None, :], model.bb.rotary_emb(x, pos)) | |
| def note_logits(model, h_slots): | |
| return model.base.lm_head(model.bb.norm(h_slots)).float() | |
| def train(a): | |
| device = torch.device("cuda") | |
| torch.manual_seed(a.seed) | |
| tok, mp, mid, mt = tokenizer_and_mapping() | |
| rend = Renderer(tok, mid, mt, a.T) | |
| items, _ = load_canonical(a.train_file or DATA + "/train.jsonl", rend, a.max_length + a.T, a.train_limit) | |
| notes, too_long = {}, 0 | |
| for l in open(a.notes): | |
| n = json.loads(l) | |
| if n["reject"] is None and (not a.kept or (n["case_id"], n["view_id"]) in a.kept_set): | |
| if len(tok.encode(n["note"], add_special_tokens=False)) > a.T: | |
| too_long += 1 # left out of E5 training (never cut) | |
| continue | |
| notes[(n["case_id"], n["view_id"])] = note_ids(tok, n["note"], a.T) | |
| print(json.dumps(dict(event="notes", used=len(notes), left_out_longer_than_T=too_long)), flush=True) | |
| for x in items: | |
| x["note"] = notes.get((x["case_id"], x["view_id"])) | |
| print(json.dumps(dict(event="start", train=len(items), with_note=sum(x["note"] is not None for x in items), T=a.T)), | |
| flush=True) | |
| base = load_backbone(S1_PATH, torch.float32, device) | |
| model = ThinkDecide(base, mp["Yes"], mp["No"], n_think=a.T, use_head=True).to(device).train() | |
| for p in (model.gate, model.fb_proj.weight, model.fb_norm.weight, model.fb_norm.bias): | |
| p.requires_grad_(False) # no latent loop in E5 | |
| model.compile_layers() | |
| opt = torch.optim.AdamW(model.param_groups(a.lr_bb, a.lr_new, 0.01), fused=True) | |
| total = len(items) // 32 | |
| warm = math.ceil(0.05 * total) | |
| peaks = [g["lr"] for g in opt.param_groups] | |
| rng, order = random.Random(a.seed), list(range(len(items))) | |
| rng.shuffle(order) | |
| log = open(os.path.join(a.out, "train_log.jsonl"), "a") | |
| t0 = time.time() | |
| for step in range(total): | |
| prog = max(0.0, (step - warm) / max(1, total - warm)) | |
| for g, pk in zip(opt.param_groups, peaks): | |
| f = (step + 1) / warm if step < warm else 1e-6 / pk + (1 - 1e-6 / pk) * 0.5 * (1 + math.cos(math.pi * prog)) | |
| g["lr"] = pk * f | |
| batch_items = [] | |
| for i in order[step * 32:(step + 1) * 32]: | |
| x = dict(items[i]) | |
| ids = list(x["ids"]) | |
| masked = [] | |
| if x["note"] is not None: | |
| t = rng.random() | |
| for j, (p, tid) in enumerate(zip(x["think"], x["note"])): | |
| if rng.random() < t: | |
| masked.append((j, tid)) | |
| else: | |
| ids[p] = tid | |
| x["ids"], x["masked"] = ids, masked | |
| batch_items.append(x) | |
| sums = dict(dec=0.0, mdlm=0.0, n_tok=0) | |
| for mb in micro_batches(batch_items, 1, a.budget, 32): | |
| b = collate(mb, device, tok.pad_token_id) | |
| with torch.autocast("cuda", dtype=torch.bfloat16): | |
| h = forward_h(model, b) | |
| s = model.decide(h, b) | |
| dec = sum(supervised_terms(lg, torch.tensor([x["target"]], device=device), x["primitive"], 0.25)["loss"] | |
| for lg, x in zip(item_logits(s, b), mb)) | |
| rows = [(bi, x["think"][j], tid) for bi, x in enumerate(mb) for j, tid in x["masked"]] | |
| if rows: | |
| hb = h[[r[0] for r in rows], [r[1] for r in rows]] | |
| tgt = torch.tensor([r[2] for r in rows], device=device) | |
| # per-item mean over its masked tokens, summed over items | |
| per = torch.zeros(len(mb), device=device) | |
| cnt = torch.zeros(len(mb), device=device) | |
| lossv = F.cross_entropy(note_logits(model, hb), tgt, reduction="none") | |
| idx = torch.tensor([r[0] for r in rows], device=device) | |
| per.index_add_(0, idx, lossv) | |
| cnt.index_add_(0, idx, torch.ones_like(lossv)) | |
| mdlm = (per / cnt.clamp_min(1)).sum() | |
| else: | |
| mdlm = h.new_zeros(()) | |
| loss = (dec + a.mdlm_weight * mdlm) / 32 | |
| loss.backward() | |
| sums["dec"] += float(dec) | |
| sums["mdlm"] += float(mdlm) | |
| sums["n_tok"] += len(rows) | |
| gn = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) | |
| opt.step() | |
| opt.zero_grad(set_to_none=True) | |
| if (step + 1) % 10 == 0: | |
| rec = dict(step=step + 1, total=total, dec=round(sums["dec"] / 32, 5), mdlm=round(sums["mdlm"] / 32, 5), | |
| masked_tokens=sums["n_tok"], grad_norm=round(float(gn), 3), elapsed=round(time.time() - t0, 1), | |
| peak_gib=round(torch.cuda.max_memory_allocated() / 2**30, 2)) | |
| log.write(json.dumps(rec) + "\n") | |
| log.flush() | |
| print(json.dumps(rec), flush=True) | |
| from safetensors.torch import save_file | |
| sd = {k: v.detach().to(torch.bfloat16).contiguous() for k, v in model.state_dict().items()} | |
| sd.pop("base.lm_head.weight", None) | |
| sd = {k: v for k, v in sd.items() if not k.startswith("bb.")} | |
| save_file(sd, os.path.join(a.out, "model.safetensors")) | |
| def write_and_decide(model, batch, think, K, T, mask_id): | |
| """K unmasking steps over the note slots (most confident first), then a final pass decides. | |
| Returns option scores after the final pass and the written note ids [B,T].""" | |
| ids = batch["ids"].clone() | |
| B = ids.shape[0] | |
| tidx = batch["think_idx"] | |
| b = torch.arange(B, device=ids.device)[:, None] | |
| written = torch.full((B, T), -1, dtype=torch.long, device=ids.device) | |
| per_step = [math.ceil(T * (k + 1) / K) - math.ceil(T * k / K) for k in range(K)] | |
| for k in range(K): | |
| h = forward_h(model, dict(batch, ids=ids)) | |
| lg = note_logits(model, h[b, tidx]) # [B,T,V] | |
| lg[..., mask_id] = -1e9 | |
| conf, tokv = lg.softmax(-1).max(-1) | |
| conf = conf.masked_fill(written >= 0, -1.0) | |
| n = per_step[k] | |
| top = conf.topk(n, dim=-1).indices # [B,n] | |
| written.scatter_(1, top, tokv.gather(1, top)) | |
| ids[b.expand(-1, n), tidx.gather(1, top)] = tokv.gather(1, top) | |
| h = forward_h(model, dict(batch, ids=ids)) | |
| return model.decide(h, dict(batch, ids=ids)), written | |
| def evaluate(a): | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| from eval import build_model, probs_from_scores, summarize | |
| device = torch.device("cuda") | |
| tok, mp, mid, mt = tokenizer_and_mapping() | |
| rend = Renderer(tok, mid, mt, a.T) | |
| if a.split == "core": | |
| items, _ = load_core(rend, 16384) | |
| else: | |
| path = {"dev": DATA + "/dev.jsonl", "proc_test": PROC + "/test.jsonl", "proc_dev": PROC + "/dev.jsonl"}[a.split] | |
| items, _ = load_canonical(path, rend, 10**9) | |
| for x in items: | |
| x.update(id="%s|%s" % (x["case_id"], x["view_id"]), suite=x["source"], | |
| gold_idx=int(np.argmax(x["target"])) if x["kind"] == "hard" else None) | |
| oracle = {} | |
| if a.oracle_notes: # diagnostic: decide with given (gold-informed) notes in the slots instead of writing them | |
| for l in open(a.oracle_notes): | |
| n = json.loads(l) | |
| ids = tok.encode(n["note"], add_special_tokens=False) | |
| if len(ids) <= a.T and n["reject"] is None: | |
| oracle["%s|%s" % (n["case_id"], n["view_id"])] = ids + [PAD_NOTE] * (a.T - len(ids)) | |
| items = [x for x in items if x["id"] in oracle] | |
| for x in items: | |
| x["ids"] = list(x["ids"]) | |
| for p, t in zip(x["think"], oracle[x["id"]]): | |
| x["ids"][p] = t | |
| print(json.dumps(dict(event="oracle", items=len(items))), flush=True) | |
| model = build_model(a.model, a.T, True, device, mp) | |
| Ks = [int(k) for k in a.ks.split(",")] | |
| rows = [dict(id=x["id"], suite=x["suite"], primitive=x["primitive"], gold_idx=x["gold_idx"], level=x.get("level"), | |
| probs={}, notes={}) for x in items] | |
| if a.oracle_notes: | |
| Ks = [0] | |
| order = sorted(range(len(items)), key=lambda i: len(items[i]["ids"])) | |
| for K in Ks: | |
| for s in range(0, len(order), 32): | |
| idx = order[s:s + 32] | |
| mb = [items[i] for i in idx] | |
| batch = collate(mb, device, tok.pad_token_id) | |
| if K == 0: # oracle: notes already in the canvas; one pass decides | |
| sc = model.decide(forward_h(model, batch), batch) | |
| written = torch.full((len(mb), a.T), -1, dtype=torch.long) | |
| else: | |
| sc, written = write_and_decide(model, batch, None, K, a.T, mid) | |
| sc = sc.float().cpu().numpy() | |
| for j, i in enumerate(idx): | |
| x = items[i] | |
| rows[i]["probs"][str(K + 1)] = [float(v) for v in probs_from_scores(sc[j], x["primitive"], len(x["pos"]))] | |
| if j < 2 or a.save_notes: | |
| w = [t for t in written[j].tolist() if t not in (PAD_NOTE, -1)] | |
| rows[i]["notes"][str(K)] = tok.decode(w) | |
| os.makedirs(a.out, exist_ok=True) | |
| with open(os.path.join(a.out, "predictions.jsonl"), "w") as f: | |
| for r in rows: | |
| f.write(json.dumps(r) + "\n") | |
| hard = [r for r in rows if r["gold_idx"] is not None] | |
| rep = dict(model=a.model, T=a.T, passes={str(K + 1): K for K in Ks}, | |
| by_K={str(K + 1): summarize(hard, str(K + 1)) for K in Ks}) | |
| json.dump(rep, open(os.path.join(a.out, "report.json"), "w"), indent=1) | |
| for K in Ks: | |
| key = str(K + 1) | |
| lv = {l: [r for r in hard if r["level"] == l] for l in range(5)} | |
| accf = lambda rs: float(np.mean([np.argmax(r["probs"][key]) == r["gold_idx"] for r in rs])) if rs else float("nan") | |
| rep.setdefault("proc", {})[key] = dict(all=accf(hard), by_level={str(l): accf(v) for l, v in lv.items()}) | |
| print("unmask K=%d (passes %d) all %.2f | " % (K, K + 1, 100 * accf(hard)) + | |
| " ".join("L%d %.2f" % (l, 100 * accf(v)) for l, v in lv.items())) | |
| json.dump(rep, open(os.path.join(a.out, "report.json"), "w"), indent=1) | |
| if __name__ == "__main__": | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("mode", choices=["train", "eval"]) | |
| ap.add_argument("--out", required=True) | |
| ap.add_argument("--notes", default="/workspace/think-decide/runs/e3/notes.jsonl") | |
| ap.add_argument("--kept", default="", help="teacher_cache report keys: train only on notes that helped the teacher") | |
| ap.add_argument("--T", type=int, default=128) | |
| ap.add_argument("--lr-bb", type=float, default=1e-5) | |
| ap.add_argument("--lr-new", type=float, default=1e-4) | |
| ap.add_argument("--mdlm-weight", type=float, default=1.0) | |
| ap.add_argument("--budget", type=int, default=40000) | |
| ap.add_argument("--seed", type=int, default=42) | |
| ap.add_argument("--model", default="") | |
| ap.add_argument("--split", default="core") | |
| ap.add_argument("--ks", default="3,7") | |
| ap.add_argument("--save-notes", action="store_true") | |
| ap.add_argument("--train-file", default="") | |
| ap.add_argument("--train-limit", type=int, default=0) | |
| ap.add_argument("--max-length", type=int, default=2048) | |
| ap.add_argument("--oracle-notes", default="") | |
| a = ap.parse_args() | |
| a.kept_set = set() | |
| if a.kept: | |
| a.kept_set = {tuple(k.split("|", 1)) for k in json.load(open(a.kept))} | |
| os.makedirs(a.out, exist_ok=True) | |
| train(a) if a.mode == "train" else evaluate(a) | |
Xet Storage Details
- Size:
- 13.2 kB
- Xet hash:
- f6e75bfcd0f417d2a7453d44cd2ff776c1c393dc1c6b70bc49fd028c10d22801
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.