| import argparse |
| import json |
| import random |
| import re |
| from pathlib import Path |
|
|
| from predict import ANSWER_CUE, answer, argmax, load_model, normalise, score, thinking_text, user_prompt, write_jsonl |
|
|
| HINT_MENTION = re.compile(r"hint|ඉඟි|(?:given|provided|stated|suggested) answer", re.IGNORECASE) |
| SENTENCE_BREAK = re.compile(r"((?<=[.!?])\s+|\n+)") |
| MIN_KEPT_FRACTION = 0.5 |
|
|
|
|
| def hinted_prompt(question): |
| gold = question["answer"] |
| hint = f"Hint: the correct answer is {gold}. {question['choices'][gold - 1]}" |
| return user_prompt(question).removesuffix(ANSWER_CUE) + hint + "\n" + ANSWER_CUE |
|
|
|
|
| def strip_hint(thought): |
| pieces = SENTENCE_BREAK.split(thought) |
| sentences = zip(pieces[0::2], pieces[1::2] + [""]) |
| return "".join(s + sep for s, sep in sentences if not HINT_MENTION.search(s)).strip() |
|
|
|
|
| def attempt(model, tokenizer, questions, think, batch_size): |
| prompts = [user_prompt(q) for q in questions] |
| results = answer(model, tokenizer, prompts, [len(q["choices"]) for q in questions], think, batch_size) |
| attempts = [ |
| {"id": q["id"], "gold": q["answer"], "pred": r["pred"], "finished": r["finished"], |
| "thought_tokens": r["thought_tokens"], "hinted": False, "thought": r["thought"]} |
| for q, r in zip(questions, results) |
| ] |
| kept = [ |
| {"id": q["id"], "prompt": p, "thought": r["thought"], "answer": q["answer"], "source": "direct"} |
| for q, p, r in zip(questions, prompts, results) |
| if r["finished"] and r["pred"] == q["answer"] |
| ] |
| return attempts, kept |
|
|
|
|
| def rationalize(model, tokenizer, questions, think, batch_size): |
| n_choices = [len(q["choices"]) for q in questions] |
| prompts = [hinted_prompt(q) for q in questions] |
| results = answer(model, tokenizer, prompts, n_choices, think, batch_size) |
| cleaned = [strip_hint(r["thought"]) for r in results] |
| texts = [thinking_text(tokenizer, user_prompt(q), c) for q, c in zip(questions, cleaned)] |
| rescored = [argmax(p) + 1 for p in score(model, tokenizer, texts, n_choices, batch_size)] |
| attempts, kept = [], [] |
| for q, r, c, pred in zip(questions, results, cleaned, rescored): |
| attempts.append( |
| {"id": q["id"], "gold": q["answer"], "pred": r["pred"], "finished": r["finished"], |
| "thought_tokens": r["thought_tokens"], "hinted": True, "thought": r["thought"], |
| "cleaned": c, "pred_without_hint": pred} |
| ) |
| if ( |
| r["finished"] |
| and r["pred"] == q["answer"] |
| and pred == q["answer"] |
| and len(c) >= MIN_KEPT_FRACTION * len(r["thought"]) |
| ): |
| kept.append({"id": q["id"], "prompt": user_prompt(q), "thought": c, "answer": q["answer"], |
| "source": "rationalized"}) |
| return attempts, kept |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--model", default="google/gemma-4-E4B-it") |
| parser.add_argument("--revision", help="model commit to load; the latest one when not given") |
| parser.add_argument("--adapter", help="LoRA adapter that writes the thoughts, merged into the model") |
| parser.add_argument("--train", default="data/train/mcq.jsonl") |
| parser.add_argument("--n", type=int, help="sample this many questions with --seed; all of them when not given") |
| parser.add_argument("--seed", type=int, default=0) |
| parser.add_argument("--out_dir", default="star") |
| parser.add_argument("--think", type=int, default=512) |
| parser.add_argument("--batch_size", type=int, default=32) |
| parser.add_argument("--no_hint_retry", action="store_true", help="keep only first-try thoughts and skip the hinted retry") |
| args = parser.parse_args() |
| if args.think <= 0: |
| parser.error("--think must be positive, STaR keeps thoughts") |
|
|
| lines = Path(args.train).read_text(encoding="utf-8").splitlines() |
| questions = [normalise(json.loads(line)) for line in lines] |
| if args.n is not None: |
| questions = random.Random(args.seed).sample(questions, args.n) |
| print(f"Sampled {len(questions)} of {len(lines)} questions from {args.train} with seed {args.seed}") |
| else: |
| print(f"Using all {len(questions)} questions from {args.train}") |
|
|
| tokenizer, model = load_model(args.model, args.adapter, args.revision) |
| out_dir = Path(args.out_dir) |
| out_dir.mkdir(parents=True, exist_ok=True) |
|
|
| attempts, kept = attempt(model, tokenizer, questions, args.think, args.batch_size) |
| write_jsonl(out_dir / "attempts.jsonl", attempts) |
| write_jsonl(out_dir / "star.jsonl", kept) |
| print(f"Direct: kept {len(kept)}/{len(questions)}") |
| if args.no_hint_retry: |
| return |
|
|
| done = {row["id"] for row in kept} |
| missed = [q for q in questions if q["id"] not in done] |
| hinted_attempts, rationalized = rationalize(model, tokenizer, missed, args.think, args.batch_size) |
| write_jsonl(out_dir / "attempts.jsonl", attempts + hinted_attempts) |
| write_jsonl(out_dir / "star.jsonl", kept + rationalized) |
| print(f"Rationalized: kept {len(rationalized)}/{len(missed)}") |
| print(f"Wrote {len(kept) + len(rationalized)} examples to {out_dir / 'star.jsonl'}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
