| """Build data/{train,val,test,heldout}.jsonl + data/STATS.md from public HF datasets. |
| |
| python -m jevlike.data.build_dataset # full build (seed 13) |
| python -m jevlike.data.build_dataset --seed 7 --max-per-source 8000 |
| python -m jevlike.data.build_dataset --sources ag_news,boolq --out /tmp/x # quick subset |
| |
| Pipeline: per-source converters (jevlike/data/sources_*.py) -> Records (validated, states |
| cleaned + truncated to 2000 chars) -> dedupe -> type-mix balancing (water-filling, so the |
| largest sources are trimmed first) -> split by example key (all questions about one state |
| land in the same split) with val/test quotas proportional to each source -> leak removal |
| (no state shared between splits; heldout states also checked by 8-word shingles) -> JSONL. |
| |
| Deterministic given --seed: every random decision uses random.Random(f"{seed}|..."). |
| Only writes data/{train,val,test,heldout}.jsonl and data/STATS.md (never data/synth, data/bench). |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import os |
| import random |
| import re |
| import sys |
| import time |
| import traceback |
| from dataclasses import dataclass |
|
|
| from jevlike.data import sources_choice, sources_heldout, sources_noul, sources_score |
| from jevlike.data.common import clean, norm_key, truncate |
| from jevlike.data.sources_choice import Ctx |
| from jevlike.types import QTYPES, Record, write_jsonl |
|
|
| IN_DOMAIN = {**sources_choice.SOURCES, **sources_score.SOURCES, **sources_noul.SOURCES} |
| HELDOUT = dict(sources_heldout.SOURCES) |
| PRIMARY = {**{k: "choice" for k in sources_choice.SOURCES}, **{k: "score" for k in sources_score.SOURCES}, |
| **{k: "noul" for k in sources_noul.SOURCES}, "trec": "choice", "pavlick_formality": "score", |
| "enron_spam": "noul"} |
| SPLIT_RANK = {"train": 0, "val": 1, "test": 2, "heldout": 3} |
| _UNFILLED = re.compile(r"\{(x|thing|target)\}") |
|
|
|
|
| @dataclass |
| class Item: |
| rec: Record |
| key: str |
| skey: str |
| split: str = "train" |
|
|
|
|
| def log(*a): |
| print(*a, file=sys.stderr, flush=True) |
|
|
|
|
| |
|
|
| def run_source(name: str, fn, seed: int, budget: int) -> list[Item]: |
| ctx = Ctx(random.Random(f"{seed}|{name}"), name) |
| items: list[Item] = [] |
| n_ex = 0 |
| for ex in fn(ctx): |
| state = truncate(clean(ex.state)) |
| if not state: |
| continue |
| key = ex.key or norm_key(state) |
| skey = norm_key(state) |
| for j, (q, label) in enumerate(ex.qs): |
| q.instructions = q.instructions.strip() |
| assert q.instructions, (name, state[:80]) |
| assert not _UNFILLED.search(q.instructions), (name, q.instructions) |
| rec = Record(f"{name}-{n_ex:06d}-{j}", name, state, q, int(label)) |
| rec.validate() |
| items.append(Item(rec, key, skey)) |
| n_ex += 1 |
| if len(items) >= budget: |
| break |
| return items |
|
|
|
|
| def dedupe(items: list[Item]) -> list[Item]: |
| """Drop exact duplicate questions about the same state; drop both copies if the |
| duplicates disagree on the answer.""" |
| seen: dict[tuple, int] = {} |
| answer: dict[tuple, str] = {} |
| bad: set = set() |
| keep: list[Item] = [] |
| for it in items: |
| q = it.rec.question |
| sig = (it.skey, q.type, q.instructions.lower(), tuple((o.key, o.text) for o in q.options)) |
| ans = q.options[it.rec.label].key if q.options else str(it.rec.label) |
| if sig in seen: |
| if answer[sig] != ans: |
| bad.add(sig) |
| continue |
| seen[sig] = len(keep) |
| answer[sig] = ans |
| keep.append(it) |
| return [it for it in keep if (it.skey, it.rec.question.type, it.rec.question.instructions.lower(), |
| tuple((o.key, o.text) for o in it.rec.question.options)) not in bad] |
|
|
|
|
| def balance_types(items: list[Item], mix: dict[str, float], total: int, seed: int, |
| native_weight: float = 3.0) -> list[Item]: |
| """Trim to `total` records with type shares `mix`. Within a type, every source is capped |
| at the same level (water-filling, so the largest sources are trimmed first), except that |
| sources whose native task is that type get a `native_weight`x higher cap (e.g. boolq / |
| NLI statements outweigh noul statements derived from classification labels).""" |
| by_t: dict[str, list[Item]] = {t: [] for t in QTYPES} |
| for it in items: |
| by_t[it.rec.question.type].append(it) |
| feasible = min(len(by_t[t]) / mix[t] for t in QTYPES if mix[t] > 0) |
| grand = min(total, int(feasible)) |
| out: list[Item] = [] |
| for t in QTYPES: |
| target = int(round(mix[t] * grand)) |
| per_src: dict[str, list[Item]] = {} |
| for it in by_t[t]: |
| per_src.setdefault(it.rec.source, []).append(it) |
| srcs = sorted(per_src) |
| weights = [native_weight if PRIMARY.get(s) == t else 1.0 for s in srcs] |
| level = _water_level([len(per_src[s]) for s in srcs], weights, target) |
| for src, wgt in zip(srcs, weights): |
| v = per_src[src] |
| cap = int(level * wgt) |
| if len(v) > cap: |
| rng = random.Random(f"{seed}|balance|{src}|{t}") |
| idx = sorted(rng.sample(range(len(v)), cap)) |
| v = [v[i] for i in idx] |
| out.extend(v) |
| return out |
|
|
|
|
| def _water_level(sizes: list[int], weights: list[float], target: int) -> float: |
| """Largest level c with sum(min(size_i, int(c * w_i))) <= target.""" |
| def used(c): |
| return sum(min(n, int(c * w)) for n, w in zip(sizes, weights)) |
| hi = max((n / w for n, w in zip(sizes, weights)), default=0.0) |
| if used(hi) <= target: |
| return hi |
| lo = 0.0 |
| for _ in range(60): |
| mid = (lo + hi) / 2 |
| if used(mid) <= target: |
| lo = mid |
| else: |
| hi = mid |
| return lo |
|
|
|
|
| |
|
|
| def assign_splits(items: list[Item], val_size: int, test_size: int, seed: int) -> None: |
| by_src: dict[str, list[Item]] = {} |
| for it in items: |
| by_src.setdefault(it.rec.source, []).append(it) |
| n_total = len(items) |
| for src in sorted(by_src): |
| its = by_src[src] |
| groups: dict[str, list[Item]] = {} |
| for it in its: |
| groups.setdefault(it.key, []).append(it) |
| keys = list(groups) |
| random.Random(f"{seed}|split|{src}").shuffle(keys) |
| q_test = round(test_size * len(its) / n_total) |
| q_val = round(val_size * len(its) / n_total) |
| n_test = n_val = 0 |
| for k in keys: |
| if n_test < q_test: |
| split = "test" |
| n_test += len(groups[k]) |
| elif n_val < q_val: |
| split = "val" |
| n_val += len(groups[k]) |
| else: |
| split = "train" |
| for it in groups[k]: |
| it.split = split |
|
|
|
|
| def remove_leaks(items: list[Item], heldout: list[Item], shingle_n: int = 8) -> tuple[list[Item], dict]: |
| """No example key / normalized state may occur in two splits: keep it in the |
| highest-ranked split (heldout > test > val > train) and drop the other copies. |
| In-domain states that share an 8-word shingle with a heldout state are dropped too.""" |
| rank: dict[str, int] = {} |
| for it in items + heldout: |
| r = SPLIT_RANK[it.split] |
| for k in (it.key, it.skey): |
| rank[k] = max(rank.get(k, 0), r) |
| held_sh: set = set() |
| for it in heldout: |
| held_sh.update(_shingles(it.skey, shingle_n)) |
| dropped = {"cross_split": 0, "heldout_shingle": 0} |
| keep = [] |
| for it in items: |
| r = SPLIT_RANK[it.split] |
| if rank[it.key] > r or rank[it.skey] > r: |
| dropped["cross_split"] += 1 |
| continue |
| if held_sh and not held_sh.isdisjoint(_shingles(it.skey, shingle_n)): |
| dropped["heldout_shingle"] += 1 |
| continue |
| keep.append(it) |
| return keep, dropped |
|
|
|
|
| def _shingles(skey: str, n: int) -> list[int]: |
| toks = skey.split() |
| if len(toks) < n: |
| return [_h(skey)] if len(toks) >= 4 else [] |
| return [_h(" ".join(toks[i:i + n])) for i in range(len(toks) - n + 1)] |
|
|
|
|
| def _h(s: str) -> int: |
| return int.from_bytes(hashlib.blake2b(s.encode(), digest_size=8).digest(), "little") |
|
|
|
|
| |
|
|
| def main(argv=None) -> int: |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("--out", default="data", help="output directory (default: data)") |
| ap.add_argument("--seed", type=int, default=13) |
| ap.add_argument("--max-per-source", type=int, default=15000, help="record budget per in-domain source") |
| ap.add_argument("--total", type=int, default=215000, help="max in-domain records (train+val+test)") |
| ap.add_argument("--mix", default="0.45,0.20,0.35", help="choice,score,noul shares") |
| ap.add_argument("--val-size", type=int, default=5000) |
| ap.add_argument("--test-size", type=int, default=5000) |
| ap.add_argument("--heldout-per-source", type=int, default=1000) |
| ap.add_argument("--sources", default="", help="comma-separated subset (debug)") |
| ap.add_argument("--list", action="store_true", help="list sources and exit") |
| args = ap.parse_args(argv) |
|
|
| if args.list: |
| for n in IN_DOMAIN: |
| print(f"{n:24s} {PRIMARY[n]}") |
| for n in HELDOUT: |
| print(f"{n:24s} {PRIMARY[n]} (heldout)") |
| return 0 |
|
|
| mix = dict(zip(QTYPES, (float(x) for x in args.mix.split(",")))) |
| only = [s for s in args.sources.split(",") if s] |
| t0 = time.time() |
| items: list[Item] = [] |
| status: dict[str, str] = {} |
| for name, fn in IN_DOMAIN.items(): |
| if only and name not in only: |
| continue |
| try: |
| got = run_source(name, fn, args.seed, args.max_per_source) |
| items.extend(got) |
| status[name] = f"ok ({len(got)} generated)" |
| log(f"[{time.time() - t0:6.0f}s] {name}: {len(got)}") |
| except Exception as e: |
| status[name] = f"FAILED: {type(e).__name__}: {str(e)[:200]}" |
| log(f"[{time.time() - t0:6.0f}s] {name}: FAILED\n{traceback.format_exc()}") |
| held: list[Item] = [] |
| for name, fn in HELDOUT.items(): |
| if only and name not in only: |
| continue |
| try: |
| got = dedupe(run_source(name, fn, args.seed, args.heldout_per_source)) |
| for it in got: |
| it.split = "heldout" |
| held.extend(got) |
| status[name] = f"ok ({len(got)} generated, heldout)" |
| log(f"[{time.time() - t0:6.0f}s] {name} (heldout): {len(got)}") |
| except Exception as e: |
| status[name] = f"FAILED: {type(e).__name__}: {str(e)[:200]}" |
| log(f"[{time.time() - t0:6.0f}s] {name}: FAILED\n{traceback.format_exc()}") |
|
|
| n_gen = len(items) |
| items = dedupe(items) |
| n_dedup = len(items) |
| items = balance_types(items, mix, args.total, args.seed) |
| n_bal = len(items) |
| assign_splits(items, args.val_size, args.test_size, args.seed) |
| items, dropped = remove_leaks(items, held) |
| log(f"generated {n_gen}, after dedupe {n_dedup}, after type balance {n_bal}, after leak removal {len(items)} " |
| f"(dropped {dropped})") |
|
|
| os.makedirs(args.out, exist_ok=True) |
| splits = {s: [it for it in items if it.split == s] for s in ("train", "val", "test")} |
| splits["heldout"] = held |
| for s, its in splits.items(): |
| recs = [it.rec for it in its] |
| random.Random(f"{args.seed}|shuffle|{s}").shuffle(recs) |
| write_jsonl(os.path.join(args.out, f"{s}.jsonl"), recs) |
| log(f"wrote {s}: {len(recs)}") |
|
|
| from jevlike.data.stats import write_stats |
| info = {"seed": args.seed, "argv": " ".join(argv if argv is not None else sys.argv[1:]), |
| "generated": n_gen, "deduped": n_dedup, "balanced": n_bal, "dropped": dropped, |
| "status": status, "primary": PRIMARY, "mix": mix} |
| write_stats(os.path.join(args.out, "STATS.md"), {s: [it.rec for it in its] for s, its in splits.items()}, info) |
| log(f"done in {time.time() - t0:.0f}s") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
