Buckets:
| """Numbers behind the data / training / RL report figures (fig_train_plot.py draws them). | |
| Run with the project venv on CPU: | |
| CUDA_VISIBLE_DEVICES="" HF_HUB_OFFLINE=1 .venv/bin/python analysis/fig_train_compute.py | |
| Writes results/report_figs/train_figs_data.json. Reads only run files (no GPU, no network): | |
| runs/v2/data/{train_pool,gen,gen_v1,rl_prompts}.jsonl, runs/v2/checkpoints/rl_grpo*/completions/*.parquet, | |
| the tokenizer of runs/v2/checkpoints/rl60_merged (the RL reward counted tokens with the same tokenizer). | |
| The chain selection below mirrors train.py build_examples with V2's flags (min_correct 2, second_chain_min 3, | |
| rep_max 20, drop_recall), so its counts must equal the V2 joint run's train_meta.json counts (checked at the end). | |
| """ | |
| from __future__ import annotations | |
| import collections | |
| import glob | |
| import json | |
| import os | |
| import re | |
| import sys | |
| import numpy as np | |
| R = "/workspace/decision-reasoner" | |
| sys.path.insert(0, R) | |
| from common import read_jsonl # noqa: E402 | |
| from train import RECALL_RE, repeated_line # noqa: E402 | |
| WAIT_RE = re.compile(r"\bWait\b") # as in dashboard/build_data.py rl_progress | |
| OUT = os.path.join(R, "results/report_figs/train_figs_data.json") | |
| def pct(a, q): | |
| return float(np.percentile(a, q)) if len(a) else None | |
| def pool_composition(pool): | |
| soft = lambda x: max(x["target"]) < 1.0 # soft-label (annotator-share) items | |
| groups = collections.defaultdict(collections.Counter) | |
| sources = collections.defaultdict(set) | |
| for x in pool.values(): | |
| if x["set"] == "train_proc": | |
| g = "proc" | |
| else: | |
| g = "ts_soft" if soft(x) else "ts_hard" | |
| groups[g][x["qtype"]] += 1 | |
| sources[g].add(x["family"]) | |
| return {g: {"by_qtype": dict(c), "n": sum(c.values()), "sources": len(sources[g])} for g, c in groups.items()} | |
| def chain_sources(pool, gen, v1_ids): | |
| """Per author/set: questions generated, chains sampled, correct-of-4 spread, chains trained (train.py V2 rule).""" | |
| out = {} | |
| trained_total, items_total = 0, 0 | |
| for rec in gen: | |
| it = pool.get(rec["id"]) | |
| if it is None: | |
| continue | |
| author = "teacher" if rec.get("source") == "teacher" else ("4b_v1" if rec["id"] in v1_ids else "4b_v2") | |
| st = "proc" if it["set"] == "train_proc" else "ts" | |
| a = out.setdefault(f"{author}:{st}", {"questions": 0, "samples": 0, "correct_of_4": [0] * 5, "sample_tokens": [], | |
| "trained": 0, "trained_questions": 0, "trained_tokens": [], | |
| "trained_waits": [], "sample_waits": [], "dropped_recall": 0, "dropped_rep": 0}) | |
| a["questions"] += 1 | |
| a["samples"] += len(rec["samples"]) | |
| a["sample_tokens"] += [s["tokens"] for s in rec["samples"]] | |
| a["sample_waits"] += [len(WAIT_RE.findall(s["text"])) for s in rec["samples"] if s["closed"]] | |
| correct = sorted((s for s in rec["samples"] if s["closed"] and s["answer"] is not None | |
| and rec["order"][s["answer"]] == it["gold"]), key=lambda s: s["tokens"]) | |
| a["correct_of_4"][min(len(correct), 4)] += 1 | |
| if len(correct) < 2: # min_correct 2 | |
| continue | |
| keep = [s for s in correct if not RECALL_RE.search(s["text"])] | |
| a["dropped_recall"] += len(correct) - len(keep) | |
| correct = keep | |
| if not correct: | |
| continue | |
| keep = [s for s in correct if repeated_line(s["text"]) < 20] # rep_max 20 | |
| a["dropped_rep"] += len(correct) - len(keep) | |
| correct = keep | |
| if not correct: | |
| continue | |
| chosen = correct[:2] if len(correct) >= 3 else correct[:1] # second_chain_min 3 | |
| a["trained"] += len(chosen) | |
| a["trained_questions"] += 1 | |
| a["trained_tokens"] += [s["tokens"] for s in chosen] | |
| a["trained_waits"] += [len(WAIT_RE.findall(s["text"])) for s in chosen] | |
| trained_total += len(chosen) | |
| items_total += 1 | |
| summ = {} | |
| for k, a in out.items(): | |
| summ[k] = {"questions": a["questions"], "samples": a["samples"], "correct_of_4": a["correct_of_4"], | |
| "trained": a["trained"], "trained_questions": a["trained_questions"], | |
| "dropped_recall": a["dropped_recall"], "dropped_rep": a["dropped_rep"], | |
| "sample_tokens_median": pct(a["sample_tokens"], 50), "sample_tokens_p90": pct(a["sample_tokens"], 90), | |
| "trained_tokens_median": pct(a["trained_tokens"], 50), "trained_tokens_p90": pct(a["trained_tokens"], 90), | |
| "trained_waits_median": pct(a["trained_waits"], 50), "trained_waits_mean": float(np.mean(a["trained_waits"])) if a["trained_waits"] else None, | |
| "sample_waits_mean": float(np.mean(a["sample_waits"])) if a["sample_waits"] else None, | |
| "sample_tokens": a["sample_tokens"], "trained_tokens": a["trained_tokens"], "trained_waits": a["trained_waits"]} | |
| return summ, trained_total, items_total | |
| def rl_items(): | |
| """RL prompt ids -> (kind, matching key). Prompts in the parquet files are rendered text; an item is found by its | |
| criterion and its evidence, both JSON-escaped as in common.user_message.""" | |
| esc = lambda t: json.dumps(t or "", ensure_ascii=False)[1:-1] | |
| rlp = read_jsonl(os.path.join(R, "runs/v2/data/rl_prompts.jsonl")) | |
| pool = {x["id"]: x for x in read_jsonl(os.path.join(R, "runs/v2/data/train_pool.jsonl"))} | |
| fresh = {x["id"]: x for x in read_jsonl(os.path.join(R, "runs/v2/data/rl_fresh_pool.jsonl"))} | |
| out = [] | |
| for r in rlp: | |
| it = fresh.get(r["id"]) if r.get("fresh") else pool.get(r["id"]) | |
| kind = ("fresh_reasoning" if r.get("reasoning_source") else "fresh_other") if r.get("fresh") else "in_domain" | |
| out.append({"id": r["id"], "kind": kind, "p_gold": r.get("v2_readout_p_gold"), | |
| "keys": ['"criterion": "' + esc(it["criterion"]), '{"evidence": "' + esc(it["evidence"])] if it else None}) | |
| return out | |
| def rl_rounds(tok, items): | |
| """One row per generation round (4 optimizer steps share its 128 chains). Categories from the logged reward | |
| (rl_grpo.py): no answer by the 8,192-token cap = -1.5 minus the length penalty (<= -1.5); a finished chain is | |
| >= -1.25; right = reward > 0 (+1 on exact/hard labels, the annotator share rescaled on soft ones, minus the | |
| length penalty of at most 0.25).""" | |
| phases = [("runs/v2/checkpoints/rl_grpo/completions", 0, 40), ("runs/v2/checkpoints/rl_grpo_b/completions", 40, 20)] | |
| rows, per_q = [], [] | |
| for comp, offset, keep in phases: | |
| for f in sorted(glob.glob(os.path.join(R, comp, "completions_*.parquet"))): | |
| import pandas as pd | |
| df = pd.read_parquet(f, columns=["step", "prompt", "completion", "reward"]) | |
| st = int(df["step"].iloc[0]) | |
| if (st - 1) % 4 or st > keep: | |
| continue | |
| rew = df["reward"].to_numpy() | |
| texts = df["completion"].tolist() | |
| ntok = np.array([len(tok(t, add_special_tokens=False)["input_ids"]) for t in texts]) | |
| fin = rew > -1.4 | |
| waits = np.array([len(WAIT_RE.findall(t)) for t in texts]) | |
| groups = df.assign(r=rew > 0).groupby("prompt")["r"].agg(["mean", "size"]) | |
| for prompt, g in df.assign(r=rew > 0, fin=fin, ntok=ntok).groupby("prompt"): | |
| hit = [it for it in items if it["keys"] and all(k in prompt for k in it["keys"])] | |
| per_q.append({"step": offset + st, "kind": hit[0]["kind"] if len(hit) == 1 else None, | |
| "p_gold": hit[0]["p_gold"] if len(hit) == 1 else None, "n_match": len(hit), | |
| "right_of_8": int(g["r"].sum()), "chains": int(len(g)), "unfinished": int((~g["fin"]).sum())}) | |
| rows.append({"step": offset + st, "round": len(rows) + 1, "phase": "b" if offset else "a", "chains": int(len(df)), | |
| "questions": int(len(groups)), "mean_reward": float(rew.mean()), | |
| "right": float((rew > 0).mean()), "right_among_finished": float((rew[fin] > 0).mean()), | |
| "unfinished": float((~fin).mean()), "tokens_median_all": float(np.median(ntok)), | |
| "tokens_median_finished": float(np.median(ntok[fin])), "tokens_mean_all": float(ntok.mean()), | |
| "waits_per_finished": float(waits[fin].mean()), "waits_median_finished": float(np.median(waits[fin])), | |
| "mixed_groups": float(((groups["mean"] > 0) & (groups["mean"] < 1)).mean())}) | |
| return rows, per_q | |
| def main(): | |
| pool = {x["id"]: x for x in read_jsonl(os.path.join(R, "runs/v2/data/train_pool.jsonl"))} | |
| gen = read_jsonl(os.path.join(R, "runs/v2/data/gen.jsonl")) | |
| v1_ids = {r["id"] for r in read_jsonl(os.path.join(R, "runs/v2/data/gen_v1.jsonl"))} | |
| comp = pool_composition(pool) | |
| chains, n_tr, n_items = chain_sources(pool, gen, v1_ids) | |
| meta = json.load(open(os.path.join(R, "runs/v2/results/adapters/v2_joint/train_meta.json"))) | |
| check = {"trained_chains": n_tr, "meta_reasoner_examples": meta["counts"]["reasoner_examples"], | |
| "trained_questions": n_items, "meta_reasoner_items": meta["counts"]["reasoner_items"]} | |
| print("chain check", check) | |
| v1_pool = {x["id"]: x for x in read_jsonl(os.path.join(R, "data/train_pool.jsonl"))} | |
| v1_gen_ids = [r["id"] for r in read_jsonl(os.path.join(R, "data/gen.jsonl"))] | |
| v1_q = collections.Counter(v1_pool[i]["qtype"] for i in v1_gen_ids if i in v1_pool) | |
| from transformers import AutoTokenizer | |
| tok = AutoTokenizer.from_pretrained(os.path.join(R, "runs/v2/checkpoints/rl60_merged")) | |
| rl, rl_q = rl_rounds(tok, rl_items()) | |
| print("rl questions", len(rl_q), "matched once", sum(q["n_match"] == 1 for q in rl_q), | |
| collections.Counter(q["n_match"] for q in rl_q)) | |
| rlp = read_jsonl(os.path.join(R, "runs/v2/data/rl_prompts.jsonl")) | |
| rl_mix = collections.Counter(("fresh_reasoning" if r.get("reasoning_source") else "fresh_other") if r.get("fresh") else "in_domain" | |
| for r in rlp) | |
| data = {"pool": comp, "pool_total": len(pool), "v1_generated_questions": len(v1_gen_ids), "v1_by_qtype": dict(v1_q), | |
| "chains": chains, "chain_check": check, "rl_rounds": rl, "rl_questions": rl_q, "rl_prompt_mix": dict(rl_mix), "rl_prompts": len(rlp)} | |
| os.makedirs(os.path.dirname(OUT), exist_ok=True) | |
| with open(OUT, "w") as f: | |
| json.dump(data, f) | |
| for k, v in chains.items(): | |
| print(k, {kk: vv for kk, vv in v.items() if not isinstance(vv, list) or kk == "correct_of_4"}) | |
| for r in rl: | |
| print(r) | |
| print(comp, rl_mix) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 10.8 kB
- Xet hash:
- ca30e3e8f7819004c3326d09fe0e1fa2c8b9023d5aff5f91da1261c894688914
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.