67 lines
2.4 KiB
Python
67 lines
2.4 KiB
Python
"""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()
|