Buckets:
| """V2 data: a ~45K-question readout pool (a superset of v1's), a generation order for new chains, and a dev slice. | |
| Writes into --out (a variant directory's data/): | |
| train_pool.jsonl every V2 readout question; items with chains come first in v1's order, then the rest shuffled. | |
| Generation walks this file in order and skips items already in gen.jsonl. | |
| dev.jsonl 1,000 training-source questions never trained on (learning-rate sweep only) | |
| summary_v2.json counts and exclusions | |
| Rules as in v1: train splits only, eval/calib content and ids excluded, whole-item exclusions counted. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import collections | |
| import json | |
| import os | |
| import random | |
| import numpy as np | |
| import pandas as pd | |
| from transformers import AutoTokenizer | |
| from common import MAX_PROMPT_TOKENS, read_jsonl, render, write_jsonl | |
| from prep_data import fetch, finish, key, noul_options, typed_question, PROC, TS | |
| SEED = 23 | |
| def soft_source(df_rows) -> bool: | |
| """A tasksource source whose targets are human vote shares (not one-hot).""" | |
| t = df_rows["target"] | |
| soft = t.map(lambda x: (len(x) > 1 and max(x) < 0.999) or (len(x) == 1 and 0.001 < x[0] < 0.999)) | |
| return soft.mean() > 0.5 | |
| def main() -> None: | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--v1", default="data") | |
| ap.add_argument("--out", required=True) | |
| ap.add_argument("--raw", default="data/raw") | |
| ap.add_argument("--model", default="Qwen/Qwen3.5-4B") | |
| ap.add_argument("--proc-total", type=int, default=18000) | |
| ap.add_argument("--ts-total", type=int, default=24000) | |
| ap.add_argument("--ts-cap", type=int, default=40) | |
| ap.add_argument("--ts-soft-cap", type=int, default=150) | |
| ap.add_argument("--dev", type=int, default=1000) | |
| args = ap.parse_args() | |
| os.makedirs(args.out, exist_ok=True) | |
| rng = random.Random(SEED) | |
| counts = collections.Counter() | |
| v1_pool = read_jsonl(os.path.join(args.v1, "train_pool.jsonl")) | |
| v1_gen_ids = [r["id"] for r in read_jsonl(os.path.join(args.v1, "gen.jsonl"))] | |
| held = {key(x) for x in read_jsonl(os.path.join(args.v1, "eval.jsonl")) + read_jsonl(os.path.join(args.v1, "calib.jsonl"))} | |
| held_ids = {x["id"] for x in read_jsonl(os.path.join(args.v1, "eval.jsonl")) + read_jsonl(os.path.join(args.v1, "calib.jsonl"))} | |
| have_ids = {x["id"] for x in v1_pool} | |
| have_keys = {key(x) for x in v1_pool} | |
| # ---------------- procedural: top up each family to proc_total / 18, one question per new state | |
| proc_paths = fetch(PROC, "all/train", args.raw) | |
| df = pd.concat([pd.read_parquet(p) for p in proc_paths], ignore_index=True) | |
| per_family_target = args.proc_total // df["task"].nunique() | |
| have_fam = collections.Counter(x["family"] for x in v1_pool if x["set"] == "train_proc") | |
| used_states = {x["id"].split(":")[1] + ":" + x["id"].split(":")[2] + ":" + x["id"].split(":")[3] for x in v1_pool if x["set"] == "train_proc"} | |
| new_proc, dev_proc = [], [] | |
| for task, g in df.groupby("task", sort=True): | |
| g = g.sample(frac=1.0, random_state=rng.randrange(1 << 30)) | |
| need, need_dev = per_family_target - have_fam[task], args.dev // 2 // df["task"].nunique() + 1 | |
| for _, row in g.iterrows(): | |
| if need <= 0 and need_dev <= 0: | |
| break | |
| if row["id"] in used_states: | |
| continue | |
| qs, ans = json.loads(row["questions"]), json.loads(row["answers"]) | |
| names = list(qs) | |
| rng.shuffle(names) | |
| for name in names: | |
| opts, target = typed_question(name, qs[name], ans[name]) | |
| item = finish({"id": f"proc:{row['id']}:{name}", "set": "train_proc", "family": task, "level": int(row["level"]), | |
| "tier": None, "qtype": qs[name]["type"], "question_name": name, "evidence": row["state"], | |
| "criterion": qs[name]["instructions"], "options": opts, "target": target}, counts, "proc") | |
| if item is None or key(item) in held or item["id"] in have_ids: | |
| continue | |
| if need_dev > 0: | |
| item["set"] = "dev_proc" | |
| dev_proc.append(item) | |
| need_dev -= 1 | |
| elif need > 0: | |
| new_proc.append(item) | |
| need -= 1 | |
| break | |
| # ---------------- tasksource: raise the per-source cap (soft-label sources higher) | |
| ts_paths = fetch(TS, "data/train", args.raw) | |
| tdf = pd.concat([pd.read_parquet(p, columns=["id", "state", "kind", "options", "target", "question", "source"]) for p in ts_paths], ignore_index=True) | |
| tdf = tdf[~tdf["source"].str.startswith("procedural-typed-decisions/")] | |
| soft = {s for s, g in tdf.groupby("source") if soft_source(g.head(200))} | |
| have_src = collections.Counter(x["family"] for x in v1_pool if x["set"] == "train_ts") | |
| tdf = tdf.sample(frac=1.0, random_state=rng.randrange(1 << 30)) | |
| new_ts, dev_ts, per = [], [], collections.Counter(have_src) | |
| need = args.ts_total - sum(have_src.values()) | |
| need_dev = args.dev - len(dev_proc) | |
| for row in tdf.itertuples(index=False): | |
| if need <= 0 and need_dev <= 0: | |
| break | |
| if f"ts:{row.id}" in have_ids: | |
| continue | |
| cap = args.ts_soft_cap if row.source in soft else args.ts_cap | |
| to_dev = need_dev > 0 and rng.random() < 0.05 | |
| if not to_dev and per[row.source] >= cap: | |
| continue | |
| if row.kind == "noul": | |
| p = float(row.target[0]) | |
| opts, target = noul_options(None), [p, 1 - p] | |
| else: | |
| opts = ([{"id": str(i), "desc": o} for i, o in enumerate(row.options)] if row.kind == "score" | |
| else [{"id": None, "desc": o} for o in row.options]) | |
| target = list(row.target) | |
| item = finish({"id": f"ts:{row.id}", "set": "train_ts", "family": row.source, "level": None, "tier": None, | |
| "qtype": row.kind, "evidence": row.state, "criterion": row.question, "options": opts, "target": target}, counts, "ts") | |
| if item is None or key(item) in held or key(item) in have_keys: | |
| counts["ts:excluded_overlap_or_invalid"] += item is not None | |
| continue | |
| if to_dev: | |
| item["set"] = "dev_ts" | |
| dev_ts.append(item) | |
| need_dev -= 1 | |
| continue | |
| new_ts.append(item) | |
| per[row.source] += 1 | |
| need -= 1 | |
| # ---------------- length check, order, write | |
| tok = AutoTokenizer.from_pretrained(args.model) | |
| def fits(rows, name): | |
| lens = [len(t) for t in tok([render(tok, x, think=False) for x in rows], add_special_tokens=False)["input_ids"]] | |
| out = [] | |
| for x, n in zip(rows, lens): | |
| if n > MAX_PROMPT_TOKENS: | |
| counts[f"{name}:prompt_over_{MAX_PROMPT_TOKENS}_tokens"] += 1 | |
| continue | |
| x["prompt_tokens"] = n | |
| out.append(x) | |
| return out | |
| new_items = fits(new_proc + new_ts, "train_new") | |
| dev = fits(dev_proc + dev_ts, "dev") | |
| by_id = {x["id"]: x for x in v1_pool} | |
| generated = [by_id[i] for i in v1_gen_ids if i in by_id] | |
| rest_v1 = [x for x in v1_pool if x["id"] not in set(v1_gen_ids)] | |
| rest = rest_v1 + new_items | |
| rng.shuffle(rest) | |
| pool = generated + rest # chains exist for the head; generation continues from there | |
| ids = [x["id"] for x in pool] | |
| assert len(ids) == len(set(ids)) and not (set(ids) & held_ids) and not ({x["id"] for x in dev} & set(ids)) | |
| write_jsonl(os.path.join(args.out, "train_pool.jsonl"), pool) | |
| write_jsonl(os.path.join(args.out, "dev.jsonl"), dev) | |
| by = lambda rows, f: dict(collections.Counter(f(x) for x in rows)) | |
| summ = {"pool": len(pool), "with_v1_chains": len(generated), "by_set": by(pool, lambda x: x["set"]), "by_type": by(pool, lambda x: x["qtype"]), | |
| "ts_sources": len({x["family"] for x in pool if x["set"] == "train_ts"}), "ts_soft_sources": len(soft), | |
| "ts_soft_items": sum(1 for x in pool if x["set"] == "train_ts" and x["family"] in soft), | |
| "dev": len(dev), "dev_by_set": by(dev, lambda x: x["set"]), "exclusions": dict(sorted(counts.items()))} | |
| json.dump(summ, open(os.path.join(args.out, "summary_v2.json"), "w"), indent=1) | |
| print(json.dumps(summ, indent=1)) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 8.43 kB
- Xet hash:
- 40707f7c67e8ffae6001f1cacb5b152950982e8c710db0a7e45b79889585105e
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.