Buckets:
| """MQ: answer diffusion over all questions of a procedural state ("thinking" = committing supervised sub-answers). | |
| Canvas: S1's chat format with every question of the state, each in S1's own block style (option lines ending in a | |
| mask; noul -> "Answer:" mask). A question is *revealed* by writing gold (train) or predicted (inference) answer | |
| tokens into its masks: Yes at the chosen option, No elsewhere (noul: Yes/No). Readout per question = S1's native | |
| Yes/No readout (no new parameters). | |
| Train (masked diffusion over answers): per state t ~ U(0,1), each question revealed with prob 1-t (>= 1 stays | |
| masked); loss = S1's decision loss on the masked questions (mean per state). | |
| --soft adds two-pass self-conditioning (p = 0.8): a no-grad pass gives beliefs; each still-masked answer slot's | |
| input becomes the expected answer embedding sum_y p(y) e_y (y in {Yes, No}); loss on the second pass. | |
| Eval modes (same model, same questions): parallel (all masked, 1 pass), sequential (commit the most confident | |
| question per pass, hard tokens), soft (K passes of soft-belief feedback, no commits; --soft models), oracle (all | |
| other questions revealed with gold). | |
| """ | |
| import argparse | |
| import glob | |
| import json | |
| import math | |
| import os | |
| import random | |
| import sys | |
| import time | |
| from collections import Counter, defaultdict | |
| import numpy as np | |
| import pandas as pd | |
| import torch | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| from common import PROC, S1_PATH, tokenizer_and_mapping # noqa: E402 | |
| from model import ThinkDecide, load_backbone # noqa: E402 | |
| from prep_procedural import parse_state, to_row # noqa: E402 | |
| from train import micro_batches # noqa: E402 | |
| from bench_local import to_internal # noqa: E402 | |
| from laya.common import render_options, serialize_state # noqa: E402 | |
| from shared_yesno_supervised import supervised_terms # noqa: E402 | |
| RAW = "/workspace/data/procedural/all" | |
| def render(tok, state, qlist, mask_id, mask_text): | |
| """qlist: [(qid, qdef)] -> ids, per-question mask positions, options.""" | |
| user = "State:\n%s\n\n" % serialize_state(state) | |
| counts, blocks = [], [] | |
| for i, (qid, q) in enumerate(qlist, 1): | |
| qi = to_internal(q) | |
| opts = render_options(qi) | |
| b = "Question %d:\n%s\n\n" % (i, qi["ins"]) | |
| if qi["t"] == "noul": | |
| b += "Answer Yes or No.\nAnswer:" + mask_text | |
| counts.append(1) | |
| else: | |
| b += "For each option, mark Yes if it answers the question, otherwise No.\n" | |
| b += "\n".join(o + ": " + mask_text for o in opts) | |
| counts.append(len(opts)) | |
| blocks.append(b) | |
| user += "\n\n".join(blocks) | |
| prompt = tok.apply_chat_template([{"role": "system", "content": "You are a helpful assistant."}, | |
| {"role": "user", "content": user}], tokenize=False, | |
| add_generation_prompt=False, enable_thinking=False) | |
| parts = prompt.split(mask_text) | |
| assert len(parts) == sum(counts) + 1 | |
| ids, pos = [], [] | |
| for i, part in enumerate(parts): | |
| ids.extend(tok.encode(part, add_special_tokens=False)) | |
| if i < len(parts) - 1: | |
| pos.append(len(ids)) | |
| ids.append(mask_id) | |
| out, k = [], 0 | |
| for c in counts: | |
| out.append(pos[k:k + c]) | |
| k += c | |
| return ids, out | |
| def load_states(split, state_ids=None, limit=0, shuffle_questions=False, seed=0): | |
| files = sorted(glob.glob("%s/%s-*.parquet" % (RAW, split))) | |
| df = pd.concat([pd.read_parquet(f) for f in files], ignore_index=True) | |
| if state_ids is not None: | |
| have = set(df["id"].astype(str)) | |
| keep = [s for s in state_ids if s in have] | |
| df = df.set_index(df["id"].astype(str)).loc[keep].reset_index(drop=True) | |
| rng = random.Random(seed) | |
| states = [] | |
| for _, r in df.iterrows(): | |
| qs, ans = json.loads(r["questions"]), json.loads(r["answers"]) | |
| qlist = [] | |
| for qid in qs: | |
| row = to_row(r, qid, qs[qid], ans[qid], "x") | |
| if row is None: | |
| continue | |
| probs = row["gold"][qid]["probabilities"] | |
| keys = row["option_ids"][qid] | |
| qd = row["questions"][qid] | |
| qlist.append(dict(qid=qid, q=qd, prim=qd["type"], target=[probs[k] for k in keys])) | |
| if not qlist: | |
| continue | |
| if shuffle_questions: | |
| rng.shuffle(qlist) | |
| states.append(dict(id=str(r["id"]), task=r["task"], level=int(r["level"]), state=parse_state(r["state"]), | |
| qs=qlist)) | |
| if limit and len(states) >= limit: | |
| break | |
| return states | |
| def prepare(states, tok, mask_id, mask_text, max_len): | |
| out, dropped = [], 0 | |
| for s in states: | |
| ids, pos = render(tok, s["state"], [(x["qid"], x["q"]) for x in s["qs"]], mask_id, mask_text) | |
| if len(ids) > max_len: | |
| dropped += 1 | |
| continue | |
| for x, p in zip(s["qs"], pos): | |
| x["pos"] = p | |
| n = 1 if x["prim"] == "noul" else len(x["target"]) | |
| assert len(p) == n, (len(p), n, x["prim"]) | |
| out.append(dict(s, ids=ids)) | |
| return out, dropped | |
| def answer_tokens(x, choice_idx, yes, no): | |
| """Tokens to write into question x's masks for answer index choice_idx.""" | |
| if x["prim"] == "noul": | |
| return [yes if choice_idx == 1 else no] | |
| return [yes if j == choice_idx else no for j in range(len(x["pos"]))] | |
| class MQModel: | |
| def __init__(self, model, yes, no, mask_id): | |
| self.m, self.yes, self.no, self.mask_id = model, yes, no, mask_id | |
| E = model.bb.embed_tokens.weight | |
| self.E = E | |
| def forward(self, ids_list, soft=None): | |
| """ids_list: list of id lists (same batch); soft: list of {pos: p_yes} for soft answer slots. -> pre-norm h.""" | |
| B, L = len(ids_list), max(map(len, ids_list)) | |
| dev = self.E.device | |
| ids = torch.full((B, L), 151643, dtype=torch.long) | |
| attn = torch.zeros((B, L), dtype=torch.bool) | |
| for i, x in enumerate(ids_list): | |
| ids[i, :len(x)] = torch.tensor(x) | |
| attn[i, :len(x)] = True | |
| ids, attn = ids.to(dev), attn.to(dev) | |
| x = self.m.bb.embed_tokens(ids) | |
| if soft is not None: | |
| rows, cols, py = [], [], [] | |
| for i, d in enumerate(soft): | |
| for p, v in d.items(): | |
| rows.append(i) | |
| cols.append(p) | |
| py.append(v) | |
| if rows: | |
| py = torch.tensor(py, device=dev, dtype=x.dtype)[:, None] | |
| emb = py * self.E[self.yes].to(x.dtype) + (1 - py) * self.E[self.no].to(x.dtype) | |
| x = x.index_put((torch.tensor(rows, device=dev), torch.tensor(cols, device=dev)), emb) | |
| pos = torch.arange(L, device=dev)[None] | |
| return self.m.layers(x, attn[:, None, None, :], self.m.bb.rotary_emb(x, pos)) | |
| def scores(self, h, items_pos): | |
| """h [B,L,d]; items_pos: list (per batch row) of list of mask-position lists -> list of list of score tensors.""" | |
| rows, cols = [], [] | |
| for i, qs in enumerate(items_pos): | |
| for p in qs: | |
| rows += [i] * len(p) | |
| cols += p | |
| s = self.m.readout(h[torch.tensor(rows, device=h.device), torch.tensor(cols, device=h.device)]) | |
| out, k = [], 0 | |
| for qs in items_pos: | |
| per = [] | |
| for p in qs: | |
| per.append(s[k:k + len(p)]) | |
| k += len(p) | |
| out.append(per) | |
| return out | |
| def q_logits(sc, prim): | |
| return torch.stack([sc.new_zeros(()), sc[0]])[None] if prim == "noul" else sc[None] | |
| def q_probs(sc, prim): | |
| return torch.softmax(q_logits(sc, prim).float(), -1)[0] | |
| def beliefs(x, pr): | |
| """Per-mask p(Yes) from a question's option distribution (noul: p(true)).""" | |
| if x["prim"] == "noul": | |
| return {x["pos"][0]: float(pr[1])} | |
| return {p: float(v) for p, v in zip(x["pos"], pr.tolist())} | |
| def train(a): | |
| device = torch.device("cuda") | |
| torch.manual_seed(a.seed) | |
| tok, mp, mid, mt = tokenizer_and_mapping() | |
| yes, no = mp["Yes"], mp["No"] | |
| ids = [json.loads(l)["group_id"] for _, l in zip(range(a.n_states), open(PROC + "/train.jsonl"))] | |
| if a.n_states > len(ids): # beyond the prepared 78.5k: add raw train states (seeded, never dev/test splits) | |
| raw = pd.concat([pd.read_parquet(f, columns=["id"]) for f in sorted(glob.glob(RAW + "/train-*.parquet"))]) | |
| have = set(ids) | |
| extra = [str(x) for x in raw["id"].astype(str) if str(x) not in have] | |
| random.Random(a.seed).shuffle(extra) | |
| ids += extra[:a.n_states - len(ids)] | |
| states = load_states("train", ids, shuffle_questions=True, seed=a.seed) | |
| states, dropped = prepare(states, tok, mid, mt, a.max_length) | |
| print(json.dumps(dict(event="start", states=len(states), dropped=dropped, | |
| questions=sum(len(s["qs"]) for s in states), soft=a.soft)), flush=True) | |
| base = load_backbone(S1_PATH, torch.float32, device) | |
| model = ThinkDecide(base, yes, no, n_think=0, use_head=False).to(device).train() | |
| model.compile_layers() | |
| mq = MQModel(model, yes, no, mid) | |
| opt = torch.optim.AdamW(model.param_groups(a.lr_bb, a.lr_bb, 0.01), fused=True) | |
| total = len(states) // 32 | |
| if a.max_steps: | |
| total = min(total, a.max_steps) | |
| warm = math.ceil(0.05 * total) | |
| peak = a.lr_bb | |
| rng = random.Random(a.seed) | |
| order = list(range(len(states))) | |
| 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)) | |
| f = (step + 1) / warm if step < warm else 1e-6 / peak + (1 - 1e-6 / peak) * 0.5 * (1 + math.cos(math.pi * prog)) | |
| for g in opt.param_groups: | |
| g["lr"] = peak * f | |
| batch = [] | |
| for i in order[step * 32:(step + 1) * 32]: | |
| s = states[i] | |
| t = rng.random() | |
| rev = [rng.random() >= t for _ in s["qs"]] | |
| if all(rev): | |
| rev[rng.randrange(len(rev))] = False | |
| ids_ = list(s["ids"]) | |
| for x, r in zip(s["qs"], rev): | |
| if r: | |
| for p, tk in zip(x["pos"], answer_tokens(x, int(np.argmax(x["target"])), yes, no)): | |
| ids_[p] = tk | |
| batch.append(dict(ids=ids_, s=s, rev=rev)) | |
| stats = Counter() | |
| for mb in micro_batches(batch, 2 if a.soft else 1, a.budget, 64, waste=a.waste, min_tokens=a.min_tokens): | |
| masked = [[x for x, r in zip(e["s"]["qs"], e["rev"]) if not r] for e in mb] | |
| soft = None | |
| with torch.autocast("cuda", dtype=torch.bfloat16): | |
| if a.soft and rng.random() < 0.8: | |
| model.eval() | |
| with torch.no_grad(): | |
| h0 = mq.forward([e["ids"] for e in mb]) | |
| sc0 = mq.scores(h0, [[x["pos"] for x in m] for m in masked]) | |
| model.train() | |
| soft = [{} for _ in mb] | |
| for i, (m, scs) in enumerate(zip(masked, sc0)): | |
| for x, sc in zip(m, scs): | |
| soft[i].update(beliefs(x, q_probs(sc, x["prim"]))) | |
| stats["soft"] += 1 | |
| h = mq.forward([e["ids"] for e in mb], soft) | |
| scs = mq.scores(h, [[x["pos"] for x in m] for m in masked]) | |
| loss = 0.0 | |
| for m, sc_l in zip(masked, scs): | |
| per = [supervised_terms(q_logits(sc, x["prim"]), torch.tensor([x["target"]], device=device), | |
| x["prim"], 0.25)["loss"] for x, sc in zip(m, sc_l)] | |
| loss = loss + sum(per) / len(per) | |
| stats["q"] += len(per) | |
| loss = loss / 32 | |
| loss.backward() | |
| stats["loss"] += float(loss) * 32 | |
| 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, loss=round(stats["loss"] / 32, 5), masked_q=stats["q"], | |
| soft_mb=stats["soft"], 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) | |
| if a.no_save: | |
| return | |
| 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")) | |
| print(json.dumps(dict(event="saved")), flush=True) | |
| def evaluate(a): | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| from eval import build_model | |
| device = torch.device("cuda") | |
| tok, mp, mid, mt = tokenizer_and_mapping() | |
| yes, no = mp["Yes"], mp["No"] | |
| split = {"test": "test", "dev": "validation"}[a.split] | |
| canon = PROC + ("/test.jsonl" if a.split == "test" else "/dev.jsonl") | |
| target_q = {} | |
| for l in open(canon): | |
| r = json.loads(l) | |
| target_q[r["group_id"]] = next(iter(r["questions"])) | |
| states = load_states(split, list(target_q)) | |
| states, dropped = prepare(states, tok, mid, mt, 10**9) | |
| if a.limit_states: | |
| states = states[:a.limit_states] | |
| model = build_model(a.model, 0, False, device, mp) | |
| if a.compile: | |
| model.compile_layers() | |
| mq = MQModel(model, yes, no, mid) | |
| recs = {(s["id"], x["qid"]): dict(state=s["id"], qid=x["qid"], task=s["task"], level=s["level"], prim=x["prim"], | |
| gold=int(np.argmax(x["target"])), n_q=len(s["qs"]), | |
| is_target=x["qid"] == target_q.get(s["id"]), pred={}) | |
| for s in states for x in s["qs"]} | |
| order = sorted(range(len(states)), key=lambda i: len(states[i]["ids"])) | |
| chunks = [order[i:i + a.batch] for i in range(0, len(order), a.batch)] | |
| t0 = time.time() | |
| for ch in chunks: | |
| S = [states[i] for i in ch] | |
| allpos = [[x["pos"] for x in s["qs"]] for s in S] | |
| # parallel | |
| scs = mq.scores(mq.forward([s["ids"] for s in S]), allpos) | |
| for s, sc_l in zip(S, scs): | |
| for x, sc in zip(s["qs"], sc_l): | |
| recs[(s["id"], x["qid"])]["pred"]["parallel"] = q_probs(sc, x["prim"]).tolist() | |
| # sequential: commit the most confident (top-2 margin) open question per pass | |
| ids = [list(s["ids"]) for s in S] | |
| open_q = [set(range(len(s["qs"]))) for s in S] | |
| for _ in range(max(len(s["qs"]) for s in S)): | |
| scs = mq.scores(mq.forward(ids), allpos) | |
| for i, (s, sc_l) in enumerate(zip(S, scs)): | |
| if not open_q[i]: | |
| continue | |
| best, bm, bp = None, -1, None | |
| for j in open_q[i]: | |
| pr = q_probs(sc_l[j], s["qs"][j]["prim"]) | |
| top = torch.topk(pr, 2).values if len(pr) > 1 else pr | |
| m = float(top[0] - (top[1] if len(top) > 1 else 0)) | |
| if m > bm: | |
| best, bm, bp = j, m, pr | |
| x = s["qs"][best] | |
| recs[(s["id"], x["qid"])]["pred"]["sequential"] = bp.tolist() | |
| recs[(s["id"], x["qid"])]["commit_rank"] = len(s["qs"]) - len(open_q[i]) | |
| for p, tk in zip(x["pos"], answer_tokens(x, int(bp.argmax()), yes, no)): | |
| ids[i][p] = tk | |
| open_q[i].discard(best) | |
| # step budgets: same questions, B passes total; each pass commits ceil(open / passes left) most-confident answers | |
| for B in a.step_budgets: | |
| ids = [list(s["ids"]) for s in S] | |
| open_q = [set(range(len(s["qs"]))) for s in S] | |
| for k in range(B): | |
| scs = mq.scores(mq.forward(ids), allpos) | |
| for i, (s, sc_l) in enumerate(zip(S, scs)): | |
| if not open_q[i]: | |
| continue | |
| n_commit = math.ceil(len(open_q[i]) / (B - k)) | |
| cand = [] | |
| for j in open_q[i]: | |
| pr = q_probs(sc_l[j], s["qs"][j]["prim"]) | |
| top = torch.topk(pr, 2).values if len(pr) > 1 else pr | |
| cand.append((float(top[0] - (top[1] if len(top) > 1 else 0)), j, pr)) | |
| for _, j, pr in sorted(cand, key=lambda c: -c[0])[:n_commit]: | |
| x = s["qs"][j] | |
| recs[(s["id"], x["qid"])]["pred"]["steps%d" % B] = pr.tolist() | |
| for p, tk in zip(x["pos"], answer_tokens(x, int(pr.argmax()), yes, no)): | |
| ids[i][p] = tk | |
| open_q[i].discard(j) | |
| # soft refinement: K passes, still-masked answer slots carry soft beliefs from the previous pass | |
| if a.soft_k: | |
| soft = None | |
| for k in range(1, a.soft_k + 1): | |
| scs = mq.scores(mq.forward([s["ids"] for s in S], soft), allpos) | |
| soft = [{} for _ in S] | |
| for i, (s, sc_l) in enumerate(zip(S, scs)): | |
| for x, sc in zip(s["qs"], sc_l): | |
| pr = q_probs(sc, x["prim"]) | |
| recs[(s["id"], x["qid"])]["pred"]["soft%d" % k] = pr.tolist() | |
| soft[i].update(beliefs(x, pr)) | |
| # oracle: each question with all others revealed (gold) | |
| canv, meta = [], [] | |
| for s in (S if not a.no_oracle else []): | |
| for j, x in enumerate(s["qs"]): | |
| ids_ = list(s["ids"]) | |
| for jj, y in enumerate(s["qs"]): | |
| if jj != j: | |
| for p, tk in zip(y["pos"], answer_tokens(y, int(np.argmax(y["target"])), yes, no)): | |
| ids_[p] = tk | |
| canv.append(ids_) | |
| meta.append((s, x)) | |
| for i in range(0, len(canv), a.batch): | |
| scs = mq.scores(mq.forward(canv[i:i + a.batch]), [[x["pos"]] for _, x in meta[i:i + a.batch]]) | |
| for (s, x), sc_l in zip(meta[i:i + a.batch], scs): | |
| recs[(s["id"], x["qid"])]["pred"]["oracle"] = q_probs(sc_l[0], x["prim"]).tolist() | |
| os.makedirs(a.out, exist_ok=True) | |
| rows = list(recs.values()) | |
| with open(os.path.join(a.out, "predictions.jsonl"), "w") as f: | |
| for r in rows: | |
| f.write(json.dumps(r) + "\n") | |
| modes = sorted(rows[0]["pred"]) | |
| acc = lambda rs, m: float(np.mean([np.argmax(r["pred"][m]) == r["gold"] for r in rs])) if rs else float("nan") | |
| rep = dict(model=a.model, n_states=len(states), n_questions=len(rows), seconds=round(time.time() - t0, 1), modes={}) | |
| multi = [r for r in rows if r["n_q"] > 1] | |
| for m in modes: | |
| rep["modes"][m] = dict(all=acc(rows, m), multi_q=acc(multi, m), target=acc([r for r in rows if r["is_target"]], m), | |
| by_level={str(l): acc([r for r in rows if r["level"] == l], m) for l in range(5)}, | |
| by_task={t: acc([r for r in rows if r["task"] == t], m) for t in sorted({r["task"] for r in rows})}) | |
| mm = rep["modes"][m] | |
| print("%-11s all %.2f multi-q %.2f target %.2f | %s" % (m, 100 * mm["all"], 100 * mm["multi_q"], 100 * mm["target"], | |
| " ".join("L%s %.2f" % (l, 100 * v) for l, v in mm["by_level"].items())), flush=True) | |
| 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("--n-states", type=int, default=32000) | |
| ap.add_argument("--max-length", type=int, default=8192) | |
| ap.add_argument("--lr-bb", type=float, default=2.5e-5) | |
| ap.add_argument("--budget", type=int, default=40000) | |
| ap.add_argument("--soft", action="store_true") | |
| ap.add_argument("--waste", type=float, default=0.15) | |
| ap.add_argument("--min-tokens", type=int, default=2048) | |
| ap.add_argument("--max-steps", type=int, default=0) | |
| ap.add_argument("--no-save", action="store_true") | |
| ap.add_argument("--seed", type=int, default=42) | |
| ap.add_argument("--model", default="") | |
| ap.add_argument("--split", choices=["test", "dev"], default="test") | |
| ap.add_argument("--soft-k", type=int, default=0) | |
| ap.add_argument("--batch", type=int, default=16) | |
| ap.add_argument("--compile", action="store_true") | |
| ap.add_argument("--limit-states", type=int, default=0) | |
| ap.add_argument("--step-budgets", default="", help="e.g. 2,3,4: evaluate with B passes total") | |
| ap.add_argument("--no-oracle", action="store_true") | |
| a = ap.parse_args() | |
| a.step_budgets = [int(x) for x in a.step_budgets.split(",") if x] | |
| os.makedirs(a.out, exist_ok=True) | |
| train(a) if a.mode == "train" else evaluate(a) | |
Xet Storage Details
- Size:
- 21.1 kB
- Xet hash:
- be625dca6490056302b59260af2dcb383857f5a82dab10e66f13117f45b2a105
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.