Buckets:
| """E3 step 2: frozen S1 reads input + notes (native readout). Keep a note only if it lowers the teacher's | |
| CE to the target vs no note (= raises p(gold) for hard targets). Cache, for kept items, the teacher's | |
| option-slot hidden states (outputs of blocks 14 and 28, pre-norm, LayerNorm-normalised, bf16) and probs. | |
| --placement before: notes between question and options (plan); after: right after the last answer mask | |
| (where the student's thought slots sit). | |
| """ | |
| import argparse | |
| import json | |
| import os | |
| import sys | |
| from collections import Counter, defaultdict | |
| 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, Renderer, tokenizer_and_mapping # noqa: E402 | |
| from eval import build_model # noqa: E402 | |
| from model import collate # noqa: E402 | |
| from bench_local import to_internal # noqa: E402 | |
| from laya.common import render_options, serialize_state # noqa: E402 | |
| class NotesRenderer: | |
| """jev-dllm's decision_input with a Notes block; with note=None it reproduces decision_input exactly.""" | |
| def __init__(self, tok, mask_id, mask_text, placement): | |
| self.tok, self.mask_id, self.mask_text, self.placement = tok, mask_id, mask_text, placement | |
| def __call__(self, state, qdef, note=None): | |
| q = to_internal(qdef) | |
| options = render_options(q) | |
| user = "State:\n%s\n\nQuestion:\n%s\n\n" % (serialize_state(state), q["ins"]) | |
| if note is not None and self.placement == "before": | |
| user += "Notes:\n%s\n\n" % note | |
| if q["t"] == "noul": | |
| user += "Answer Yes or No.\nAnswer:" + self.mask_text | |
| count = 1 | |
| else: | |
| user += "For each option, mark Yes if it answers the question, otherwise No.\n" | |
| user += "\n".join(o + ": " + self.mask_text for o in options) | |
| count = len(options) | |
| if note is not None and self.placement == "after": | |
| user += "\nNotes:\n%s" % note | |
| prompt = self.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(self.mask_text) | |
| assert len(parts) == count + 1 | |
| ids, pos = [], [] | |
| for i, part in enumerate(parts): | |
| ids.extend(self.tok.encode(part, add_special_tokens=False)) | |
| if i < count: | |
| pos.append(len(ids)) | |
| ids.append(self.mask_id) | |
| return ids, pos | |
| def ce(p, t): | |
| return float(-(np.asarray(t) * np.log(np.maximum(p, 1e-12))).sum()) | |
| def teacher_pass(model, items, device, pad_id, capture=(14, 28)): | |
| """items with ids/pos/primitive -> per item probs [n] and LN-normalised hidden [n_slots, d] per layer.""" | |
| out = [None] * len(items) | |
| order = sorted(range(len(items)), key=lambda i: len(items[i]["ids"])) | |
| chunks, cur, L = [], [], 0 | |
| for i in order: | |
| n = len(items[i]["ids"]) | |
| if cur and (n * (len(cur) + 1) > 65536 or len(cur) >= 32): | |
| chunks.append(cur) | |
| cur = [] | |
| cur.append(i) | |
| if cur: | |
| chunks.append(cur) | |
| for idx in chunks: | |
| mb = [dict(items[i], think=[]) for i in idx] | |
| batch = collate(mb, device, pad_id) | |
| outs, caps = model(batch, K=1, capture=capture) | |
| s = outs[0].float() | |
| for j, i in enumerate(idx): | |
| n = len(items[i]["pos"]) | |
| if items[i]["primitive"] == "noul": | |
| p1 = torch.sigmoid(s[j, 0]).item() | |
| probs = np.array([1 - p1, p1]) | |
| else: | |
| probs = torch.softmax(s[j, :n], -1).cpu().numpy().astype(np.float64) | |
| hs = {l: F.layer_norm(caps[l][j, :n].float(), (caps[l].shape[-1],)).to(torch.bfloat16).cpu() | |
| for l in capture} | |
| out[i] = (probs, hs) | |
| return out | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--notes", default="/workspace/think-decide/runs/e3/notes.jsonl") | |
| ap.add_argument("--out", default="/workspace/think-decide/runs/e3/teacher_cache.pt") | |
| ap.add_argument("--placement", choices=["before", "after"], default="before") | |
| ap.add_argument("--limit", type=int, default=0) | |
| ap.add_argument("--probe-only", action="store_true", help="report keep rate/teacher accuracy, no cache") | |
| a = ap.parse_args() | |
| device = torch.device("cuda") | |
| tok, mapping, mask_id, mask_text = tokenizer_and_mapping() | |
| nr = NotesRenderer(tok, mask_id, mask_text, a.placement) | |
| notes = {} | |
| for l in open(a.notes): | |
| n = json.loads(l) | |
| notes[(n["case_id"], n["view_id"])] = n | |
| rows = {} | |
| for l in open(DATA + "/train.jsonl"): | |
| r = json.loads(l) | |
| if (r["case_id"], r["view_id"]) in notes: | |
| rows[(r["case_id"], r["view_id"])] = r | |
| keys = [k for k in notes if notes[k]["reject"] is None] | |
| if a.limit: | |
| keys = keys[: a.limit] | |
| # check: with note=None the renderer equals jev-dllm's decision_input | |
| plain = Renderer(tok, mask_id, mask_text, 0) | |
| for k in keys[:200]: | |
| r = rows[k] | |
| (_, q), = r["questions"].items() | |
| q = {kk: q[kk] for kk in ("type", "instructions", "criteria") if kk in q} | |
| assert nr(r["state"], q)[0] == plain(r["state"], q)[0], "renderer mismatch" | |
| with_n, without = [], [] | |
| for k in keys: | |
| r = rows[k] | |
| (qid, q), = r["questions"].items() | |
| q = {kk: q[kk] for kk in ("type", "instructions", "criteria") if kk in q} | |
| crit = q["criteria"] | |
| ks = list(crit) if q["type"] == "choice" else ([str(i) for i in range(len(crit))] if q["type"] == "score" | |
| else ["false", "true"]) | |
| target = [float(r["gold"][qid]["probabilities"][x]) for x in ks] | |
| for lst, note in ((with_n, notes[k]["note"]), (without, None)): | |
| ids, pos = nr(r["state"], q, note) | |
| lst.append(dict(ids=ids, pos=pos, primitive=q["type"], target=target, key=k, | |
| source=r["source"]["name"], kind=r["gold"][qid]["kind"])) | |
| model = build_model("s1", 0, False, device, mapping) | |
| tw = teacher_pass(model, with_n, device, tok.pad_token_id) | |
| t0 = teacher_pass(model, without, device, tok.pad_token_id, capture=(28,)) | |
| stats = defaultdict(Counter) | |
| cache = {} | |
| for x, (pw, hw), (p0, _) in zip(with_n, tw, t0): | |
| keep = ce(pw, x["target"]) < ce(p0, x["target"]) | |
| g = int(np.argmax(x["target"])) | |
| for s in (x["source"], "ALL"): | |
| st = stats[s] | |
| st["n"] += 1 | |
| st["keep"] += keep | |
| if x["kind"] == "hard": | |
| st["hard"] += 1 | |
| st["acc_with"] += int(pw.argmax() == g) | |
| st["acc_without"] += int(p0.argmax() == g) | |
| if keep and not a.probe_only: | |
| cache["%s|%s" % x["key"]] = dict(p=torch.tensor(pw, dtype=torch.float32), h14=hw[14], h28=hw[28]) | |
| report = {s: dict(n=v["n"], keep_rate=round(v["keep"] / v["n"], 3), | |
| acc_with=round(v["acc_with"] / max(1, v["hard"]), 3), | |
| acc_without=round(v["acc_without"] / max(1, v["hard"]), 3)) for s, v in stats.items()} | |
| print(json.dumps(dict(placement=a.placement, report=report), indent=1), flush=True) | |
| json.dump(dict(placement=a.placement, report=report, n=len(keys)), | |
| open(os.path.join(os.path.dirname(a.out), "teacher_probe_%s%s.json" % (a.placement, "" if a.probe_only else "_full")), "w"), indent=1) | |
| if not a.probe_only: | |
| torch.save(cache, a.out) | |
| json.dump(dict(placement=a.placement, report=report, kept=len(cache)), | |
| open(a.out.replace(".pt", "_report.json"), "w"), indent=1) | |
| print(json.dumps(dict(event="saved", kept=len(cache), path=a.out)), flush=True) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 7.98 kB
- Xet hash:
- 21ee8a533922c884dc1e123fdfd166b191f9d94aa2f3deb60f54f44f3b8f35e3
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.