item 20 scored; item 21 (GSM carry-CoT + control) pre-registered and queued

2026-07-16 02:42:46 +02:00
parent 4b9e168838
commit 1618c206ae
9 changed files with 7431 additions and 0 deletions
@@ -0,0 +1,66 @@
+"""E2 stage A prep (item 21): harvest TERSE verified CoTs for GSM8K train.
+
+For each non-drop train item, sample a compressed scratchpad ("at most 3
+short steps"), keep it only if the final number matches gold (STaR
+filter). Output: results-loop/gsm_cot_data.json rows
+{idx, label, cot} — the dense supervision the latent-mode arms never had.
+"""
+
+import json
+import os
+import sys
+import time
+from pathlib import Path
+
+import torch
+
+from loop_common import chat_prompt, last_number, num_eq
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from jlens.core import load_model  # noqa: E402
+
+OUT = Path(os.environ.get("LOOP_OUT",
+    Path(__file__).resolve().parent.parent / "results-loop"))
+TERSE_SUFFIX = ("\nSolve in at most 3 short steps, one line each, digits "
+                "only (like '4*6=24'). Then give the last line exactly as "
+                "'Answer: N'.")
+BATCH = 16
+MAX_NEW = 120
+
+
+@torch.no_grad()
+def main():
+    model, tok = load_model(dtype=torch.bfloat16)
+    tok.padding_side = "left"
+    items = [it for it in json.load(open(OUT / "star_data.json"))
+             if it["split"] == "train" and it["label"] != "drop"]
+    print(f"harvesting terse CoTs for {len(items)} train items", flush=True)
+    rows, kept = [], 0
+    t0 = time.time()
+    for i in range(0, len(items), BATCH):
+        chunk = items[i : i + BATCH]
+        enc = tok([chat_prompt(tok, it["question"], TERSE_SUFFIX)
+                   for it in chunk], return_tensors="pt", padding=True,
+                  add_special_tokens=False).to("cuda")
+        gen = model.generate(**enc, max_new_tokens=MAX_NEW, do_sample=False)
+        for j, it in enumerate(chunk):
+            txt = tok.decode(gen[j, enc["input_ids"].shape[1]:],
+                             skip_special_tokens=True).strip()
+            ok = num_eq(last_number(txt), it["gold"])
+            if ok:
+                kept += 1
+                rows.append({"idx": it["idx"], "label": it["label"],
+                             "cot": txt})
+        if i % 80 == 0:
+            print(f"[{i+len(chunk)}/{len(items)}] kept={kept} "
+                  f"({time.time()-t0:.0f}s)", flush=True)
+    json.dump(rows, open(OUT / "gsm_cot_data.json", "w"), indent=1)
+    by = {}
+    for r in rows:
+        by[r["label"]] = by.get(r["label"], 0) + 1
+    print(f"wrote {len(rows)} verified terse CoTs {by} -> gsm_cot_data.json",
+          flush=True)
+
+
+if __name__ == "__main__":
+    main()