Buckets:
| """S1MB with routing: the fast readout answers every judgment; the lowest-margin share goes to thinking mode. | |
| python s1mb_think.py run --name rl60_routed30 --fast rl60 --model runs/v2/checkpoints/rl60_merged --reason-ref rl60 | |
| python s1mb_think.py run --name v2_2b_hybrid30 --fast v2_2b_ro --model openbmb/MiniCPM5-2B --reason-ref base_2b | |
| python s1mb_think.py combine --name rl60_routed30 --fast rl60 --reason-ref rl60 # CPU: writes data/s1mb/runs/<name>.json | |
| Routing = the deployed rule: rank all 26,269 judgments by the fast readout's margin (top-1 minus top-2 probability, | |
| after the fast readout's per-type temperatures from our calibration items) and send the lowest --frac to thinking. | |
| Thinking sees the same options as the fast readout's final list (for >26 options: the tournament's last list). | |
| The routed answer is the letter readout after the chain, with post-chain temperatures fitted on our calibration chains | |
| (--reason-ref names that model's thinking predictions); an unfinished chain keeps the fast answer. The combined run is | |
| written as already-calibrated log-probs, so `s1mb_eval.py score` reads it like any other run (its "raw" mode). | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import asyncio | |
| import json | |
| import os | |
| import numpy as np | |
| from common import read_jsonl | |
| from s1mb_eval import D, QT, to_item | |
| THINK = f"{D}/think" | |
| def log_softmax(lp: np.ndarray, T: float) -> np.ndarray: | |
| z = np.asarray(lp, dtype=np.float64) / T | |
| z = z - z.max() | |
| return z - np.log(np.exp(z).sum()) | |
| def routed(rows, fast, temps, frac: float) -> list[int]: | |
| margins = [] | |
| for r, a in zip(rows, fast["answers"]): | |
| p = np.sort(np.exp(log_softmax(a["lp"], temps[QT[r["task"]]])))[::-1] | |
| margins.append(p[0] - (p[1] if len(p) > 1 else 0.0)) | |
| k = int(round(frac * len(rows))) | |
| return sorted(np.argsort(np.asarray(margins), kind="stable")[:k].tolist()) | |
| def refs(fast_name: str, reason_ref: str): | |
| from jevals_eval import PRED, REASON, temps_for_model, temps_post_for_model | |
| return temps_for_model(PRED[fast_name]), temps_post_for_model(REASON[reason_ref]) | |
| def run(args) -> None: | |
| from transformers import AutoTokenizer | |
| from vllm_utils import make_async_engine, think_async | |
| rows = read_jsonl(f"{D}/items.jsonl") | |
| fast = json.load(open(f"{D}/runs/{args.fast}.json")) | |
| t_fast, _ = refs(args.fast, args.reason_ref) | |
| idx = routed(rows, fast, t_fast, args.frac) | |
| os.makedirs(THINK, exist_ok=True) | |
| raw = f"{THINK}/{args.name}.raw.jsonl" | |
| done = {r["id"] for r in read_jsonl(raw)} if os.path.exists(raw) else set() | |
| jobs = [] | |
| for i in idx: | |
| it, cands = to_item(rows[i]), fast["answers"][i]["cands"] | |
| if it["id"] in done: | |
| continue | |
| sub = {**it, "options": [it["options"][j] for j in cands]} | |
| jobs.append((sub, list(range(len(cands))))) | |
| print(json.dumps({"routed": len(idx), "already_done": len(done), "to_run": len(jobs), | |
| "by_task": {t: sum(rows[i]["task"] == t for i in idx) for t in ("noul", "choice", "score")}}), flush=True) | |
| if jobs: | |
| tok = AutoTokenizer.from_pretrained(args.model) | |
| engine = make_async_engine(args.model, gpu_mem=args.gpu_mem) | |
| stats = asyncio.run(think_async(engine, tok, jobs, raw, n=1, deadline=float("inf"), | |
| concurrency=args.concurrency, post_readout=True)) | |
| print(json.dumps({"think_stats": stats}, default=str), flush=True) | |
| combine(args) | |
| def combine(args) -> None: | |
| rows = read_jsonl(f"{D}/items.jsonl") | |
| fast = json.load(open(f"{D}/runs/{args.fast}.json")) | |
| t_fast, t_post = refs(args.fast, args.reason_ref) | |
| idx = set(routed(rows, fast, t_fast, args.frac)) | |
| chains = {r["id"]: r["samples"][0] for r in read_jsonl(f"{THINK}/{args.name}.raw.jsonl")} | |
| answers, used, unfinished, missing, tokens = [], 0, 0, 0, [] | |
| for i, (r, a) in enumerate(zip(rows, fast["answers"])): | |
| qt = QT[r["task"]] | |
| lp = log_softmax(a["lp"], t_fast[qt]) | |
| if i in idx: | |
| s = chains.get(to_item(r)["id"]) | |
| if s is None: | |
| missing += 1 | |
| else: | |
| tokens.append(s["tokens"]) | |
| if s["closed"] and s.get("post_lp") is not None: | |
| lp, used = log_softmax(s["post_lp"], t_post[qt]), used + 1 | |
| else: | |
| unfinished += 1 | |
| answers.append({"bench": r["bench"], "case_id": r["case_id"], "qid": r["qid"], "cands": a["cands"], "lp": lp.tolist()}) | |
| if missing: | |
| raise SystemExit(f"{missing} routed judgments have no chain yet; run `run` first") | |
| json.dump({"model": args.model, "name": args.name, "mode": f"routed {args.frac:.0%} to thinking", "fast": args.fast, | |
| "temps_fast": t_fast, "temps_post": t_post, "answers": answers}, open(f"{D}/runs/{args.name}.json", "w")) | |
| summary = {"name": args.name, "routed": len(idx), "post_chain_used": used, "unfinished": unfinished, | |
| "tokens_median": float(np.median(tokens)), "tokens_mean": float(np.mean(tokens)), | |
| "thinking_tokens_per_judgment": float(np.sum(tokens) / len(rows))} | |
| json.dump(summary, open(f"{THINK}/{args.name}.summary.json", "w"), indent=1) | |
| print(json.dumps(summary), flush=True) | |
| def main() -> None: | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("cmd", choices=["run", "combine"]) | |
| ap.add_argument("--name", required=True) | |
| ap.add_argument("--fast", required=True, help="fast-readout run in data/s1mb/runs (also picks its temperatures)") | |
| ap.add_argument("--reason-ref", required=True, help="model whose calibration chains give the post-chain temperatures") | |
| ap.add_argument("--model", default="") | |
| ap.add_argument("--frac", type=float, default=0.3) | |
| ap.add_argument("--gpu-mem", type=float, default=0.92) | |
| ap.add_argument("--concurrency", type=int, default=256) | |
| args = ap.parse_args() | |
| run(args) if args.cmd == "run" else combine(args) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 6.04 kB
- Xet hash:
- 8ca132a2220247eb068e59e325cd38e9430e98c6346430351e1b87dc9e6c13ea
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.