File size: 6,042 Bytes
0052e99 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | """Build the reasoning-effort SFT set: Beam traces at low, medium and high effort for the same prompts, each
tagged the way Dory-Beam's chat template renders `reasoning_effort` (Dory's own `{reasoning effort: ...}`
convention, appended to the user turn).
.venv/bin/python -I scripts/build_effort_dataset.py --efforts low,medium,high --pool xhigh:high --longest high
Inputs: data/effort_prompts.jsonl (the 20k prompts), low/high/xhigh generations in data/selfhost_effort/*.jsonl,
and the existing medium generations (data/gen, data/selfhost). Per prompt and level, one randomly drawn
verified-correct trace, or the longest with --longest (used for high, pooled from Beam's high and xhigh). Medium examples are tagged half the time and left untagged otherwise, so the untagged
default keeps behaving like medium. Writes data/final_effort/sft_effort.jsonl and stats.json.
"""
import argparse
import glob
import json
import random
import statistics
import sys
from collections import defaultdict
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "scripts"))
from verify import check # noqa: E402
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--prompts", default=str(ROOT / "data" / "effort_prompts.jsonl"))
ap.add_argument("--effort-inputs", default=str(ROOT / "data" / "selfhost_effort" / "*.jsonl"))
ap.add_argument("--medium-inputs", nargs="+",
default=[str(ROOT / "data" / "gen" / "*.jsonl"), str(ROOT / "data" / "selfhost" / "*.jsonl")])
ap.add_argument("--untagged-medium", type=float, default=0.5)
ap.add_argument("--out-dir", default=str(ROOT / "data" / "final_effort"))
ap.add_argument("--seed", type=int, default=0)
ap.add_argument("--efforts", default="low,medium,high", help="Beam effort levels to include")
ap.add_argument("--pool", default="", help="use one Beam level's traces as candidates for another, e.g. xhigh:high")
ap.add_argument("--longest", default="", help="levels that take the longest correct candidate instead of a random one")
args = ap.parse_args()
EFFORTS = tuple(args.efforts.split(","))
pool = dict(x.split(":") for x in args.pool.split(",") if x)
longest = set(filter(None, args.longest.split(",")))
prompts = {json.loads(l)["id"]: json.loads(l) for l in open(args.prompts)}
eval_ids = {l.strip() for l in open(ROOT / "data" / "eval_ids.txt") if l.strip()}
assert not set(prompts) & eval_ids
gens = defaultdict(list) # (id, effort) -> verified records
teachers = defaultdict(set)
files = sorted(glob.glob(args.effort_inputs))
files += sorted(p for pat in args.medium_inputs for p in glob.glob(pat))
for path in files:
for line in open(path):
r = json.loads(line)
if "error" in r or r["id"] not in prompts:
continue
beam_eff = r.get("effort", "medium")
if (beam_eff == "medium") != ("selfhost_effort" not in path):
continue # medium only from the original pool, other levels only from the effort runs
eff = pool.get(beam_eff, beam_eff)
if eff not in EFFORTS:
continue
r["beam_effort"] = beam_eff
ok, _ = check(r["content"], r["gold"], r["answer_type"])
r["correct"] = ok and r.get("finish_reason") == "stop" and bool(r.get("reasoning", "").strip())
gens[(r["id"], eff)].append(r)
teachers[eff].add(r.get("teacher", "beam-api"))
rng = random.Random(args.seed)
out = Path(args.out_dir)
out.mkdir(parents=True, exist_ok=True)
stats = {e: {"generations": 0, "correct": 0, "examples": 0, "chars": [], "tokens": []} for e in EFFORTS}
n_untagged = 0
with open(out / "sft_effort.jsonl", "w") as f:
for pid in sorted(prompts):
for eff in EFFORTS:
rs = gens.get((pid, eff), [])
good = sorted((r for r in rs if r["correct"]), key=lambda r: len(r["reasoning"]))
st = stats[eff]
st["generations"] += len(rs)
st["correct"] += len(good)
for r in rs:
st["tokens"].append(r["usage"].get("completion_tokens", 0))
if not good:
continue
r = good[-1] if eff in longest else rng.choice(good)
tagged = eff != "medium" or rng.random() >= args.untagged_medium
n_untagged += not tagged
user = r["prompt"] + (f"\n\n{{reasoning effort: {eff}}}" if tagged else "")
st["examples"] += 1
st["chars"].append(len(r["reasoning"]) + len(r["content"]))
f.write(json.dumps({
"id": f"{pid}#{eff}", "prompt_id": pid, "effort": eff, "beam_effort": r["beam_effort"],
"tagged": tagged,
"domain": r["domain"], "teacher": r.get("teacher", "beam-api"),
"messages": [{"role": "user", "content": user},
{"role": "assistant", "reasoning_content": r["reasoning"], "content": r["content"]}],
"gold": r["gold"], "chars": st["chars"][-1],
}, ensure_ascii=False) + "\n")
summary = {"prompts": len(prompts), "untagged_medium": n_untagged, "pool": pool, "longest": sorted(longest),
"teachers": {e: sorted(t) for e, t in teachers.items()}}
for e, st in stats.items():
summary[e] = {"generations": st["generations"], "correct": st["correct"],
"accuracy": round(st["correct"] / max(1, st["generations"]), 3), "examples": st["examples"],
"chars_mean": round(statistics.mean(st["chars"])) if st["chars"] else None,
"tokens_mean": round(statistics.mean(st["tokens"])) if st["tokens"] else None}
(out / "stats.json").write_text(json.dumps(summary, indent=1))
print(json.dumps(summary, indent=1))
if __name__ == "__main__":
main()
|