"""Keep the teacher's reasoning traces that reach the right answer, and throw the rest away.
Purpose: reasoning distillation's filter. The teacher was asked each problem several
times; this reads all of those attempts, extracts the final answer from each,
checks it against the answer the problem generator computed, and keeps at most a
stated number of correct traces per problem. Wrong traces are dropped, traces that
reach the right answer by an unusable route are dropped, and what survives is
written in the shape train-student.py reads. The verifier is the same idea as
Part 14's reward functions: a program decides, not a model.
Platform: all (standard library only; no model, no accelerator, no network)
Minimum memory: 12 GB nominally, and far less in practice: this is text in memory
Assumes: Python 3.10 or newer. The input is the JSON Lines file written by
generate-teacher-data.py over a problem file that carries an "answer" field, such
as seeds/maths.jsonl from make-seed-prompts.py. distillog.py sits next to this file.
Usage: python3 rejection-sample.py --raw raw/maths-traces.jsonl --out-dir . \\
--keep-per-problem 1 --labbook labbook.md
python3 rejection-sample.py --raw raw/maths-traces.jsonl --out-dir . \\
--keep-per-problem 2 --max-words 500 --prefer shortest --tasks my-tasks.json
Written under --out-dir:
data-reasoning/{train,valid}.jsonl TRL conversational prompt-completion
data-reasoning-mlx/{train,valid}.jsonl mlx-lm completions
rejection-report.json per-problem attempts, correct, kept
"""
from __future__ import annotations
import argparse
import json
import random
import re
from collections import Counter
from pathlib import Path
from typing import Any
import distillog
WORD_RE = re.compile(r"[a-z0-9]+")
# Three ways a model marks its final answer, in the order this course prefers them,
# the same order Part 14's reward library uses. GSM8K's own solutions end with
# "#### 18", the maths literature uses \boxed{18}, and the course's own prompt asks
# for "Answer: 18".
FINAL_PATTERNS = (
re.compile(r"####\s*(-?[0-9][0-9,]*(?:\.[0-9]+)?)"),
re.compile(r"\\boxed\{\s*(-?[0-9][0-9,]*(?:\.[0-9]+)?)\s*\}"),
re.compile(r"(?i)\banswer\s*[:=]\s*\$?(-?[0-9][0-9,]*(?:\.[0-9]+)?)"),
)
THINK_BLOCK = re.compile(r"(.*?)", re.DOTALL)
UNCLOSED_THINK = re.compile(r"(?!.*)", re.DOTALL)
def extract_final_answer(text: str) -> float | None:
"""The last match of the most specific pattern that appears at all.
Taking the last match matters: a trace that says "so the total is 42, but wait,
Answer: 47" has two numbers and only the second is its answer.
"""
for pattern in FINAL_PATTERNS:
matches = pattern.findall(text)
if matches:
try:
return float(matches[-1].replace(",", ""))
except ValueError:
continue
return None
def is_correct(predicted: float | None, expected: float, tolerance: float = 1e-6) -> bool:
return predicted is not None and abs(predicted - expected) <= tolerance
def normalise(text: str) -> list[str]:
return WORD_RE.findall(text.lower())
def ngrams(words: list[str], n: int) -> set[tuple[str, ...]]:
if len(words) < n:
return {tuple(words)} if words else set()
return {tuple(words[i:i + n]) for i in range(len(words) - n + 1)}
def containment(a: set[tuple[str, ...]], b: set[tuple[str, ...]]) -> float:
return len(a & b) / len(a) if a else 0.0
def trace_problem(text: str, args: argparse.Namespace) -> str | None:
"""Reasons to reject a trace that nonetheless reached the right answer."""
if UNCLOSED_THINK.search(text):
return "unclosed-thinking-block"
blocks = THINK_BLOCK.findall(text)
if args.require_thinking and not blocks:
return "no-thinking-block"
if len(blocks) > 1:
return "several-thinking-blocks"
if blocks:
working = normalise(blocks[0])
if len(working) < args.min_working_words:
# A trace whose thinking block is two words long reached the answer by
# asserting it. Training on that teaches assertion, not reasoning, and
# it is checked before the length filter so the reason names the fault
# rather than the symptom.
return "no-working-shown"
words = normalise(text)
if len(words) < args.min_words:
return "too-short"
if len(words) > args.max_words:
return "too-long"
counts = Counter(tuple(words[i:i + 8]) for i in range(max(0, len(words) - 7)))
if counts and counts.most_common(1)[0][1] >= 3:
return "repetition-loop"
return None
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
parser.add_argument("--raw", required=True, help="traces from generate-teacher-data.py")
parser.add_argument("--out-dir", default=".")
parser.add_argument("--keep-per-problem", type=int, default=1,
help="how many correct traces to keep for each problem")
parser.add_argument("--prefer", choices=["shortest", "longest", "first"], default="shortest",
help="which correct traces to keep when there are more than needed")
parser.add_argument("--min-words", type=int, default=20)
parser.add_argument("--max-words", type=int, default=600)
parser.add_argument("--min-working-words", type=int, default=15)
parser.add_argument("--require-thinking", action="store_true", default=True,
help="reject a trace with no block (default)")
parser.add_argument("--allow-no-thinking", dest="require_thinking", action="store_false")
parser.add_argument("--tasks", default=None, help="evaluation set to decontaminate against")
parser.add_argument("--contamination", type=float, default=0.5)
parser.add_argument("--n", type=int, default=13)
parser.add_argument("--valid-fraction", type=float, default=0.1)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--labbook", default=None)
args = parser.parse_args()
raw_path = Path(args.raw)
rows = [json.loads(line) for line in raw_path.read_text(encoding="utf-8").splitlines() if line.strip()]
missing_answer = [r for r in rows if "answer" not in r]
if missing_answer:
raise SystemExit(
f"{len(missing_answer)} trace(s) carry no 'answer' field, so nothing can verify them. "
"Generate from a problem file that has one, such as seeds/maths.jsonl."
)
print(f"read {len(rows)} trace(s) from {raw_path}")
by_problem: dict[str, list[dict]] = {}
for row in rows:
by_problem.setdefault(row["id"], []).append(row)
rejected: Counter[str] = Counter()
kept: list[dict] = []
per_problem_stats = []
for pid, attempts in sorted(by_problem.items()):
expected = float(attempts[0]["answer"])
correct = []
for attempt in attempts:
completion = (attempt.get("completion") or "").strip()
predicted = extract_final_answer(completion)
if predicted is None:
rejected["no-answer-found"] += 1
continue
if not is_correct(predicted, expected):
rejected["wrong-answer"] += 1
continue
reason = trace_problem(completion, args)
if reason:
rejected[reason] += 1
continue
correct.append({**attempt, "completion": completion,
"words": len(normalise(completion))})
if args.prefer == "shortest":
correct.sort(key=lambda a: a["words"])
elif args.prefer == "longest":
correct.sort(key=lambda a: -a["words"])
chosen = correct[:args.keep_per_problem]
kept.extend(chosen)
per_problem_stats.append({"id": pid, "attempts": len(attempts),
"correct": len(correct), "kept": len(chosen)})
solved = sum(1 for s in per_problem_stats if s["kept"])
print(f"problems: {len(by_problem)}, at least one usable trace for {solved}")
print(f"traces kept: {len(kept)}")
for reason, count in rejected.most_common():
print(f" rejected {count:>5} {reason}")
contaminated = 0
if args.tasks:
spec = json.loads(Path(args.tasks).read_text(encoding="utf-8"))
eval_grams = [ngrams(normalise(t.get("prompt", "")), args.n) for t in spec.get("tasks", [])]
clean = []
for row in kept:
grams = ngrams(normalise(f"{row['prompt']}\n{row['completion']}"), args.n)
if any(containment(e, grams) >= args.contamination for e in eval_grams if e):
contaminated += 1
continue
clean.append(row)
kept = clean
print(f"decontamination: dropped {contaminated} trace(s) overlapping the evaluation set")
else:
print("decontamination: SKIPPED, no --tasks given.")
if not kept:
raise SystemExit("no trace survived; look at rejection-report.json before loosening a threshold")
rng = random.Random(args.seed)
rng.shuffle(kept)
cut = max(1, int(len(kept) * args.valid_fraction))
valid, train = kept[:cut], kept[cut:]
out = Path(args.out_dir)
written: dict[str, Any] = {}
for name, part in (("train", train), ("valid", valid)):
trl_path = out / "data-reasoning" / f"{name}.jsonl"
trl_path.parent.mkdir(parents=True, exist_ok=True)
with trl_path.open("w", encoding="utf-8") as handle:
for row in part:
handle.write(json.dumps({
"prompt": [{"role": "user", "content": row["prompt"]}],
"completion": [{"role": "assistant", "content": row["completion"]}],
}, ensure_ascii=False) + "\n")
mlx_path = out / "data-reasoning-mlx" / f"{name}.jsonl"
mlx_path.parent.mkdir(parents=True, exist_ok=True)
with mlx_path.open("w", encoding="utf-8") as handle:
for row in part:
handle.write(json.dumps({"prompt": row["prompt"], "completion": row["completion"]},
ensure_ascii=False) + "\n")
written[name] = {"count": len(part), "sha256": distillog.file_sha256(trl_path)}
report = {
"raw": str(raw_path),
"raw_sha256": distillog.file_sha256(raw_path),
"traces": len(rows),
"problems": len(by_problem),
"problems_with_a_usable_trace": solved,
"kept": len(kept),
"contaminated_dropped": contaminated,
"rejected_by_reason": dict(rejected.most_common()),
"per_problem": per_problem_stats,
"thresholds": {
"keep_per_problem": args.keep_per_problem, "prefer": args.prefer,
"min_words": args.min_words, "max_words": args.max_words,
"min_working_words": args.min_working_words,
"require_thinking": args.require_thinking,
},
"splits": written,
}
(out / "rejection-report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
print(f"\ntrain {len(train)} valid {len(valid)}")
print(f"report written to {out / 'rejection-report.json'}")
print("The problems with no usable trace are the interesting ones: they are the "
"problems the teacher could not solve, and the student will not learn them here.")
if args.labbook:
rec = distillog.record(
labbook=args.labbook,
lab="part-15/rejection-sample",
stage="reject",
teacher=(rows[0].get("teacher") if rows else None),
student=None,
dataset={
"raw": str(raw_path), "raw_sha256": report["raw_sha256"],
"traces": len(rows), "problems": len(by_problem),
"problems_with_a_usable_trace": solved,
"kept": len(kept), "train": len(train), "valid": len(valid),
"train_sha256": written["train"]["sha256"],
"tasks": args.tasks, "contaminated_dropped": contaminated,
},
hyperparameters=report["thresholds"],
seed=args.seed,
scores={"rejected_by_reason": report["rejected_by_reason"],
"solve_rate": round(solved / len(by_problem), 3) if by_problem else None},
config_path=__file__,
notes=None if args.tasks else "decontamination skipped: no evaluation set supplied",
)
print(f"recorded {rec['stage']} as {rec['run_id']} in {args.labbook}")
if __name__ == "__main__":
main()