| """Rejection-sampled reasoning traces from the base model itself (STaR-style), for SinhalaMMLU dev questions. |
| |
| For every question the base model answers with thinking mode and temperature sampling; traces whose final answer equals |
| the gold key are kept. Output: <out>.all.jsonl (every sample) and <out>.jsonl (correct ones, one record per kept sample: |
| the question item + "completion" = raw model output including the thought channel and the end-of-turn token). |
| |
| Time-boxed and resumable: samples are generated in rounds (all questions once, then again), results are appended after |
| every chunk, and the run stops when --max-minutes is reached. |
| |
| python scripts/gen_traces.py --data smmlu_dev_rest1450 --out data/synth/rft_traces --samples 2 --max-minutes 240 |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| import time |
| from pathlib import Path |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
| from eval.common import cap_gpu_memory, load_env, read_jsonl, set_seed |
| from eval.run_eval import resolve_data |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--data", required=True) |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--model", default="models/gemma-4-E4B-it-text") |
| ap.add_argument("--samples", type=int, default=2) |
| ap.add_argument("--temperature", type=float, default=0.8) |
| ap.add_argument("--max-new-tokens", type=int, default=3072) |
| ap.add_argument("--chunk", type=int, default=128) |
| ap.add_argument("--max-batch", type=int, default=64, help="sequences per generate call (lower it on a smaller GPU)") |
| ap.add_argument("--max-minutes", type=float, default=240) |
| ap.add_argument("--limit", type=int) |
| ap.add_argument("--seed", type=int, default=1234) |
| args = ap.parse_args() |
| load_env() |
| set_seed(args.seed) |
| cap_gpu_memory() |
| from eval.prompts import PromptBuilder |
| from eval.scorers import Scorer, load_model |
|
|
| items = read_jsonl(resolve_data(args.data))[: args.limit] |
| all_path, ok_path = Path(f"{args.out}.all.jsonl"), Path(f"{args.out}.jsonl") |
| all_path.parent.mkdir(parents=True, exist_ok=True) |
| done = {(r["uid"], r["k"]) for r in read_jsonl(all_path)} if all_path.exists() else set() |
| jobs = [(it, k) for k in range(args.samples) for it in items if (it["uid"], k) not in done] |
| print(f"{len(items)} questions x {args.samples} samples; {len(done)} already done, {len(jobs)} to go", flush=True) |
|
|
| model, tok = load_model(args.model) |
| pb = PromptBuilder(tok, lang="en", chat=True, chat_kwargs={"enable_thinking": True}) |
| scorer = Scorer(model, tok, pb.add_special_tokens, token_budget=32768, max_batch=args.max_batch, max_new_tokens=args.max_new_tokens) |
| scorer.keep_raw = True |
| scorer.sampling = {"do_sample": True, "temperature": args.temperature, "top_p": 0.95, "top_k": 64} |
|
|
| t0 = time.time() |
| n_ok = n_all = 0 |
| for i in range(0, len(jobs), args.chunk): |
| if (time.time() - t0) / 60 > args.max_minutes: |
| print(f"time box of {args.max_minutes} min reached after {n_all} samples", flush=True) |
| break |
| chunk = jobs[i:i + args.chunk] |
| res = scorer.generate([pb.mcq(it, it["options"]) for it, _ in chunk], [len(it["options"]) for it, _ in chunk]) |
| with open(all_path, "a", encoding="utf-8") as fa, open(ok_path, "a", encoding="utf-8") as fo: |
| for (it, k), r in zip(chunk, res): |
| rec = {"uid": it["uid"], "k": k, "pred": r["pred"], "gold": it["answer"], "new_tokens": r["new_tokens"], |
| "correct": r["pred"] == it["answer"]} |
| fa.write(json.dumps(rec, ensure_ascii=False) + "\n") |
| n_all += 1 |
| if rec["correct"] and r["raw"].endswith("<turn|>") and "<|channel>" in r["raw"]: |
| fo.write(json.dumps(it | {"completion": r["raw"], "k": k}, ensure_ascii=False) + "\n") |
| n_ok += 1 |
| print(f"{n_all}/{len(jobs)} samples, {n_ok} kept ({100 * n_ok / n_all:.0f}%), {(time.time() - t0) / 60:.0f} min", flush=True) |
| print(f"done: {n_all} samples, {n_ok} kept traces -> {ok_path}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
