Buckets:
| """Thinking-mode sampling with vLLM: the throughput probe and the time-boxed generation stage. | |
| python gen_reasoning.py probe --model M # 4 samples on the first 128 pool items -> results/probe.json | |
| python gen_reasoning.py generate --model M --deadline <unix time> | |
| Both write data/gen.jsonl, one line per finished item (option order, samples, parsed answers, token counts); | |
| generation walks the pool in its fixed order, skips items already written, and stops taking new items at the | |
| deadline while in-flight items finish. Chains are never cut: max_tokens is the context minus the prompt. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import asyncio | |
| import collections | |
| import json | |
| import os | |
| import time | |
| import numpy as np | |
| from transformers import AutoTokenizer | |
| import random | |
| from common import read_jsonl, shuffled_order | |
| from vllm_utils import think_async | |
| def summarize(recs: list[dict], pool: dict, seconds: float | None = None) -> dict: | |
| samples = [s for r in recs for s in r["samples"]] | |
| toks = np.array([s["tokens"] for s in samples]) if samples else np.zeros(1) | |
| correct_any, vote_ok = 0, 0 | |
| for r in recs: | |
| it = pool[r["id"]] | |
| answers = [r["order"][s["answer"]] for s in r["samples"] if s["answer"] is not None and s["closed"]] | |
| correct_any += any(a == it["gold"] for a in answers) | |
| if answers: | |
| vote_ok += collections.Counter(answers).most_common(1)[0][0] == it["gold"] | |
| out = {"items": len(recs), "samples": len(samples), "tokens_total": int(toks.sum()), | |
| "tokens_median": float(np.median(toks)), "tokens_p90": float(np.percentile(toks, 90)), | |
| "tokens_max": int(toks.max()), "parse_rate": float(np.mean([s["answer"] is not None for s in samples])), | |
| "looped": int(sum(s["looped"] for s in samples)), | |
| "items_with_correct_chain": correct_any, "majority_accuracy": vote_ok / max(len(recs), 1)} | |
| if seconds: | |
| out["seconds"] = seconds | |
| out["gen_tok_per_s"] = float(toks.sum() / seconds) | |
| return out | |
| def gen_order(item: dict) -> list[int]: | |
| """One fixed shuffled option order per item (score levels keep theirs); saved with each record.""" | |
| return shuffled_order(item, random.Random(f"gen:{item['id']}")) | |
| def main() -> None: | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("stage", choices=["probe", "generate"]) | |
| ap.add_argument("--model", default="Qwen/Qwen3.5-4B") | |
| ap.add_argument("--pool", default="data/train_pool.jsonl") | |
| ap.add_argument("--out", default="data/gen.jsonl") | |
| ap.add_argument("--n", type=int, default=4) | |
| ap.add_argument("--probe-items", type=int, default=128) | |
| ap.add_argument("--budget-min", type=float, default=165.0, help="generation budget, used by the probe to fix N") | |
| ap.add_argument("--deadline", type=float, default=None) | |
| ap.add_argument("--concurrency", type=int, default=256) | |
| ap.add_argument("--limit", type=int, default=None, help="preflight only") | |
| args = ap.parse_args() | |
| tok = AutoTokenizer.from_pretrained(args.model) | |
| pool_rows = read_jsonl(args.pool) | |
| pool = {x["id"]: x for x in pool_rows} | |
| done = {r["id"] for r in read_jsonl(args.out)} if os.path.exists(args.out) else set() | |
| if args.stage == "probe": | |
| jobs = [(x, gen_order(x)) for x in pool_rows[:args.probe_items] if x["id"] not in done] | |
| t0 = time.time() | |
| stats = asyncio.run(think_async(args.model, tok, jobs, args.out, args.n, deadline=float("inf"), | |
| concurrency=args.concurrency)) | |
| recs = [r for r in read_jsonl(args.out) if r["id"] in {x["id"] for x in pool_rows[:args.probe_items]}] | |
| summ = summarize(recs, pool, time.time() - t0) | {"engine": stats} | |
| per_item_tokens = summ["tokens_total"] / max(summ["items"], 1) | |
| summ["N_fit"] = int(summ["gen_tok_per_s"] * args.budget_min * 60 / max(per_item_tokens, 1)) | |
| os.makedirs("results", exist_ok=True) | |
| with open("results/probe.json", "w") as f: | |
| json.dump(summ, f, indent=1) | |
| print(json.dumps(summ, indent=1)) | |
| return | |
| rows = pool_rows if args.limit is None else pool_rows[:args.limit] | |
| jobs = [(x, gen_order(x)) for x in rows if x["id"] not in done] | |
| deadline = args.deadline or (time.time() + args.budget_min * 60) | |
| print(json.dumps({"to_generate": len(jobs), "already_done": len(done), | |
| "minutes_left": (deadline - time.time()) / 60}), flush=True) | |
| t0 = time.time() | |
| stats = asyncio.run(think_async(args.model, tok, jobs, args.out, args.n, deadline=deadline, | |
| concurrency=args.concurrency)) | |
| recs = read_jsonl(args.out) | |
| summ = summarize(recs, pool) | {"this_run": stats, "seconds_this_run": time.time() - t0} | |
| os.makedirs("results", exist_ok=True) | |
| with open("results/generation.json", "w") as f: | |
| json.dump(summ, f, indent=1) | |
| print(json.dumps(summ, indent=1)) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 4.99 kB
- Xet hash:
- 410397c4cd0913b6f2d12648c4b5026acda140f70bd858f0acea7b294a2d5fc0
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.