| """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")))) | |