File size: 4,473 Bytes
67ebcef | 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 | """Builds the SFT set from teach_gen.py output: what the module system knows how to do, as plain model outputs.
Per problem (one example each, so easy problems do not dominate):
* "short": the shortest correct natural solution among the N samples (efficient reasoning that still ends right);
* "stop": if no natural sample is shorter, a checked stop example (thinking cut where the answer was already
known and stable + </think> + the model's own closing answer) — teaches ending the thinking on time.
usage: build_sft.py --inp results/teach/pilot.jsonl --items data/teach/*.jsonl --out data/teach/sft_pilot.jsonl
"""
import argparse
import collections
import glob
import json
import zlib
from transformers import AutoTokenizer
EOS = 124900
ap = argparse.ArgumentParser()
ap.add_argument("--inp", nargs="+", required=True)
ap.add_argument("--items", default="data/teach/*.jsonl")
ap.add_argument("--out", required=True)
ap.add_argument("--tokenizer", default="models/LFM3-1.7B")
ap.add_argument("--max-len", type=int, default=7000)
ap.add_argument("--mode", default="short", choices=["short", "correctness"],
help="correctness: only problems the model solves sometimes (one random correct sample each) and "
"unsolved problems rescued by a checked stop; no length-only examples")
args = ap.parse_args()
tok = AutoTokenizer.from_pretrained(args.tokenizer)
items = {}
for p in glob.glob(args.items):
if "sft_" in p:
continue
for l in open(p, encoding="utf-8"):
r = json.loads(l)
items[r["id"]] = r
by = collections.defaultdict(list)
for p in args.inp:
for l in open(p, encoding="utf-8"):
r = json.loads(l)
by[r["id"]].append(r)
out, stats = [], collections.Counter()
for pid, recs in by.items():
it = items[pid]
good = [r for r in recs if r["correct"]]
stops = [r["stop"] for r in recs if r.get("stop", {}).get("ok")]
best_short = min(good, key=lambda r: r["n_tokens"]) if good else None
best_stop = min(stops, key=lambda s: len(s["tokens"])) if stops else None
p_ids = tok.apply_chat_template([{"role": "user", "content": it["prompt"]}], tokenize=True,
add_generation_prompt=True, return_dict=False)
if args.mode == "correctness":
rescue = [r["stop"] for r in recs if r.get("stop", {}).get("ok") and not r["stop"].get("from_correct")]
if good and len(good) < len(recs):
pick = good[zlib.crc32(pid.encode()) % len(good)]
comp, kind = [t for t in pick["tokens"] if t != EOS] + [EOS], "partial"
elif not good and rescue:
comp, kind = [t for t in rescue[0]["tokens"] if t != EOS] + [EOS], "rescue"
else:
stats["skipped"] += 1
continue
if len(p_ids) + len(comp) > args.max_len:
stats["too_long"] += 1
continue
out.append({"id": pid, "task": it["task"], "kind": kind, "prompt_ids": p_ids, "completion_ids": comp,
"solved": f"{len(good)}/{len(recs)}"})
stats[f"{it['task']}_{kind}"] += 1
continue
if best_stop and (not best_short or len(best_stop["tokens"]) < best_short["n_tokens"]):
comp, kind = [t for t in best_stop["tokens"] if t != EOS] + [EOS], "stop"
elif best_short:
comp, kind = [t for t in best_short["tokens"] if t != EOS] + [EOS], "short"
else:
stats["no_example"] += 1
continue
if len(p_ids) + len(comp) > args.max_len:
stats["too_long"] += 1
continue
out.append({"id": pid, "task": it["task"], "kind": kind, "prompt_ids": p_ids, "completion_ids": comp,
"solved": f"{len(good)}/{len(recs)}"})
stats[f"{it['task']}_{kind}"] += 1
stats[f"len_{kind}"] += len(comp)
stats[f"natural_mean_{kind}"] += sum(r["n_tokens"] for r in recs) / len(recs)
with open(args.out, "w", encoding="utf-8") as f:
for r in out:
f.write(json.dumps(r) + "\n")
for k in ("short", "stop"):
n = sum(v for kk, v in stats.items() if kk.endswith("_" + k) and not kk.startswith(("len", "natural")))
if n:
print(f"{k}: {n} examples, mean length {stats['len_' + k] / n:.0f} vs mean sample length "
f"{stats['natural_mean_' + k] / n:.0f}")
print(dict((k, v) for k, v in stats.items() if not k.startswith(("len", "natural"))))
|