claude-opus-5

File created successfully at: /workspace/score_pool_tok.py (file state is current in your context — no need to Read it back)

[metadata] {"type": "create", "filePath": "/workspace/score_pool_tok.py", "content": "\"\"\"Stage 1b: the same per-register importance score as score_pool.py, but computed over\nGPT-2 *token* uni/bi-grams instead of lowercased word n-grams.\n\nRationale: the model is scored by next-token cross-entropy over GPT-2 BPE tokens, so the\nunits that matter for perplexity are exactly those tokens -- including casing, punctuation\nand whitespace conventions that a word-level featuriser throws away.  Scoring in token\nspace therefore matches training-data selection to the objective being evaluated.\n\nPer-feature log-ratios are clipped to +-CLIP so no single idiosyncratic token (e.g. a rare\nmarkup artefact) can dominate a document's score.\n\nSide effect: writes an exact GPT-2 token count for every pool document (work/tokcache.pkl,\nwork/tokcount.json), which lets the curator do exact budget accounting.\n\nRun: python score_pool_tok.py  ->  work/scores_tok.npz\n\"\"\"\nimport json, os, pickle\nos.environ[\"TOKENIZERS_PARALLELISM\"] = \"true\"\nimport numpy as np\n\nPOOL = \"/workspace/data/pool.jsonl\"\nDEV = \"/workspace/data/multi_dev.npy\"\nOUT = \"/workspace/work/scores_tok.npz\"\nCACHE = \"/workspace/work/tokcache.pkl\"\nTOKCNT = \"/workspace/work/tokcount.json\"\nV = 50257\nNB = 1 << 18\nNREG = 4\nALPHA = 0.5\nCLIP = 4.0\nBG_STRIDE = 11\n\n\ndef bihash(arr):\n    a = arr[:-1].astype(np.int64)\n    b = arr[1:].astype(np.int64)\n    return ((a * 50261) ^ (b * 97)) & (NB - 1)\n\n\ndef tables(arr):\n    uni = np.bincount(arr.astype(np.int64), minlength=V).astype(np.float64)\n    bic = np.bincount(bihash(arr), minlength=NB).astype(np.float64)\n    return uni, bic\n\n\ndef main():\n    from transformers import AutoTokenizer\n    tok = AutoTokenizer.from_pretrained(\"gpt2\")\n\n    ids, texts = [], []\n    for line in open(POOL):\n        r = json.loads(line)\n        ids.append(r[\"id\"]); texts.append(r[\"text\"])\n    print(f\"tokenising {len(texts)} pool docs ...\", flush=True)\n    toks = {}\n    B = 512\n    for s in range(0, len(texts), B):\n        enc = tok(texts[s:s + B], add_special_tokens=False).input_ids\n        for i, e in zip(ids[s:s + B], enc):\n            toks[i] = np.array(e, dtype=np.uint16)\n        if s % 51200 == 0:\n            print(f\"  {s}/{len(texts)}\", flush=True)\n    del texts\n    pickle.dump(toks, open(CACHE, \"wb\"), protocol=4)\n    json.dump({str(k): int(len(v)) for k, v in toks.items()}, open(TOKCNT, \"w\"))\n\n    # ---- background (1/BG_STRIDE sample of docs) and per-register reference tables\n    print(\"building tables ...\", flush=True)\n    bg_u = np.zeros(V); bg_b = np.zeros(NB)\n    for j, i in enumerate(ids):\n        if j % BG_STRIDE == 0 and len(toks[i]) > 1:\n            u, b = tables(toks[i]); bg_u += u; bg_b += b\n    d = np.load(DEV); q = len(d) // NREG\n    lut_u = np.zeros((NREG, V), dtype=np.float32)\n    lut_b = np.zeros((NREG, NB), dtype=np.float32)\n    lqu = np.log((bg_u + ALPHA) / (bg_u.sum() + ALPHA * V))\n    lqb = np.log((bg_b + ALPHA) / (bg_b.sum() + ALPHA * NB))\n    for k in range(NREG):\n        u, b = tables(d[k * q:(k + 1) * q])\n        lpu = np.log((u + ALPHA) / (u.sum() + ALPHA * V))\n        lpb = np.log((b + ALPHA) / (b.sum() + ALPHA * NB))\n        lut_u[k] = np.clip(lpu - lqu, -CLIP, CLIP).astype(np.float32)\n        lut_b[k] = np.clip(lpb - lqb, -CLIP, CLIP).astype(np.float32)\n\n    # ---- score every doc: mean clipped log-ratio over its uni- and bi-grams\n    print(\"scoring ...\", flush=True)\n    S = np.zeros((NREG, len(ids)), dtype=np.float32)\n    NT = np.zeros(len(ids), dtype=np.int64)\n    for j, i in enumerate(ids):\n        t = toks[i]\n        NT[j] = len(t)\n        if len(t) < 2:\n            continue\n        ti = t.astype(np.int64); bh = bihash(t)\n        for k in range(NREG):\n            S[k, j] = 0.5 * (lut_u[k][ti].mean() + lut_b[k][bh].mean())\n    np.savez(OUT, id=np.array(ids, dtype=np.int64), ntok=NT,\n             **{f\"s{k}\": S[k] for k in range(NREG)})\n    print(\"wrote\", OUT, flush=True)\n\n\nif __name__ == \"__main__\":\n    main()\n", "structuredPatch": [], "originalFile": null, "userModified": false}

tokenising 182016 pool docs ...
  0/182016
  51200/182016
  102400/182016
  153600/182016
building tables ...
scoring ...
wrote /workspace/work/scores_tok.npz

real	1m52.666s
user	12m5.480s
sys	0m6.590s
4399

[stdout]
tokenising 182016 pool docs ...
  0/182016
  51200/182016
  102400/182016
  153600/182016
building tables ...
scoring ...
wrote /workspace/work/scores_tok.npz

real	1m52.666s
user	12m5.480s
sys	0m6.590s
4399
[metadata] {"noOutputExpected": false}

The file /workspace/curate.py has been updated successfully. (file state is current in your context — no need to Read it back)

[metadata] {"filePath": "/workspace/curate.py", "oldString": "ap.add_argument(\"--balanced\", type=int, default=1)    # 1 = per-register quotas, 0 = global rank\na = ap.parse_args()\n\nif not os.path.exists(SCORES):\n    subprocess.check_call([sys.executable, \"/workspace/score_pool.py\"])\nz = np.load(SCORES)\nids, nw, nch = z[\"id\"], z[\"nw\"], z[\"nch\"]", "newString": "ap.add_argument(\"--balanced\", type=int, default=1)    # 1 = per-register quotas, 0 = global rank\nap.add_argument(\"--scores\", default=\"tok\", choices=[\"tok\", \"word\", \"both\"])\nap.add_argument(\"--share\", default=\"\")                # e.g. \"0.3,0.2,0.2,0.3\" register token shares\na = ap.parse_args()\n\nif not os.path.exists(SCORES):\n    subprocess.check_call([sys.executable, \"/workspace/score_pool.py\"])\nz = np.load(SCORES)\nids, nw, nch = z[\"id\"], z[\"nw\"], z[\"nch\"]\nSW = np.stack([z[f\"s{k}\"] for k in range(NREG)])       # word-level register scores\nif a.scores in (\"tok\", \"both\"):\n    TOKSC = \"/workspace/work/scores_tok.npz\"\n    if not os.path.exists(TOKSC):\n        subprocess.check_call([sys.executable, \"/workspace/score_pool_tok.py\"])\n    zt = np.load(TOKSC)\n    assert (zt[\"id\"] == ids).all()\n    ST = np.stack([zt[f\"s{k}\"] for k in range(NREG)])  # GPT-2-token-level register scores", "originalFile": "\"\"\"Curate /workspace/data/pool.jsonl -> /workspace/submission/selection.json (priority order).\n\nSTATED CRITERION\n================\nThe eval target is an equal-parts mixture of four registers (encyclopedic, high-quality\nweb prose, news, technical Q&A).  We therefore build one *reference distribution per\nregister* from the disclosed dev target and select, for each register, the pool documents\nthat are most likely under it relative to the pool's own background -- then interleave the\nfour ranked lists so that any prefix of the submission (in particular the 12M-token prefix\nthe trainer consumes) is register-balanced.\n\n  1. GATE   Drop documents failing surface-quality checks (Gopher/C4-style): too short,\n            implausible mean word length, duplicated lines, digit/symbol/non-ASCII heavy,\n            single-token domination, navigation boilerplate.  Junk web text is never\n            worth budget, whatever its topic.\n  2. SCORE  s_k(d) = per-token log-likelihood ratio log p_k / q over hashed word uni+bi-grams\n            (DSIR, Xie et al. 2023), p_k = register-k n-grams from multi_dev.npy quarter k,\n            q = pool background n-grams.  Shrunk toward 0 for short docs by n/(n+n0) so a\n            handful of lucky n-grams cannot outrank a long, consistently on-target document.\n  3. ASSIGN Each surviving doc goes to its best-matching register (argmax of the\n            per-register standardised score), so the four lists are disjoint.\n  4. DEDUP  Exact (normalised-text hash) and near-duplicate (min-of-word-8-gram-hashes\n            sketch) collapse; the highest-scoring representative survives.\n  5. ORDER  Greedy balanced interleave: repeatedly emit the next-best document of whichever\n            register currently has the fewest tokens emitted.  Equal-parts target => equal\n            token shares, and the balance holds at *every* prefix length.\n\nUsage:  PYTHONHASHSEED=0 python curate.py [--out selection.json] [--minw 128] [--n0 64]\nRequires work/scores.npz (built automatically by score_pool.py if absent);\nwork/tokcount.json (exact GPT-2 token counts for candidates) is used for the token\naccounting when present, otherwise len(chars)/4 is used as the estimate.\n\"\"\"\nimport argparse, json, os, subprocess, sys\nimport numpy as np\n\nSCORES = \"/workspace/work/scores.npz\"\nTOKCNT = \"/workspace/work/tokcount.json\"\nNREG = 4\nBUDGET = 12_000_000\n\nap = argparse.ArgumentParser()\nap.add_argument(\"--out\", default=\"/workspace/submission/selection.json\")\nap.add_argument(\"--minw\", type=int, default=128)      # min words per doc\nap.add_argument(\"--maxw\", type=int, default=12000)\nap.add_argument(\"--n0\", type=int, default=64)         # length-shrinkage constant\nap.add_argument(\"--emit_tokens\", type=int, default=48_000_000)  # how much to list (>> budget)\nap.add_argument(\"--balanced\", type=int, default=1)    # 1 = per-register quotas, 0 = global rank\na = ap.parse_args()\n\nif not os.path.exists(SCORES):\n    subprocess.check_call([sys.executable, \"/workspace/score_pool.py\"])\nz = np.load(SCORES)\nids, nw, nch = z[\"id\"], z[\"nw\"], z[\"nch\"]\n\n# ---------------------------------------------------------------- 1. quality gate\ngate = ((nw >= a.minw) & (nw <= a.maxw) & (z[\"wlen\"] >= 3.0) & (z[\"wlen\"] <= 9.0)\n        & (z[\"dupl\"] <= 0.25) & (z[\"digit\"] <= 0.20) & (z[\"nonascii\"] <= 0.05)\n        & (z[\"toprep\"] <= 0.16) & (z[\"symb\"] <= 0.02) & (z[\"boiler\"] <= 2))\n\n# ---------------------------------------------------------------- 2. scores\nshrink = nw / (nw + a.n0)\nS = np.stack([z[f\"s{k}\"] * shrink for k in range(NREG)])          # (NREG, N)\nSz = (S - S.mean(axis=1, keepdims=True)) / S.std(axis=1, keepdims=True)\n\n# exact token counts when available, else 4 chars/token\ntokc = np.maximum(1, nch // 4).astype(np.int64)\nif os.path.exists(TOKCNT):\n    d = json.load(open(TOKCNT))\n    pos = {int(v): i for i, v in enumerate(ids)}\n    for k, v in d.items():\n        i = pos.get(int(k))\n        if i is not None:\n            tokc[i] = v\n\n# ---------------------------------------------------------------- 3. assign + 4. dedup\nbest = Sz.argmax(axis=0)\ncand = np.where(gate)[0]\nlists = [[] for _ in range(NREG)]\nfor k in range(NREG):\n    sel = cand[best[cand] == k]\n    sel = sel[np.argsort(-S[k][sel])]                            # best-first\n    lists[k] = sel\n\nseen_ex, seen_nd = set(), set()\nex, nd = z[\"exhash\"], z[\"ndhash\"]\nded = [[] for _ in range(NREG)]\norder = np.concatenate([np.stack([np.full(len(l), k), l]) for k, l in enumerate(lists)], axis=1)\nrank = np.argsort(-np.array([Sz[k, i] for k, i in order.T]))      # dedup globally, best-first\nkeep = np.zeros(len(ids), dtype=bool)\nfor r in rank:\n    k, i = order[0, r], order[1, r]\n    e, n = int(ex[i]), int(nd[i])\n    if e in seen_ex or n in seen_nd:\n        continue\n    seen_ex.add(e); seen_nd.add(n); keep[i] = True\nfor k in range(NREG):\n    ded[k] = [i for i in lists[k] if keep[i]]\n\n# ---------------------------------------------------------------- 5. balanced interleave\nout, ptr, tok = [], [0] * NREG, [0] * NREG\nif a.balanced:\n    total = 0\n    while total < a.emit_tokens:\n        live = [k for k in range(NREG) if ptr[k] < len(ded[k])]\n        if not live:\n            break\n        k = min(live, key=lambda k: tok[k])\n        i = ded[k][ptr[k]]; ptr[k] += 1\n        out.append(int(ids[i])); tok[k] += int(tokc[i]); total += int(tokc[i])\nelse:  # ablation: one global ranking by best register score, no balancing\n    flat = np.concatenate(ded)\n    flat = flat[np.argsort(-Sz.max(axis=0)[flat])]\n    total = 0\n    for i in flat:\n        out.append(int(ids[i])); total += int(tokc[i])\n        tok[int(best[i])] += int(tokc[i])\n        if total >= a.emit_tokens:\n            break\n\nassert len(out) == len(set(out)), \"duplicate ids\"\nos.makedirs(os.path.dirname(a.out), exist_ok=True)\njson.dump(out, open(a.out, \"w\"))\n\n# report the composition of the prefix the trainer will actually use\ncum, comp, nd_ = 0, [0] * NREG, 0\npos = {int(v): i for i, v in enumerate(ids)}\nfor did in out:\n    i = pos[did]\n    comp[int(best[i])] += int(tokc[i]); cum += int(tokc[i]); nd_ += 1\n    if cum >= BUDGET:\n        break\nprint(f\"gate kept {gate.sum()}/{len(ids)}; dedup kept {keep.sum()}\")\nprint(f\"emitted {len(out)} ids (~{total/1e6:.1f}M tok); first {BUDGET/1e6:.0f}M tok = {nd_} docs\")\nprint(\"  register token mix in budget prefix:\", [round(c / max(1, cum), 3) for c in comp])\nprint(\"  wrote\", a.out)\n", "structuredPatch": [{"oldStart": 45, "oldLines": 12, "newStart": 45, "newLines": 22, "lines": [" ap.add_argument(\"--n0\", type=int, default=64)         # length-shrinkage constant", " ap.add_argument(\"--emit_tokens\", type=int, default=48_000_000)  # how much to list (>> budget)", " ap.add_argument(\"--balanced\", type=int, default=1)    # 1 = per-register quotas, 0 = global rank", "+ap.add_argument(\"--scores\", default=\"tok\", choices=[\"tok\", \"word\", \"both\"])", "+ap.add_argument(\"--share\", default=\"\")                # e.g. \"0.3,0.2,0.2,0.3\" register token shares", " a = ap.parse_args()", " ", " if not os.path.exists(SCORES):", "     subprocess.check_call([sys.executable, \"/workspace/score_pool.py\"])", " z = np.load(SCORES)", " ids, nw, nch = z[\"id\"], z[\"nw\"], z[\"nch\"]", "+SW = np.stack([z[f\"s{k}\"] for k in range(NREG)])       # word-level register scores", "+if a.scores in (\"tok\", \"both\"):", "+    TOKSC = \"/workspace/work/scores_tok.npz\"", "+    if not os.path.exists(TOKSC):", "+        subprocess.check_call([sys.executable, \"/workspace/score_pool_tok.py\"])", "+    zt = np.load(TOKSC)", "+    assert (zt[\"id\"] == ids).all()", "+    ST = np.stack([zt[f\"s{k}\"] for k in range(NREG)])  # GPT-2-token-level register scores", " ", " # ---------------------------------------------------------------- 1. quality gate", " gate = ((nw >= a.minw) & (nw <= a.maxw) & (z[\"wlen\"] >= 3.0) & (z[\"wlen\"] <= 9.0)"]}], "userModified": false, "replaceAll": false}

The file /workspace/curate.py has been updated successfully. (file state is current in your context — no need to Read it back)

[metadata] {"filePath": "/workspace/curate.py", "oldString": "# ---------------------------------------------------------------- 2. scores\nshrink = nw / (nw + a.n0)\nS = np.stack([z[f\"s{k}\"] * shrink for k in range(NREG)])          # (NREG, N)\nSz = (S - S.mean(axis=1, keepdims=True)) / S.std(axis=1, keepdims=True)", "newString": "# ---------------------------------------------------------------- 2. scores\ndef zsc(M):\n    return (M - M.mean(axis=1, keepdims=True)) / M.std(axis=1, keepdims=True)\n\nshrink = nw / (nw + a.n0)\nif a.scores == \"word\":\n    S = SW * shrink\nelif a.scores == \"tok\":\n    S = ST * shrink\nelse:                                    # average the two standardised views\n    S = 0.5 * (zsc(SW) + zsc(ST)) * shrink\nSz = zsc(S)", "originalFile": "\"\"\"Curate /workspace/data/pool.jsonl -> /workspace/submission/selection.json (priority order).\n\nSTATED CRITERION\n================\nThe eval target is an equal-parts mixture of four registers (encyclopedic, high-quality\nweb prose, news, technical Q&A).  We therefore build one *reference distribution per\nregister* from the disclosed dev target and select, for each register, the pool documents\nthat are most likely under it relative to the pool's own background -- then interleave the\nfour ranked lists so that any prefix of the submission (in particular the 12M-token prefix\nthe trainer consumes) is register-balanced.\n\n  1. GATE   Drop documents failing surface-quality checks (Gopher/C4-style): too short,\n            implausible mean word length, duplicated lines, digit/symbol/non-ASCII heavy,\n            single-token domination, navigation boilerplate.  Junk web text is never\n            worth budget, whatever its topic.\n  2. SCORE  s_k(d) = per-token log-likelihood ratio log p_k / q over hashed word uni+bi-grams\n            (DSIR, Xie et al. 2023), p_k = register-k n-grams from multi_dev.npy quarter k,\n            q = pool background n-grams.  Shrunk toward 0 for short docs by n/(n+n0) so a\n            handful of lucky n-grams cannot outrank a long, consistently on-target document.\n  3. ASSIGN Each surviving doc goes to its best-matching register (argmax of the\n            per-register standardised score), so the four lists are disjoint.\n  4. DEDUP  Exact (normalised-text hash) and near-duplicate (min-of-word-8-gram-hashes\n            sketch) collapse; the highest-scoring representative survives.\n  5. ORDER  Greedy balanced interleave: repeatedly emit the next-best document of whichever\n            register currently has the fewest tokens emitted.  Equal-parts target => equal\n            token shares, and the balance holds at *every* prefix length.\n\nUsage:  PYTHONHASHSEED=0 python curate.py [--out selection.json] [--minw 128] [--n0 64]\nRequires work/scores.npz (built automatically by score_pool.py if absent);\nwork/tokcount.json (exact GPT-2 token counts for candidates) is used for the token\naccounting when present, otherwise len(chars)/4 is used as the estimate.\n\"\"\"\nimport argparse, json, os, subprocess, sys\nimport numpy as np\n\nSCORES = \"/workspace/work/scores.npz\"\nTOKCNT = \"/workspace/work/tokcount.json\"\nNREG = 4\nBUDGET = 12_000_000\n\nap = argparse.ArgumentParser()\nap.add_argument(\"--out\", default=\"/workspace/submission/selection.json\")\nap.add_argument(\"--minw\", type=int, default=128)      # min words per doc\nap.add_argument(\"--maxw\", type=int, default=12000)\nap.add_argument(\"--n0\", type=int, default=64)         # length-shrinkage constant\nap.add_argument(\"--emit_tokens\", type=int, default=48_000_000)  # how much to list (>> budget)\nap.add_argument(\"--balanced\", type=int, default=1)    # 1 = per-register quotas, 0 = global rank\nap.add_argument(\"--scores\", default=\"tok\", choices=[\"tok\", \"word\", \"both\"])\nap.add_argument(\"--share\", default=\"\")                # e.g. \"0.3,0.2,0.2,0.3\" register token shares\na = ap.parse_args()\n\nif not os.path.exists(SCORES):\n    subprocess.check_call([sys.executable, \"/workspace/score_pool.py\"])\nz = np.load(SCORES)\nids, nw, nch = z[\"id\"], z[\"nw\"], z[\"nch\"]\nSW = np.stack([z[f\"s{k}\"] for k in range(NREG)])       # word-level register scores\nif a.scores in (\"tok\", \"both\"):\n    TOKSC = \"/workspace/work/scores_tok.npz\"\n    if not os.path.exists(TOKSC):\n        subprocess.check_call([sys.executable, \"/workspace/score_pool_tok.py\"])\n    zt = np.load(TOKSC)\n    assert (zt[\"id\"] == ids).all()\n    ST = np.stack([zt[f\"s{k}\"] for k in range(NREG)])  # GPT-2-token-level register scores\n\n# ---------------------------------------------------------------- 1. quality gate\ngate = ((nw >= a.minw) & (nw <= a.maxw) & (z[\"wlen\"] >= 3.0) & (z[\"wlen\"] <= 9.0)\n        & (z[\"dupl\"] <= 0.25) & (z[\"digit\"] <= 0.20) & (z[\"nonascii\"] <= 0.05)\n        & (z[\"toprep\"] <= 0.16) & (z[\"symb\"] <= 0.02) & (z[\"boiler\"] <= 2))\n\n# ---------------------------------------------------------------- 2. scores\nshrink = nw / (nw + a.n0)\nS = np.stack([z[f\"s{k}\"] * shrink for k in range(NREG)])          # (NREG, N)\nSz = (S - S.mean(axis=1, keepdims=True)) / S.std(axis=1, keepdims=True)\n\n# exact token counts when available, else 4 chars/token\ntokc = np.maximum(1, nch // 4).astype(np.int64)\nif os.path.exists(TOKCNT):\n    d = json.load(open(TOKCNT))\n    pos = {int(v): i for i, v in enumerate(ids)}\n    for k, v in d.items():\n        i = pos.get(int(k))\n        if i is not None:\n            tokc[i] = v\n\n# ---------------------------------------------------------------- 3. assign + 4. dedup\nbest = Sz.argmax(axis=0)\ncand = np.where(gate)[0]\nlists = [[] for _ in range(NREG)]\nfor k in range(NREG):\n    sel = cand[best[cand] == k]\n    sel = sel[np.argsort(-S[k][sel])]                            # best-first\n    lists[k] = sel\n\nseen_ex, seen_nd = set(), set()\nex, nd = z[\"exhash\"], z[\"ndhash\"]\nded = [[] for _ in range(NREG)]\norder = np.concatenate([np.stack([np.full(len(l), k), l]) for k, l in enumerate(lists)], axis=1)\nrank = np.argsort(-np.array([Sz[k, i] for k, i in order.T]))      # dedup globally, best-first\nkeep = np.zeros(len(ids), dtype=bool)\nfor r in rank:\n    k, i = order[0, r], order[1, r]\n    e, n = int(ex[i]), int(nd[i])\n    if e in seen_ex or n in seen_nd:\n        continue\n    seen_ex.add(e); seen_nd.add(n); keep[i] = True\nfor k in range(NREG):\n    ded[k] = [i for i in lists[k] if keep[i]]\n\n# ---------------------------------------------------------------- 5. balanced interleave\nout, ptr, tok = [], [0] * NREG, [0] * NREG\nif a.balanced:\n    total = 0\n    while total < a.emit_tokens:\n        live = [k for k in range(NREG) if ptr[k] < len(ded[k])]\n        if not live:\n            break\n        k = min(live, key=lambda k: tok[k])\n        i = ded[k][ptr[k]]; ptr[k] += 1\n        out.append(int(ids[i])); tok[k] += int(tokc[i]); total += int(tokc[i])\nelse:  # ablation: one global ranking by best register score, no balancing\n    flat = np.concatenate(ded)\n    flat = flat[np.argsort(-Sz.max(axis=0)[flat])]\n    total = 0\n    for i in flat:\n        out.append(int(ids[i])); total += int(tokc[i])\n        tok[int(best[i])] += int(tokc[i])\n        if total >= a.emit_tokens:\n            break\n\nassert len(out) == len(set(out)), \"duplicate ids\"\nos.makedirs(os.path.dirname(a.out), exist_ok=True)\njson.dump(out, open(a.out, \"w\"))\n\n# report the composition of the prefix the trainer will actually use\ncum, comp, nd_ = 0, [0] * NREG, 0\npos = {int(v): i for i, v in enumerate(ids)}\nfor did in out:\n    i = pos[did]\n    comp[int(best[i])] += int(tokc[i]); cum += int(tokc[i]); nd_ += 1\n    if cum >= BUDGET:\n        break\nprint(f\"gate kept {gate.sum()}/{len(ids)}; dedup kept {keep.sum()}\")\nprint(f\"emitted {len(out)} ids (~{total/1e6:.1f}M tok); first {BUDGET/1e6:.0f}M tok = {nd_} docs\")\nprint(\"  register token mix in budget prefix:\", [round(c / max(1, cum), 3) for c in comp])\nprint(\"  wrote\", a.out)\n", "structuredPatch": [{"oldStart": 68, "oldLines": 9, "newStart": 68, "newLines": 17, "lines": ["         & (z[\"toprep\"] <= 0.16) & (z[\"symb\"] <= 0.02) & (z[\"boiler\"] <= 2))", " ", " # ---------------------------------------------------------------- 2. scores", "+def zsc(M):", "+    return (M - M.mean(axis=1, keepdims=True)) / M.std(axis=1, keepdims=True)", "+", " shrink = nw / (nw + a.n0)", "-S = np.stack([z[f\"s{k}\"] * shrink for k in range(NREG)])          # (NREG, N)", "-Sz = (S - S.mean(axis=1, keepdims=True)) / S.std(axis=1, keepdims=True)", "+if a.scores == \"word\":", "+    S = SW * shrink", "+elif a.scores == \"tok\":", "+    S = ST * shrink", "+else:                                    # average the two standardised views", "+    S = 0.5 * (zsc(SW) + zsc(ST)) * shrink", "+Sz = zsc(S)", " ", " # exact token counts when available, else 4 chars/token", " tokc = np.maximum(1, nch // 4).astype(np.int64)"]}], "userModified": false, "replaceAll": false}

The file /workspace/curate.py has been updated successfully. (file state is current in your context — no need to Read it back)

[metadata] {"filePath": "/workspace/curate.py", "oldString": "if a.balanced:\n    total = 0\n    while total < a.emit_tokens:\n        live = [k for k in range(NREG) if ptr[k] < len(ded[k])]\n        if not live:\n            break\n        k = min(live, key=lambda k: tok[k])", "newString": "share = [float(x) for x in a.share.split(\",\")] if a.share else [1.0 / NREG] * NREG\nshare = [s / sum(share) for s in share]\nif a.balanced:\n    total = 0\n    while total < a.emit_tokens:\n        live = [k for k in range(NREG) if ptr[k] < len(ded[k])]\n        if not live:\n            break\n        k = min(live, key=lambda k: tok[k] / share[k])", "originalFile": "\"\"\"Curate /workspace/data/pool.jsonl -> /workspace/submission/selection.json (priority order).\n\nSTATED CRITERION\n================\nThe eval target is an equal-parts mixture of four registers (encyclopedic, high-quality\nweb prose, news, technical Q&A).  We therefore build one *reference distribution per\nregister* from the disclosed dev target and select, for each register, the pool documents\nthat are most likely under it relative to the pool's own background -- then interleave the\nfour ranked lists so that any prefix of the submission (in particular the 12M-token prefix\nthe trainer consumes) is register-balanced.\n\n  1. GATE   Drop documents failing surface-quality checks (Gopher/C4-style): too short,\n            implausible mean word length, duplicated lines, digit/symbol/non-ASCII heavy,\n            single-token domination, navigation boilerplate.  Junk web text is never\n            worth budget, whatever its topic.\n  2. SCORE  s_k(d) = per-token log-likelihood ratio log p_k / q over hashed word uni+bi-grams\n            (DSIR, Xie et al. 2023), p_k = register-k n-grams from multi_dev.npy quarter k,\n            q = pool background n-grams.  Shrunk toward 0 for short docs by n/(n+n0) so a\n            handful of lucky n-grams cannot outrank a long, consistently on-target document.\n  3. ASSIGN Each surviving doc goes to its best-matching register (argmax of the\n            per-register standardised score), so the four lists are disjoint.\n  4. DEDUP  Exact (normalised-text hash) and near-duplicate (min-of-word-8-gram-hashes\n            sketch) collapse; the highest-scoring representative survives.\n  5. ORDER  Greedy balanced interleave: repeatedly emit the next-best document of whichever\n            register currently has the fewest tokens emitted.  Equal-parts target => equal\n            token shares, and the balance holds at *every* prefix length.\n\nUsage:  PYTHONHASHSEED=0 python curate.py [--out selection.json] [--minw 128] [--n0 64]\nRequires work/scores.npz (built automatically by score_pool.py if absent);\nwork/tokcount.json (exact GPT-2 token counts for candidates) is used for the token\naccounting when present, otherwise len(chars)/4 is used as the estimate.\n\"\"\"\nimport argparse, json, os, subprocess, sys\nimport numpy as np\n\nSCORES = \"/workspace/work/scores.npz\"\nTOKCNT = \"/workspace/work/tokcount.json\"\nNREG = 4\nBUDGET = 12_000_000\n\nap = argparse.ArgumentParser()\nap.add_argument(\"--out\", default=\"/workspace/submission/selection.json\")\nap.add_argument(\"--minw\", type=int, default=128)      # min words per doc\nap.add_argument(\"--maxw\", type=int, default=12000)\nap.add_argument(\"--n0\", type=int, default=64)         # length-shrinkage constant\nap.add_argument(\"--emit_tokens\", type=int, default=48_000_000)  # how much to list (>> budget)\nap.add_argument(\"--balanced\", type=int, default=1)    # 1 = per-register quotas, 0 = global rank\nap.add_argument(\"--scores\", default=\"tok\", choices=[\"tok\", \"word\", \"both\"])\nap.add_argument(\"--share\", default=\"\")                # e.g. \"0.3,0.2,0.2,0.3\" register token shares\na = ap.parse_args()\n\nif not os.path.exists(SCORES):\n    subprocess.check_call([sys.executable, \"/workspace/score_pool.py\"])\nz = np.load(SCORES)\nids, nw, nch = z[\"id\"], z[\"nw\"], z[\"nch\"]\nSW = np.stack([z[f\"s{k}\"] for k in range(NREG)])       # word-level register scores\nif a.scores in (\"tok\", \"both\"):\n    TOKSC = \"/workspace/work/scores_tok.npz\"\n    if not os.path.exists(TOKSC):\n        subprocess.check_call([sys.executable, \"/workspace/score_pool_tok.py\"])\n    zt = np.load(TOKSC)\n    assert (zt[\"id\"] == ids).all()\n    ST = np.stack([zt[f\"s{k}\"] for k in range(NREG)])  # GPT-2-token-level register scores\n\n# ---------------------------------------------------------------- 1. quality gate\ngate = ((nw >= a.minw) & (nw <= a.maxw) & (z[\"wlen\"] >= 3.0) & (z[\"wlen\"] <= 9.0)\n        & (z[\"dupl\"] <= 0.25) & (z[\"digit\"] <= 0.20) & (z[\"nonascii\"] <= 0.05)\n        & (z[\"toprep\"] <= 0.16) & (z[\"symb\"] <= 0.02) & (z[\"boiler\"] <= 2))\n\n# ---------------------------------------------------------------- 2. scores\ndef zsc(M):\n    return (M - M.mean(axis=1, keepdims=True)) / M.std(axis=1, keepdims=True)\n\nshrink = nw / (nw + a.n0)\nif a.scores == \"word\":\n    S = SW * shrink\nelif a.scores == \"tok\":\n    S = ST * shrink\nelse:                                    # average the two standardised views\n    S = 0.5 * (zsc(SW) + zsc(ST)) * shrink\nSz = zsc(S)\n\n# exact token counts when available, else 4 chars/token\ntokc = np.maximum(1, nch // 4).astype(np.int64)\nif os.path.exists(TOKCNT):\n    d = json.load(open(TOKCNT))\n    pos = {int(v): i for i, v in enumerate(ids)}\n    for k, v in d.items():\n        i = pos.get(int(k))\n        if i is not None:\n            tokc[i] = v\n\n# ---------------------------------------------------------------- 3. assign + 4. dedup\nbest = Sz.argmax(axis=0)\ncand = np.where(gate)[0]\nlists = [[] for _ in range(NREG)]\nfor k in range(NREG):\n    sel = cand[best[cand] == k]\n    sel = sel[np.argsort(-S[k][sel])]                            # best-first\n    lists[k] = sel\n\nseen_ex, seen_nd = set(), set()\nex, nd = z[\"exhash\"], z[\"ndhash\"]\nded = [[] for _ in range(NREG)]\norder = np.concatenate([np.stack([np.full(len(l), k), l]) for k, l in enumerate(lists)], axis=1)\nrank = np.argsort(-np.array([Sz[k, i] for k, i in order.T]))      # dedup globally, best-first\nkeep = np.zeros(len(ids), dtype=bool)\nfor r in rank:\n    k, i = order[0, r], order[1, r]\n    e, n = int(ex[i]), int(nd[i])\n    if e in seen_ex or n in seen_nd:\n        continue\n    seen_ex.add(e); seen_nd.add(n); keep[i] = True\nfor k in range(NREG):\n    ded[k] = [i for i in lists[k] if keep[i]]\n\n# ---------------------------------------------------------------- 5. balanced interleave\nout, ptr, tok = [], [0] * NREG, [0] * NREG\nif a.balanced:\n    total = 0\n    while total < a.emit_tokens:\n        live = [k for k in range(NREG) if ptr[k] < len(ded[k])]\n        if not live:\n            break\n        k = min(live, key=lambda k: tok[k])\n        i = ded[k][ptr[k]]; ptr[k] += 1\n        out.append(int(ids[i])); tok[k] += int(tokc[i]); total += int(tokc[i])\nelse:  # ablation: one global ranking by best register score, no balancing\n    flat = np.concatenate(ded)\n    flat = flat[np.argsort(-Sz.max(axis=0)[flat])]\n    total = 0\n    for i in flat:\n        out.append(int(ids[i])); total += int(tokc[i])\n        tok[int(best[i])] += int(tokc[i])\n        if total >= a.emit_tokens:\n            break\n\nassert len(out) == len(set(out)), \"duplicate ids\"\nos.makedirs(os.path.dirname(a.out), exist_ok=True)\njson.dump(out, open(a.out, \"w\"))\n\n# report the composition of the prefix the trainer will actually use\ncum, comp, nd_ = 0, [0] * NREG, 0\npos = {int(v): i for i, v in enumerate(ids)}\nfor did in out:\n    i = pos[did]\n    comp[int(best[i])] += int(tokc[i]); cum += int(tokc[i]); nd_ += 1\n    if cum >= BUDGET:\n        break\nprint(f\"gate kept {gate.sum()}/{len(ids)}; dedup kept {keep.sum()}\")\nprint(f\"emitted {len(out)} ids (~{total/1e6:.1f}M tok); first {BUDGET/1e6:.0f}M tok = {nd_} docs\")\nprint(\"  register token mix in budget prefix:\", [round(c / max(1, cum), 3) for c in comp])\nprint(\"  wrote\", a.out)\n", "structuredPatch": [{"oldStart": 116, "oldLines": 13, "newStart": 116, "newLines": 15, "lines": [" ", " # ---------------------------------------------------------------- 5. balanced interleave", " out, ptr, tok = [], [0] * NREG, [0] * NREG", "+share = [float(x) for x in a.share.split(\",\")] if a.share else [1.0 / NREG] * NREG", "+share = [s / sum(share) for s in share]", " if a.balanced:", "     total = 0", "     while total < a.emit_tokens:", "         live = [k for k in range(NREG) if ptr[k] < len(ded[k])]", "         if not live:", "             break", "-        k = min(live, key=lambda k: tok[k])", "+        k = min(live, key=lambda k: tok[k] / share[k])", "         i = ded[k][ptr[k]]; ptr[k] += 1", "         out.append(int(ids[i])); tok[k] += int(tokc[i]); total += int(tokc[i])", " else:  # ablation: one global ranking by best register score, no balancing"]}], "userModified": false, "replaceAll": false}

gate kept 125448/182016; dedup kept 113779
emitted 63178 ids (~48.0M tok); first 12M tok = 15686 docs
  register token mix in budget prefix: [0.25, 0.25, 0.25, 0.25]
  wrote work/sel_tok.json
[51268] 'omaya Hendousa founded dumped in slums east of Khartoum , severely tortured and her hair shaved confirmed resources close to her family . Hendousa fam'
[7641] "How Will Obama's Israel Visit Play Out at Home?3/21/2013 4:49PM | President Obama made a public appeal for sacrifice in the name of peace during his vis"
[41186] 'Amarnath Yatra 2017 | Amarnath Yatra 2017 News | Jammu and Kashmir police said three people, alleged conspirators in the 10 July attack on Amarnath pilgri'
[179526] ' series | Thursday, August 09, 2007 | The Jesus and the Eyewitnesses series | Oddly, I forgot to complete my series on Richard Bauckham’s Jesus and the Eyew'
[50793] ' Majesty King Peter II of Yugoslavia was the firstborn son of King Alexander I and Queen Maria of Yugoslavia. King Peter II was born in Belgrade 6 Sep'
[58452] '<|endoftext|>New Delhi: Five new Supreme Court judges were today sworn in, taking the top court’s strength to 28, including the Chief Justice of India'
[61757] '�The United States will encourage a peace and really a great peace deal” between Israel and the Palestinians, but they have to negotiate it themselves'
[100806] ' Market<|endoftext|>Manipur officials to study exodus situation | The Manipur government will send senior officials to Karnataka and Maharashtra to look'
[76107] 'Major Bhupinder Singh | In sept. 1965, Major Bhupinder Singh commanded the "B" squadron of the 4 Horse against Pakistani forces. On the 11th Sept. he su'
[112271] ' work!<|endoftext|>It is understood that the processor is a Qualcomm Snapdragon 630 this year with the processor 660 to be accompanied by the flagship'
[28976] 'umbai: Seventy days after they parted ways, the BJP and the Shiv Sena have once again come together and formed an alliance to share power in Maharasht'
[6827] 'Slippers were hurled at RJD leader Tejashwi Yadav during a poll rally in Bihar’s Aurangabad. The leader was sitting with party workers when the incide'

[stdout]
gate kept 125448/182016; dedup kept 113779
emitted 63178 ids (~48.0M tok); first 12M tok = 15686 docs
  register token mix in budget prefix: [0.25, 0.25, 0.25, 0.25]
  wrote work/sel_tok.json
[51268] 'omaya Hendousa founded dumped in slums east of Khartoum , severely tortured and her hair shaved confirmed resources close to her family . Hendousa fam'
[7641] "How Will Obama's Israel Visit Play Out at Home?3/21/2013 4:49PM | President Obama made a public appeal for sacrifice in the name of peace during his vis"
[41186] 'Amarnath Yatra 2017 | Amarnath Yatra 2017 News | Jammu and Kashmir police said three people, alleged conspirators in the 10 July attack on Amarnath pilgri'
[179526] ' series | Thursday, August 09, 2007 | The Jesus and the Eyewitnesses series | Oddly, I forgot to complete my series on Richard Bauckham’s Jesus and the Eyew'
[50793] ' Majesty King Peter II of Yugoslavia was the firstborn son of King Alexander I and Queen Maria of Yugoslavia. King Peter II was born in Belgrade 6 Sep'
[58452] '<|endoftext|>New Delhi: Five new Supreme Court judges were today sworn in, taking the top court’s strength to 28, including the Chief Justice of India'
[61757] '�The United States will encourage a peace and really a great peace deal” between Israel and the Palestinians, but they have to negotiate it themselves'
[100806] ' Market<|endoftext|>Manipur officials to study exodus situation | The Manipur government will send senior officials to Karnataka and Maharashtra to look'
[76107] 'Major Bhupinder Singh | In sept. 1965, Major Bhupinder Singh commanded the "B" squadron of the 4 Horse against Pakistani forces. On the 11th Sept. he su'
[112271] ' work!<|endoftext|>It is understood that the processor is a Qualcomm Snapdragon 630 this year with the processor 660 to be accompanied by the flagship'
[28976] 'umbai: Seventy days after they parted ways, the BJP and the Shiv Sena have once again come together and formed an alliance to share power in Maharasht'
[6827] 'Slippers were hurled at RJD leader Tejashwi Yadav during a poll rally in Bihar’s Aurangabad. The leader was sitting with party workers when the incide'
[metadata] {"noOutputExpected": false}

register token mix in budget prefix: [0.25, 0.25, 0.25, 0.25]
  wrote work/sel_both.json
packed 12000000 tokens from 63178 ids -> /tmp/tr_tok.npy
HELDOUT_PPL 317.434  loss 5.7603  train_s 51  /tmp/tr_tok.npy
packed 12000000 tokens from 46000 ids -> /tmp/tr_both.npy
HELDOUT_PPL 335.753  loss 5.8164  train_s 51  /tmp/tr_both.npy
4238

[stdout]
  register token mix in budget prefix: [0.25, 0.25, 0.25, 0.25]
  wrote work/sel_both.json
packed 12000000 tokens from 63178 ids -> /tmp/tr_tok.npy
HELDOUT_PPL 317.434  loss 5.7603  train_s 51  /tmp/tr_tok.npy
packed 12000000 tokens from 46000 ids -> /tmp/tr_both.npy
HELDOUT_PPL 335.753  loss 5.8164  train_s 51  /tmp/tr_both.npy
4238
[metadata] {"noOutputExpected": false}

The file /workspace/curate.py has been updated successfully. (file state is current in your context — no need to Read it back)

[metadata] {"filePath": "/workspace/curate.py", "oldString": "ap.add_argument(\"--share\", default=\"\")                # e.g. \"0.3,0.2,0.2,0.3\" register token shares\na = ap.parse_args()", "newString": "ap.add_argument(\"--share\", default=\"\")                # e.g. \"0.3,0.2,0.2,0.3\" register token shares\nap.add_argument(\"--gate\", type=int, default=1)        # 0 = length gate only (ablation)\nap.add_argument(\"--temp\", type=float, default=0.0)    # >0: Gumbel-perturbed order (diversity)\na = ap.parse_args()", "originalFile": "\"\"\"Curate /workspace/data/pool.jsonl -> /workspace/submission/selection.json (priority order).\n\nSTATED CRITERION\n================\nThe eval target is an equal-parts mixture of four registers (encyclopedic, high-quality\nweb prose, news, technical Q&A).  We therefore build one *reference distribution per\nregister* from the disclosed dev target and select, for each register, the pool documents\nthat are most likely under it relative to the pool's own background -- then interleave the\nfour ranked lists so that any prefix of the submission (in particular the 12M-token prefix\nthe trainer consumes) is register-balanced.\n\n  1. GATE   Drop documents failing surface-quality checks (Gopher/C4-style): too short,\n            implausible mean word length, duplicated lines, digit/symbol/non-ASCII heavy,\n            single-token domination, navigation boilerplate.  Junk web text is never\n            worth budget, whatever its topic.\n  2. SCORE  s_k(d) = per-token log-likelihood ratio log p_k / q over hashed word uni+bi-grams\n            (DSIR, Xie et al. 2023), p_k = register-k n-grams from multi_dev.npy quarter k,\n            q = pool background n-grams.  Shrunk toward 0 for short docs by n/(n+n0) so a\n            handful of lucky n-grams cannot outrank a long, consistently on-target document.\n  3. ASSIGN Each surviving doc goes to its best-matching register (argmax of the\n            per-register standardised score), so the four lists are disjoint.\n  4. DEDUP  Exact (normalised-text hash) and near-duplicate (min-of-word-8-gram-hashes\n            sketch) collapse; the highest-scoring representative survives.\n  5. ORDER  Greedy balanced interleave: repeatedly emit the next-best document of whichever\n            register currently has the fewest tokens emitted.  Equal-parts target => equal\n            token shares, and the balance holds at *every* prefix length.\n\nUsage:  PYTHONHASHSEED=0 python curate.py [--out selection.json] [--minw 128] [--n0 64]\nRequires work/scores.npz (built automatically by score_pool.py if absent);\nwork/tokcount.json (exact GPT-2 token counts for candidates) is used for the token\naccounting when present, otherwise len(chars)/4 is used as the estimate.\n\"\"\"\nimport argparse, json, os, subprocess, sys\nimport numpy as np\n\nSCORES = \"/workspace/work/scores.npz\"\nTOKCNT = \"/workspace/work/tokcount.json\"\nNREG = 4\nBUDGET = 12_000_000\n\nap = argparse.ArgumentParser()\nap.add_argument(\"--out\", default=\"/workspace/submission/selection.json\")\nap.add_argument(\"--minw\", type=int, default=128)      # min words per doc\nap.add_argument(\"--maxw\", type=int, default=12000)\nap.add_argument(\"--n0\", type=int, default=64)         # length-shrinkage constant\nap.add_argument(\"--emit_tokens\", type=int, default=48_000_000)  # how much to list (>> budget)\nap.add_argument(\"--balanced\", type=int, default=1)    # 1 = per-register quotas, 0 = global rank\nap.add_argument(\"--scores\", default=\"tok\", choices=[\"tok\", \"word\", \"both\"])\nap.add_argument(\"--share\", default=\"\")                # e.g. \"0.3,0.2,0.2,0.3\" register token shares\na = ap.parse_args()\n\nif not os.path.exists(SCORES):\n    subprocess.check_call([sys.executable, \"/workspace/score_pool.py\"])\nz = np.load(SCORES)\nids, nw, nch = z[\"id\"], z[\"nw\"], z[\"nch\"]\nSW = np.stack([z[f\"s{k}\"] for k in range(NREG)])       # word-level register scores\nif a.scores in (\"tok\", \"both\"):\n    TOKSC = \"/workspace/work/scores_tok.npz\"\n    if not os.path.exists(TOKSC):\n        subprocess.check_call([sys.executable, \"/workspace/score_pool_tok.py\"])\n    zt = np.load(TOKSC)\n    assert (zt[\"id\"] == ids).all()\n    ST = np.stack([zt[f\"s{k}\"] for k in range(NREG)])  # GPT-2-token-level register scores\n\n# ---------------------------------------------------------------- 1. quality gate\ngate = ((nw >= a.minw) & (nw <= a.maxw) & (z[\"wlen\"] >= 3.0) & (z[\"wlen\"] <= 9.0)\n        & (z[\"dupl\"] <= 0.25) & (z[\"digit\"] <= 0.20) & (z[\"nonascii\"] <= 0.05)\n        & (z[\"toprep\"] <= 0.16) & (z[\"symb\"] <= 0.02) & (z[\"boiler\"] <= 2))\n\n# ---------------------------------------------------------------- 2. scores\ndef zsc(M):\n    return (M - M.mean(axis=1, keepdims=True)) / M.std(axis=1, keepdims=True)\n\nshrink = nw / (nw + a.n0)\nif a.scores == \"word\":\n    S = SW * shrink\nelif a.scores == \"tok\":\n    S = ST * shrink\nelse:                                    # average the two standardised views\n    S = 0.5 * (zsc(SW) + zsc(ST)) * shrink\nSz = zsc(S)\n\n# exact token counts when available, else 4 chars/token\ntokc = np.maximum(1, nch // 4).astype(np.int64)\nif os.path.exists(TOKCNT):\n    d = json.load(open(TOKCNT))\n    pos = {int(v): i for i, v in enumerate(ids)}\n    for k, v in d.items():\n        i = pos.get(int(k))\n        if i is not None:\n            tokc[i] = v\n\n# ---------------------------------------------------------------- 3. assign + 4. dedup\nbest = Sz.argmax(axis=0)\ncand = np.where(gate)[0]\nlists = [[] for _ in range(NREG)]\nfor k in range(NREG):\n    sel = cand[best[cand] == k]\n    sel = sel[np.argsort(-S[k][sel])]                            # best-first\n    lists[k] = sel\n\nseen_ex, seen_nd = set(), set()\nex, nd = z[\"exhash\"], z[\"ndhash\"]\nded = [[] for _ in range(NREG)]\norder = np.concatenate([np.stack([np.full(len(l), k), l]) for k, l in enumerate(lists)], axis=1)\nrank = np.argsort(-np.array([Sz[k, i] for k, i in order.T]))      # dedup globally, best-first\nkeep = np.zeros(len(ids), dtype=bool)\nfor r in rank:\n    k, i = order[0, r], order[1, r]\n    e, n = int(ex[i]), int(nd[i])\n    if e in seen_ex or n in seen_nd:\n        continue\n    seen_ex.add(e); seen_nd.add(n); keep[i] = True\nfor k in range(NREG):\n    ded[k] = [i for i in lists[k] if keep[i]]\n\n# ---------------------------------------------------------------- 5. balanced interleave\nout, ptr, tok = [], [0] * NREG, [0] * NREG\nshare = [float(x) for x in a.share.split(\",\")] if a.share else [1.0 / NREG] * NREG\nshare = [s / sum(share) for s in share]\nif a.balanced:\n    total = 0\n    while total < a.emit_tokens:\n        live = [k for k in range(NREG) if ptr[k] < len(ded[k])]\n        if not live:\n            break\n        k = min(live, key=lambda k: tok[k] / share[k])\n        i = ded[k][ptr[k]]; ptr[k] += 1\n        out.append(int(ids[i])); tok[k] += int(tokc[i]); total += int(tokc[i])\nelse:  # ablation: one global ranking by best register score, no balancing\n    flat = np.concatenate(ded)\n    flat = flat[np.argsort(-Sz.max(axis=0)[flat])]\n    total = 0\n    for i in flat:\n        out.append(int(ids[i])); total += int(tokc[i])\n        tok[int(best[i])] += int(tokc[i])\n        if total >= a.emit_tokens:\n            break\n\nassert len(out) == len(set(out)), \"duplicate ids\"\nos.makedirs(os.path.dirname(a.out), exist_ok=True)\njson.dump(out, open(a.out, \"w\"))\n\n# report the composition of the prefix the trainer will actually use\ncum, comp, nd_ = 0, [0] * NREG, 0\npos = {int(v): i for i, v in enumerate(ids)}\nfor did in out:\n    i = pos[did]\n    comp[int(best[i])] += int(tokc[i]); cum += int(tokc[i]); nd_ += 1\n    if cum >= BUDGET:\n        break\nprint(f\"gate kept {gate.sum()}/{len(ids)}; dedup kept {keep.sum()}\")\nprint(f\"emitted {len(out)} ids (~{total/1e6:.1f}M tok); first {BUDGET/1e6:.0f}M tok = {nd_} docs\")\nprint(\"  register token mix in budget prefix:\", [round(c / max(1, cum), 3) for c in comp])\nprint(\"  wrote\", a.out)\n", "structuredPatch": [{"oldStart": 47, "oldLines": 6, "newStart": 47, "newLines": 8, "lines": [" ap.add_argument(\"--balanced\", type=int, default=1)    # 1 = per-register quotas, 0 = global rank", " ap.add_argument(\"--scores\", default=\"tok\", choices=[\"tok\", \"word\", \"both\"])", " ap.add_argument(\"--share\", default=\"\")                # e.g. \"0.3,0.2,0.2,0.3\" register token shares", "+ap.add_argument(\"--gate\", type=int, default=1)        # 0 = length gate only (ablation)", "+ap.add_argument(\"--temp\", type=float, default=0.0)    # >0: Gumbel-perturbed order (diversity)", " a = ap.parse_args()", " ", " if not os.path.exists(SCORES):"]}], "userModified": false, "replaceAll": false}

The file /workspace/curate.py has been updated successfully. (file state is current in your context — no need to Read it back)

[metadata] {"filePath": "/workspace/curate.py", "oldString": "gate = ((nw >= a.minw) & (nw <= a.maxw) & (z[\"wlen\"] >= 3.0) & (z[\"wlen\"] <= 9.0)\n        & (z[\"dupl\"] <= 0.25) & (z[\"digit\"] <= 0.20) & (z[\"nonascii\"] <= 0.05)\n        & (z[\"toprep\"] <= 0.16) & (z[\"symb\"] <= 0.02) & (z[\"boiler\"] <= 2))", "newString": "gate = (nw >= a.minw) & (nw <= a.maxw)\nif a.gate:\n    gate &= ((z[\"wlen\"] >= 3.0) & (z[\"wlen\"] <= 9.0)\n             & (z[\"dupl\"] <= 0.25) & (z[\"digit\"] <= 0.20) & (z[\"nonascii\"] <= 0.05)\n             & (z[\"toprep\"] <= 0.16) & (z[\"symb\"] <= 0.02) & (z[\"boiler\"] <= 2))", "originalFile": "\"\"\"Curate /workspace/data/pool.jsonl -> /workspace/submission/selection.json (priority order).\n\nSTATED CRITERION\n================\nThe eval target is an equal-parts mixture of four registers (encyclopedic, high-quality\nweb prose, news, technical Q&A).  We therefore build one *reference distribution per\nregister* from the disclosed dev target and select, for each register, the pool documents\nthat are most likely under it relative to the pool's own background -- then interleave the\nfour ranked lists so that any prefix of the submission (in particular the 12M-token prefix\nthe trainer consumes) is register-balanced.\n\n  1. GATE   Drop documents failing surface-quality checks (Gopher/C4-style): too short,\n            implausible mean word length, duplicated lines, digit/symbol/non-ASCII heavy,\n            single-token domination, navigation boilerplate.  Junk web text is never\n            worth budget, whatever its topic.\n  2. SCORE  s_k(d) = per-token log-likelihood ratio log p_k / q over hashed word uni+bi-grams\n            (DSIR, Xie et al. 2023), p_k = register-k n-grams from multi_dev.npy quarter k,\n            q = pool background n-grams.  Shrunk toward 0 for short docs by n/(n+n0) so a\n            handful of lucky n-grams cannot outrank a long, consistently on-target document.\n  3. ASSIGN Each surviving doc goes to its best-matching register (argmax of the\n            per-register standardised score), so the four lists are disjoint.\n  4. DEDUP  Exact (normalised-text hash) and near-duplicate (min-of-word-8-gram-hashes\n            sketch) collapse; the highest-scoring representative survives.\n  5. ORDER  Greedy balanced interleave: repeatedly emit the next-best document of whichever\n            register currently has the fewest tokens emitted.  Equal-parts target => equal\n            token shares, and the balance holds at *every* prefix length.\n\nUsage:  PYTHONHASHSEED=0 python curate.py [--out selection.json] [--minw 128] [--n0 64]\nRequires work/scores.npz (built automatically by score_pool.py if absent);\nwork/tokcount.json (exact GPT-2 token counts for candidates) is used for the token\naccounting when present, otherwise len(chars)/4 is used as the estimate.\n\"\"\"\nimport argparse, json, os, subprocess, sys\nimport numpy as np\n\nSCORES = \"/workspace/work/scores.npz\"\nTOKCNT = \"/workspace/work/tokcount.json\"\nNREG = 4\nBUDGET = 12_000_000\n\nap = argparse.ArgumentParser()\nap.add_argument(\"--out\", default=\"/workspace/submission/selection.json\")\nap.add_argument(\"--minw\", type=int, default=128)      # min words per doc\nap.add_argument(\"--maxw\", type=int, default=12000)\nap.add_argument(\"--n0\", type=int, default=64)         # length-shrinkage constant\nap.add_argument(\"--emit_tokens\", type=int, default=48_000_000)  # how much to list (>> budget)\nap.add_argument(\"--balanced\", type=int, default=1)    # 1 = per-register quotas, 0 = global rank\nap.add_argument(\"--scores\", default=\"tok\", choices=[\"tok\", \"word\", \"both\"])\nap.add_argument(\"--share\", default=\"\")                # e.g. \"0.3,0.2,0.2,0.3\" register token shares\nap.add_argument(\"--gate\", type=int, default=1)        # 0 = length gate only (ablation)\nap.add_argument(\"--temp\", type=float, default=0.0)    # >0: Gumbel-perturbed order (diversity)\na = ap.parse_args()\n\nif not os.path.exists(SCORES):\n    subprocess.check_call([sys.executable, \"/workspace/score_pool.py\"])\nz = np.load(SCORES)\nids, nw, nch = z[\"id\"], z[\"nw\"], z[\"nch\"]\nSW = np.stack([z[f\"s{k}\"] for k in range(NREG)])       # word-level register scores\nif a.scores in (\"tok\", \"both\"):\n    TOKSC = \"/workspace/work/scores_tok.npz\"\n    if not os.path.exists(TOKSC):\n        subprocess.check_call([sys.executable, \"/workspace/score_pool_tok.py\"])\n    zt = np.load(TOKSC)\n    assert (zt[\"id\"] == ids).all()\n    ST = np.stack([zt[f\"s{k}\"] for k in range(NREG)])  # GPT-2-token-level register scores\n\n# ---------------------------------------------------------------- 1. quality gate\ngate = ((nw >= a.minw) & (nw <= a.maxw) & (z[\"wlen\"] >= 3.0) & (z[\"wlen\"] <= 9.0)\n        & (z[\"dupl\"] <= 0.25) & (z[\"digit\"] <= 0.20) & (z[\"nonascii\"] <= 0.05)\n        & (z[\"toprep\"] <= 0.16) & (z[\"symb\"] <= 0.02) & (z[\"boiler\"] <= 2))\n\n# ---------------------------------------------------------------- 2. scores\ndef zsc(M):\n    return (M - M.mean(axis=1, keepdims=True)) / M.std(axis=1, keepdims=True)\n\nshrink = nw / (nw + a.n0)\nif a.scores == \"word\":\n    S = SW * shrink\nelif a.scores == \"tok\":\n    S = ST * shrink\nelse:                                    # average the two standardised views\n    S = 0.5 * (zsc(SW) + zsc(ST)) * shrink\nSz = zsc(S)\n\n# exact token counts when available, else 4 chars/token\ntokc = np.maximum(1, nch // 4).astype(np.int64)\nif os.path.exists(TOKCNT):\n    d = json.load(open(TOKCNT))\n    pos = {int(v): i for i, v in enumerate(ids)}\n    for k, v in d.items():\n        i = pos.get(int(k))\n        if i is not None:\n            tokc[i] = v\n\n# ---------------------------------------------------------------- 3. assign + 4. dedup\nbest = Sz.argmax(axis=0)\ncand = np.where(gate)[0]\nlists = [[] for _ in range(NREG)]\nfor k in range(NREG):\n    sel = cand[best[cand] == k]\n    sel = sel[np.argsort(-S[k][sel])]                            # best-first\n    lists[k] = sel\n\nseen_ex, seen_nd = set(), set()\nex, nd = z[\"exhash\"], z[\"ndhash\"]\nded = [[] for _ in range(NREG)]\norder = np.concatenate([np.stack([np.full(len(l), k), l]) for k, l in enumerate(lists)], axis=1)\nrank = np.argsort(-np.array([Sz[k, i] for k, i in order.T]))      # dedup globally, best-first\nkeep = np.zeros(len(ids), dtype=bool)\nfor r in rank:\n    k, i = order[0, r], order[1, r]\n    e, n = int(ex[i]), int(nd[i])\n    if e in seen_ex or n in seen_nd:\n        continue\n    seen_ex.add(e); seen_nd.add(n); keep[i] = True\nfor k in range(NREG):\n    ded[k] = [i for i in lists[k] if keep[i]]\n\n# ---------------------------------------------------------------- 5. balanced interleave\nout, ptr, tok = [], [0] * NREG, [0] * NREG\nshare = [float(x) for x in a.share.split(\",\")] if a.share else [1.0 / NREG] * NREG\nshare = [s / sum(share) for s in share]\nif a.balanced:\n    total = 0\n    while total < a.emit_tokens:\n        live = [k for k in range(NREG) if ptr[k] < len(ded[k])]\n        if not live:\n            break\n        k = min(live, key=lambda k: tok[k] / share[k])\n        i = ded[k][ptr[k]]; ptr[k] += 1\n        out.append(int(ids[i])); tok[k] += int(tokc[i]); total += int(tokc[i])\nelse:  # ablation: one global ranking by best register score, no balancing\n    flat = np.concatenate(ded)\n    flat = flat[np.argsort(-Sz.max(axis=0)[flat])]\n    total = 0\n    for i in flat:\n        out.append(int(ids[i])); total += int(tokc[i])\n        tok[int(best[i])] += int(tokc[i])\n        if total >= a.emit_tokens:\n            break\n\nassert len(out) == len(set(out)), \"duplicate ids\"\nos.makedirs(os.path.dirname(a.out), exist_ok=True)\njson.dump(out, open(a.out, \"w\"))\n\n# report the composition of the prefix the trainer will actually use\ncum, comp, nd_ = 0, [0] * NREG, 0\npos = {int(v): i for i, v in enumerate(ids)}\nfor did in out:\n    i = pos[did]\n    comp[int(best[i])] += int(tokc[i]); cum += int(tokc[i]); nd_ += 1\n    if cum >= BUDGET:\n        break\nprint(f\"gate kept {gate.sum()}/{len(ids)}; dedup kept {keep.sum()}\")\nprint(f\"emitted {len(out)} ids (~{total/1e6:.1f}M tok); first {BUDGET/1e6:.0f}M tok = {nd_} docs\")\nprint(\"  register token mix in budget prefix:\", [round(c / max(1, cum), 3) for c in comp])\nprint(\"  wrote\", a.out)\n", "structuredPatch": [{"oldStart": 65, "oldLines": 9, "newStart": 65, "newLines": 11, "lines": ["     ST = np.stack([zt[f\"s{k}\"] for k in range(NREG)])  # GPT-2-token-level register scores", " ", " # ---------------------------------------------------------------- 1. quality gate", "-gate = ((nw >= a.minw) & (nw <= a.maxw) & (z[\"wlen\"] >= 3.0) & (z[\"wlen\"] <= 9.0)", "-        & (z[\"dupl\"] <= 0.25) & (z[\"digit\"] <= 0.20) & (z[\"nonascii\"] <= 0.05)", "-        & (z[\"toprep\"] <= 0.16) & (z[\"symb\"] <= 0.02) & (z[\"boiler\"] <= 2))", "+gate = (nw >= a.minw) & (nw <= a.maxw)", "+if a.gate:", "+    gate &= ((z[\"wlen\"] >= 3.0) & (z[\"wlen\"] <= 9.0)", "+             & (z[\"dupl\"] <= 0.25) & (z[\"digit\"] <= 0.20) & (z[\"nonascii\"] <= 0.05)", "+             & (z[\"toprep\"] <= 0.16) & (z[\"symb\"] <= 0.02) & (z[\"boiler\"] <= 2))", " ", " # ---------------------------------------------------------------- 2. scores", " def zsc(M):"]}], "userModified": false, "replaceAll": false}

The file /workspace/curate.py has been updated successfully. (file state is current in your context — no need to Read it back)

[metadata] {"filePath": "/workspace/curate.py", "oldString": "best = Sz.argmax(axis=0)\ncand = np.where(gate)[0]\nlists = [[] for _ in range(NREG)]\nfor k in range(NREG):\n    sel = cand[best[cand] == k]\n    sel = sel[np.argsort(-S[k][sel])]                            # best-first\n    lists[k] = sel", "newString": "best = Sz.argmax(axis=0)\ncand = np.where(gate)[0]\nSord = Sz.copy()\nif a.temp > 0:                        # DSIR-style resampling: Gumbel noise on the score\n    g = np.random.default_rng(0).gumbel(size=Sord.shape)\n    Sord = Sord + a.temp * g\nlists = [[] for _ in range(NREG)]\nfor k in range(NREG):\n    sel = cand[best[cand] == k]\n    sel = sel[np.argsort(-Sord[k][sel])]                         # best-first\n    lists[k] = sel", "originalFile": "\"\"\"Curate /workspace/data/pool.jsonl -> /workspace/submission/selection.json (priority order).\n\nSTATED CRITERION\n================\nThe eval target is an equal-parts mixture of four registers (encyclopedic, high-quality\nweb prose, news, technical Q&A).  We therefore build one *reference distribution per\nregister* from the disclosed dev target and select, for each register, the pool documents\nthat are most likely under it relative to the pool's own background -- then interleave the\nfour ranked lists so that any prefix of the submission (in particular the 12M-token prefix\nthe trainer consumes) is register-balanced.\n\n  1. GATE   Drop documents failing surface-quality checks (Gopher/C4-style): too short,\n            implausible mean word length, duplicated lines, digit/symbol/non-ASCII heavy,\n            single-token domination, navigation boilerplate.  Junk web text is never\n            worth budget, whatever its topic.\n  2. SCORE  s_k(d) = per-token log-likelihood ratio log p_k / q over hashed word uni+bi-grams\n            (DSIR, Xie et al. 2023), p_k = register-k n-grams from multi_dev.npy quarter k,\n            q = pool background n-grams.  Shrunk toward 0 for short docs by n/(n+n0) so a\n            handful of lucky n-grams cannot outrank a long, consistently on-target document.\n  3. ASSIGN Each surviving doc goes to its best-matching register (argmax of the\n            per-register standardised score), so the four lists are disjoint.\n  4. DEDUP  Exact (normalised-text hash) and near-duplicate (min-of-word-8-gram-hashes\n            sketch) collapse; the highest-scoring representative survives.\n  5. ORDER  Greedy balanced interleave: repeatedly emit the next-best document of whichever\n            register currently has the fewest tokens emitted.  Equal-parts target => equal\n            token shares, and the balance holds at *every* prefix length.\n\nUsage:  PYTHONHASHSEED=0 python curate.py [--out selection.json] [--minw 128] [--n0 64]\nRequires work/scores.npz (built automatically by score_pool.py if absent);\nwork/tokcount.json (exact GPT-2 token counts for candidates) is used for the token\naccounting when present, otherwise len(chars)/4 is used as the estimate.\n\"\"\"\nimport argparse, json, os, subprocess, sys\nimport numpy as np\n\nSCORES = \"/workspace/work/scores.npz\"\nTOKCNT = \"/workspace/work/tokcount.json\"\nNREG = 4\nBUDGET = 12_000_000\n\nap = argparse.ArgumentParser()\nap.add_argument(\"--out\", default=\"/workspace/submission/selection.json\")\nap.add_argument(\"--minw\", type=int, default=128)      # min words per doc\nap.add_argument(\"--maxw\", type=int, default=12000)\nap.add_argument(\"--n0\", type=int, default=64)         # length-shrinkage constant\nap.add_argument(\"--emit_tokens\", type=int, default=48_000_000)  # how much to list (>> budget)\nap.add_argument(\"--balanced\", type=int, default=1)    # 1 = per-register quotas, 0 = global rank\nap.add_argument(\"--scores\", default=\"tok\", choices=[\"tok\", \"word\", \"both\"])\nap.add_argument(\"--share\", default=\"\")                # e.g. \"0.3,0.2,0.2,0.3\" register token shares\nap.add_argument(\"--gate\", type=int, default=1)        # 0 = length gate only (ablation)\nap.add_argument(\"--temp\", type=float, default=0.0)    # >0: Gumbel-perturbed order (diversity)\na = ap.parse_args()\n\nif not os.path.exists(SCORES):\n    subprocess.check_call([sys.executable, \"/workspace/score_pool.py\"])\nz = np.load(SCORES)\nids, nw, nch = z[\"id\"], z[\"nw\"], z[\"nch\"]\nSW = np.stack([z[f\"s{k}\"] for k in range(NREG)])       # word-level register scores\nif a.scores in (\"tok\", \"both\"):\n    TOKSC = \"/workspace/work/scores_tok.npz\"\n    if not os.path.exists(TOKSC):\n        subprocess.check_call([sys.executable, \"/workspace/score_pool_tok.py\"])\n    zt = np.load(TOKSC)\n    assert (zt[\"id\"] == ids).all()\n    ST = np.stack([zt[f\"s{k}\"] for k in range(NREG)])  # GPT-2-token-level register scores\n\n# ---------------------------------------------------------------- 1. quality gate\ngate = (nw >= a.minw) & (nw <= a.maxw)\nif a.gate:\n    gate &= ((z[\"wlen\"] >= 3.0) & (z[\"wlen\"] <= 9.0)\n             & (z[\"dupl\"] <= 0.25) & (z[\"digit\"] <= 0.20) & (z[\"nonascii\"] <= 0.05)\n             & (z[\"toprep\"] <= 0.16) & (z[\"symb\"] <= 0.02) & (z[\"boiler\"] <= 2))\n\n# ---------------------------------------------------------------- 2. scores\ndef zsc(M):\n    return (M - M.mean(axis=1, keepdims=True)) / M.std(axis=1, keepdims=True)\n\nshrink = nw / (nw + a.n0)\nif a.scores == \"word\":\n    S = SW * shrink\nelif a.scores == \"tok\":\n    S = ST * shrink\nelse:                                    # average the two standardised views\n    S = 0.5 * (zsc(SW) + zsc(ST)) * shrink\nSz = zsc(S)\n\n# exact token counts when available, else 4 chars/token\ntokc = np.maximum(1, nch // 4).astype(np.int64)\nif os.path.exists(TOKCNT):\n    d = json.load(open(TOKCNT))\n    pos = {int(v): i for i, v in enumerate(ids)}\n    for k, v in d.items():\n        i = pos.get(int(k))\n        if i is not None:\n            tokc[i] = v\n\n# ---------------------------------------------------------------- 3. assign + 4. dedup\nbest = Sz.argmax(axis=0)\ncand = np.where(gate)[0]\nlists = [[] for _ in range(NREG)]\nfor k in range(NREG):\n    sel = cand[best[cand] == k]\n    sel = sel[np.argsort(-S[k][sel])]                            # best-first\n    lists[k] = sel\n\nseen_ex, seen_nd = set(), set()\nex, nd = z[\"exhash\"], z[\"ndhash\"]\nded = [[] for _ in range(NREG)]\norder = np.concatenate([np.stack([np.full(len(l), k), l]) for k, l in enumerate(lists)], axis=1)\nrank = np.argsort(-np.array([Sz[k, i] for k, i in order.T]))      # dedup globally, best-first\nkeep = np.zeros(len(ids), dtype=bool)\nfor r in rank:\n    k, i = order[0, r], order[1, r]\n    e, n = int(ex[i]), int(nd[i])\n    if e in seen_ex or n in seen_nd:\n        continue\n    seen_ex.add(e); seen_nd.add(n); keep[i] = True\nfor k in range(NREG):\n    ded[k] = [i for i in lists[k] if keep[i]]\n\n# ---------------------------------------------------------------- 5. balanced interleave\nout, ptr, tok = [], [0] * NREG, [0] * NREG\nshare = [float(x) for x in a.share.split(\",\")] if a.share else [1.0 / NREG] * NREG\nshare = [s / sum(share) for s in share]\nif a.balanced:\n    total = 0\n    while total < a.emit_tokens:\n        live = [k for k in range(NREG) if ptr[k] < len(ded[k])]\n        if not live:\n            break\n        k = min(live, key=lambda k: tok[k] / share[k])\n        i = ded[k][ptr[k]]; ptr[k] += 1\n        out.append(int(ids[i])); tok[k] += int(tokc[i]); total += int(tokc[i])\nelse:  # ablation: one global ranking by best register score, no balancing\n    flat = np.concatenate(ded)\n    flat = flat[np.argsort(-Sz.max(axis=0)[flat])]\n    total = 0\n    for i in flat:\n        out.append(int(ids[i])); total += int(tokc[i])\n        tok[int(best[i])] += int(tokc[i])\n        if total >= a.emit_tokens:\n            break\n\nassert len(out) == len(set(out)), \"duplicate ids\"\nos.makedirs(os.path.dirname(a.out), exist_ok=True)\njson.dump(out, open(a.out, \"w\"))\n\n# report the composition of the prefix the trainer will actually use\ncum, comp, nd_ = 0, [0] * NREG, 0\npos = {int(v): i for i, v in enumerate(ids)}\nfor did in out:\n    i = pos[did]\n    comp[int(best[i])] += int(tokc[i]); cum += int(tokc[i]); nd_ += 1\n    if cum >= BUDGET:\n        break\nprint(f\"gate kept {gate.sum()}/{len(ids)}; dedup kept {keep.sum()}\")\nprint(f\"emitted {len(out)} ids (~{total/1e6:.1f}M tok); first {BUDGET/1e6:.0f}M tok = {nd_} docs\")\nprint(\"  register token mix in budget prefix:\", [round(c / max(1, cum), 3) for c in comp])\nprint(\"  wrote\", a.out)\n", "structuredPatch": [{"oldStart": 97, "oldLines": 10, "newStart": 97, "newLines": 14, "lines": [" # ---------------------------------------------------------------- 3. assign + 4. dedup", " best = Sz.argmax(axis=0)", " cand = np.where(gate)[0]", "+Sord = Sz.copy()", "+if a.temp > 0:                        # DSIR-style resampling: Gumbel noise on the score", "+    g = np.random.default_rng(0).gumbel(size=Sord.shape)", "+    Sord = Sord + a.temp * g", " lists = [[] for _ in range(NREG)]", " for k in range(NREG):", "     sel = cand[best[cand] == k]", "-    sel = sel[np.argsort(-S[k][sel])]                            # best-first", "+    sel = sel[np.argsort(-Sord[k][sel])]                         # best-first", "     lists[k] = sel", " ", " seen_ex, seen_nd = set(), set()"]}], "userModified": false, "replaceAll": false}

register token mix in budget prefix: [0.249, 0.249, 0.249, 0.252]
  nogate: HELDOUT_PPL 315.661
  register token mix in budget prefix: [0.25, 0.25, 0.25, 0.25]
  minw64: HELDOUT_PPL 311.619
  register token mix in budget prefix: [0.25, 0.25, 0.25, 0.25]
  n0_256: HELDOUT_PPL 315.098
  register token mix in budget prefix: [0.25, 0.25, 0.25, 0.25]
  temp05: HELDOUT_PPL 324.259
3968

[stdout]
  register token mix in budget prefix: [0.249, 0.249, 0.249, 0.252]
  nogate: HELDOUT_PPL 315.661
  register token mix in budget prefix: [0.25, 0.25, 0.25, 0.25]
  minw64: HELDOUT_PPL 311.619
  register token mix in budget prefix: [0.25, 0.25, 0.25, 0.25]
  n0_256: HELDOUT_PPL 315.098
  register token mix in budget prefix: [0.25, 0.25, 0.25, 0.25]
  temp05: HELDOUT_PPL 324.259
3968
[metadata] {"noOutputExpected": false}

m32ng256: HELDOUT_PPL 321.931
     emitted 90211 ids (~48.0M tok); first 12M tok = 27649 docs
  m64ng256: HELDOUT_PPL 320.361
     emitted 85967 ids (~48.0M tok); first 12M tok = 25456 docs
  m32g64: HELDOUT_PPL 310.576
     emitted 81904 ids (~48.0M tok); first 12M tok = 21297 docs
  m16ng256: HELDOUT_PPL 316.294
     emitted 90914 ids (~48.0M tok); first 12M tok = 28255 docs
3718

[stdout]
  m32ng256: HELDOUT_PPL 321.931
     emitted 90211 ids (~48.0M tok); first 12M tok = 27649 docs
  m64ng256: HELDOUT_PPL 320.361
     emitted 85967 ids (~48.0M tok); first 12M tok = 25456 docs
  m32g64: HELDOUT_PPL 310.576
     emitted 81904 ids (~48.0M tok); first 12M tok = 21297 docs
  m16ng256: HELDOUT_PPL 316.294
     emitted 90914 ids (~48.0M tok); first 12M tok = 28255 docs
3718
[metadata] {"noOutputExpected": false}

wrote quarters 250000
 reg0: HELDOUT_PPL 713.341  loss 6.5700
 reg1: HELDOUT_PPL 261.774  loss 5.5675
 reg2: HELDOUT_PPL 298.442  loss 5.6986
 reg3: HELDOUT_PPL 166.462  loss 5.1148
3457

[stdout]
wrote quarters 250000
 reg0: HELDOUT_PPL 713.341  loss 6.5700
 reg1: HELDOUT_PPL 261.774  loss 5.5675
 reg2: HELDOUT_PPL 298.442  loss 5.6986
 reg3: HELDOUT_PPL 166.462  loss 5.1148
3457
[metadata] {"noOutputExpected": false}

<stdin>:18: RuntimeWarning: divide by zero encountered in log
<stdin>:18: RuntimeWarning: invalid value encountered in multiply
--- reg0: unigram KL(target||selection) = 1.0035 nats/token
   worst tokens: ' ,':0.250, ' .':0.135, ' @':0.048, '@':0.043, " '":0.025, ' "':0.018, ' )':0.018, '<|endoftext|>':0.013, ' the':0.010, ' The':0.008, ' was':0.007, ' and':0.006
--- reg1: unigram KL(target||selection) = 0.2717 nats/token
   worst tokens: '\n':0.018, '�':0.004, ' ...':0.003, ' that':0.003, ' you':0.002, '�':0.002, ',':0.002, ' a':0.002, ' to':0.002, ' �':0.002, '�':0.001, '�':0.001
--- reg2: unigram KL(target||selection) = 0.4597 nats/token
   worst tokens: ' off':0.007, 'gar':0.005, ' to':0.005, '0':0.005, '."':0.005, ' the':0.004, ' said':0.004, ',':0.004, 'ram':0.004, ' length':0.004, ' on':0.004, ' leg':0.004
--- reg3: unigram KL(target||selection) = 1.5355 nats/token
   worst tokens: ' ':0.286, '>':0.178, 'p':0.110, '</':0.084, '<':0.079, 'code':0.066, '\n':0.055, '.</':0.036, '&':0.031, 'gt':0.030, 'pre':0.028, ';':0.027

[stdout]
<stdin>:18: RuntimeWarning: divide by zero encountered in log
<stdin>:18: RuntimeWarning: invalid value encountered in multiply
--- reg0: unigram KL(target||selection) = 1.0035 nats/token
   worst tokens: ' ,':0.250, ' .':0.135, ' @':0.048, '@':0.043, " '":0.025, ' "':0.018, ' )':0.018, '<|endoftext|>':0.013, ' the':0.010, ' The':0.008, ' was':0.007, ' and':0.006
--- reg1: unigram KL(target||selection) = 0.2717 nats/token
   worst tokens: '\n':0.018, '�':0.004, ' ...':0.003, ' that':0.003, ' you':0.002, '�':0.002, ',':0.002, ' a':0.002, ' to':0.002, ' �':0.002, '�':0.001, '�':0.001
--- reg2: unigram KL(target||selection) = 0.4597 nats/token
   worst tokens: ' off':0.007, 'gar':0.005, ' to':0.005, '0':0.005, '."':0.005, ' the':0.004, ' said':0.004, ',':0.004, 'ram':0.004, ' length':0.004, ' on':0.004, ' leg':0.004
--- reg3: unigram KL(target||selection) = 1.5355 nats/token
   worst tokens: ' ':0.286, '>':0.178, 'p':0.110, '</':0.084, '<':0.079, 'code':0.066, '\n':0.055, '.</':0.036, '&':0.031, 'gt':0.030, 'pre':0.028, ';':0.027
[metadata] {"noOutputExpected": false}

reg0 target: spacing rate 0.0899  markup rate 0.0000
reg3 target: spacing rate 0.0016  markup rate 0.0717
pool spacing pct [0.     0.0015 0.0145 0.0483] max 0.3057
pool markup  pct [0.     0.     0.0054 0.0506] max 0.1693
  top500 spacing: 0.69M tok, min rate 0.0322
  top1000 spacing: 1.75M tok, min rate 0.0218
  top2000 spacing: 2.93M tok, min rate 0.0135
  top4000 spacing: 4.96M tok, min rate 0.0082
  top500 markup: 1.14M tok, min rate 0.0257
  top1000 markup: 1.90M tok, min rate 0.0116
  top2000 markup: 3.64M tok, min rate 0.0047

[stdout]
reg0 target: spacing rate 0.0899  markup rate 0.0000
reg3 target: spacing rate 0.0016  markup rate 0.0717
pool spacing pct [0.     0.0015 0.0145 0.0483] max 0.3057
pool markup  pct [0.     0.     0.0054 0.0506] max 0.1693
  top500 spacing: 0.69M tok, min rate 0.0322
  top1000 spacing: 1.75M tok, min rate 0.0218
  top2000 spacing: 2.93M tok, min rate 0.0135
  top4000 spacing: 4.96M tok, min rate 0.0082
  top500 markup: 1.14M tok, min rate 0.0257
  top1000 markup: 1.90M tok, min rate 0.0116
  top2000 markup: 3.64M tok, min rate 0.0047
[metadata] {"noOutputExpected": false}

The file /workspace/curate.py has been updated successfully. (file state is current in your context — no need to Read it back)

[metadata] {"filePath": "/workspace/curate.py", "oldString": "ap.add_argument(\"--temp\", type=float, default=0.0)    # >0: Gumbel-perturbed order (diversity)\na = ap.parse_args()", "newString": "ap.add_argument(\"--temp\", type=float, default=0.0)    # >0: Gumbel-perturbed order (diversity)\nap.add_argument(\"--fill_sp\", type=int, default=0)     # tokens reserved for spaced-punctuation docs\nap.add_argument(\"--fill_mk\", type=int, default=0)     # tokens reserved for markup/code-ish docs\na = ap.parse_args()", "originalFile": "\"\"\"Curate /workspace/data/pool.jsonl -> /workspace/submission/selection.json (priority order).\n\nSTATED CRITERION\n================\nThe eval target is an equal-parts mixture of four registers (encyclopedic, high-quality\nweb prose, news, technical Q&A).  We therefore build one *reference distribution per\nregister* from the disclosed dev target and select, for each register, the pool documents\nthat are most likely under it relative to the pool's own background -- then interleave the\nfour ranked lists so that any prefix of the submission (in particular the 12M-token prefix\nthe trainer consumes) is register-balanced.\n\n  1. GATE   Drop documents failing surface-quality checks (Gopher/C4-style): too short,\n            implausible mean word length, duplicated lines, digit/symbol/non-ASCII heavy,\n            single-token domination, navigation boilerplate.  Junk web text is never\n            worth budget, whatever its topic.\n  2. SCORE  s_k(d) = per-token log-likelihood ratio log p_k / q over hashed word uni+bi-grams\n            (DSIR, Xie et al. 2023), p_k = register-k n-grams from multi_dev.npy quarter k,\n            q = pool background n-grams.  Shrunk toward 0 for short docs by n/(n+n0) so a\n            handful of lucky n-grams cannot outrank a long, consistently on-target document.\n  3. ASSIGN Each surviving doc goes to its best-matching register (argmax of the\n            per-register standardised score), so the four lists are disjoint.\n  4. DEDUP  Exact (normalised-text hash) and near-duplicate (min-of-word-8-gram-hashes\n            sketch) collapse; the highest-scoring representative survives.\n  5. ORDER  Greedy balanced interleave: repeatedly emit the next-best document of whichever\n            register currently has the fewest tokens emitted.  Equal-parts target => equal\n            token shares, and the balance holds at *every* prefix length.\n\nUsage:  PYTHONHASHSEED=0 python curate.py [--out selection.json] [--minw 128] [--n0 64]\nRequires work/scores.npz (built automatically by score_pool.py if absent);\nwork/tokcount.json (exact GPT-2 token counts for candidates) is used for the token\naccounting when present, otherwise len(chars)/4 is used as the estimate.\n\"\"\"\nimport argparse, json, os, subprocess, sys\nimport numpy as np\n\nSCORES = \"/workspace/work/scores.npz\"\nTOKCNT = \"/workspace/work/tokcount.json\"\nNREG = 4\nBUDGET = 12_000_000\n\nap = argparse.ArgumentParser()\nap.add_argument(\"--out\", default=\"/workspace/submission/selection.json\")\nap.add_argument(\"--minw\", type=int, default=128)      # min words per doc\nap.add_argument(\"--maxw\", type=int, default=12000)\nap.add_argument(\"--n0\", type=int, default=64)         # length-shrinkage constant\nap.add_argument(\"--emit_tokens\", type=int, default=48_000_000)  # how much to list (>> budget)\nap.add_argument(\"--balanced\", type=int, default=1)    # 1 = per-register quotas, 0 = global rank\nap.add_argument(\"--scores\", default=\"tok\", choices=[\"tok\", \"word\", \"both\"])\nap.add_argument(\"--share\", default=\"\")                # e.g. \"0.3,0.2,0.2,0.3\" register token shares\nap.add_argument(\"--gate\", type=int, default=1)        # 0 = length gate only (ablation)\nap.add_argument(\"--temp\", type=float, default=0.0)    # >0: Gumbel-perturbed order (diversity)\na = ap.parse_args()\n\nif not os.path.exists(SCORES):\n    subprocess.check_call([sys.executable, \"/workspace/score_pool.py\"])\nz = np.load(SCORES)\nids, nw, nch = z[\"id\"], z[\"nw\"], z[\"nch\"]\nSW = np.stack([z[f\"s{k}\"] for k in range(NREG)])       # word-level register scores\nif a.scores in (\"tok\", \"both\"):\n    TOKSC = \"/workspace/work/scores_tok.npz\"\n    if not os.path.exists(TOKSC):\n        subprocess.check_call([sys.executable, \"/workspace/score_pool_tok.py\"])\n    zt = np.load(TOKSC)\n    assert (zt[\"id\"] == ids).all()\n    ST = np.stack([zt[f\"s{k}\"] for k in range(NREG)])  # GPT-2-token-level register scores\n\n# ---------------------------------------------------------------- 1. quality gate\ngate = (nw >= a.minw) & (nw <= a.maxw)\nif a.gate:\n    gate &= ((z[\"wlen\"] >= 3.0) & (z[\"wlen\"] <= 9.0)\n             & (z[\"dupl\"] <= 0.25) & (z[\"digit\"] <= 0.20) & (z[\"nonascii\"] <= 0.05)\n             & (z[\"toprep\"] <= 0.16) & (z[\"symb\"] <= 0.02) & (z[\"boiler\"] <= 2))\n\n# ---------------------------------------------------------------- 2. scores\ndef zsc(M):\n    return (M - M.mean(axis=1, keepdims=True)) / M.std(axis=1, keepdims=True)\n\nshrink = nw / (nw + a.n0)\nif a.scores == \"word\":\n    S = SW * shrink\nelif a.scores == \"tok\":\n    S = ST * shrink\nelse:                                    # average the two standardised views\n    S = 0.5 * (zsc(SW) + zsc(ST)) * shrink\nSz = zsc(S)\n\n# exact token counts when available, else 4 chars/token\ntokc = np.maximum(1, nch // 4).astype(np.int64)\nif os.path.exists(TOKCNT):\n    d = json.load(open(TOKCNT))\n    pos = {int(v): i for i, v in enumerate(ids)}\n    for k, v in d.items():\n        i = pos.get(int(k))\n        if i is not None:\n            tokc[i] = v\n\n# ---------------------------------------------------------------- 3. assign + 4. dedup\nbest = Sz.argmax(axis=0)\ncand = np.where(gate)[0]\nSord = Sz.copy()\nif a.temp > 0:                        # DSIR-style resampling: Gumbel noise on the score\n    g = np.random.default_rng(0).gumbel(size=Sord.shape)\n    Sord = Sord + a.temp * g\nlists = [[] for _ in range(NREG)]\nfor k in range(NREG):\n    sel = cand[best[cand] == k]\n    sel = sel[np.argsort(-Sord[k][sel])]                         # best-first\n    lists[k] = sel\n\nseen_ex, seen_nd = set(), set()\nex, nd = z[\"exhash\"], z[\"ndhash\"]\nded = [[] for _ in range(NREG)]\norder = np.concatenate([np.stack([np.full(len(l), k), l]) for k, l in enumerate(lists)], axis=1)\nrank = np.argsort(-np.array([Sz[k, i] for k, i in order.T]))      # dedup globally, best-first\nkeep = np.zeros(len(ids), dtype=bool)\nfor r in rank:\n    k, i = order[0, r], order[1, r]\n    e, n = int(ex[i]), int(nd[i])\n    if e in seen_ex or n in seen_nd:\n        continue\n    seen_ex.add(e); seen_nd.add(n); keep[i] = True\nfor k in range(NREG):\n    ded[k] = [i for i in lists[k] if keep[i]]\n\n# ---------------------------------------------------------------- 5. balanced interleave\nout, ptr, tok = [], [0] * NREG, [0] * NREG\nshare = [float(x) for x in a.share.split(\",\")] if a.share else [1.0 / NREG] * NREG\nshare = [s / sum(share) for s in share]\nif a.balanced:\n    total = 0\n    while total < a.emit_tokens:\n        live = [k for k in range(NREG) if ptr[k] < len(ded[k])]\n        if not live:\n            break\n        k = min(live, key=lambda k: tok[k] / share[k])\n        i = ded[k][ptr[k]]; ptr[k] += 1\n        out.append(int(ids[i])); tok[k] += int(tokc[i]); total += int(tokc[i])\nelse:  # ablation: one global ranking by best register score, no balancing\n    flat = np.concatenate(ded)\n    flat = flat[np.argsort(-Sz.max(axis=0)[flat])]\n    total = 0\n    for i in flat:\n        out.append(int(ids[i])); total += int(tokc[i])\n        tok[int(best[i])] += int(tokc[i])\n        if total >= a.emit_tokens:\n            break\n\nassert len(out) == len(set(out)), \"duplicate ids\"\nos.makedirs(os.path.dirname(a.out), exist_ok=True)\njson.dump(out, open(a.out, \"w\"))\n\n# report the composition of the prefix the trainer will actually use\ncum, comp, nd_ = 0, [0] * NREG, 0\npos = {int(v): i for i, v in enumerate(ids)}\nfor did in out:\n    i = pos[did]\n    comp[int(best[i])] += int(tokc[i]); cum += int(tokc[i]); nd_ += 1\n    if cum >= BUDGET:\n        break\nprint(f\"gate kept {gate.sum()}/{len(ids)}; dedup kept {keep.sum()}\")\nprint(f\"emitted {len(out)} ids (~{total/1e6:.1f}M tok); first {BUDGET/1e6:.0f}M tok = {nd_} docs\")\nprint(\"  register token mix in budget prefix:\", [round(c / max(1, cum), 3) for c in comp])\nprint(\"  wrote\", a.out)\n", "structuredPatch": [{"oldStart": 49, "oldLines": 6, "newStart": 49, "newLines": 8, "lines": [" ap.add_argument(\"--share\", default=\"\")                # e.g. \"0.3,0.2,0.2,0.3\" register token shares", " ap.add_argument(\"--gate\", type=int, default=1)        # 0 = length gate only (ablation)", " ap.add_argument(\"--temp\", type=float, default=0.0)    # >0: Gumbel-perturbed order (diversity)", "+ap.add_argument(\"--fill_sp\", type=int, default=0)     # tokens reserved for spaced-punctuation docs", "+ap.add_argument(\"--fill_mk\", type=int, default=0)     # tokens reserved for markup/code-ish docs", " a = ap.parse_args()", " ", " if not os.path.exists(SCORES):"]}], "userModified": false, "replaceAll": false}

The file /workspace/curate.py has been updated successfully. (file state is current in your context — no need to Read it back)

[metadata] {"filePath": "/workspace/curate.py", "oldString": "# ---------------------------------------------------------------- 5. balanced interleave\nout, ptr, tok = [], [0] * NREG, [0] * NREG", "newString": "# ------------------------------------------- 4b. surface-convention gap fillers\n# The target's encyclopedic quarter writes punctuation with a leading space (\" ,\", \" .\",\n# \"@-@\") and its technical-Q&A quarter is HTML-marked-up (\"<p>\", \"</\", \"&gt;\").  Those are\n# *token-level conventions*, not topics: a model that never sees them assigns them almost\n# no probability, and they account for ~9% / ~7% of the tokens in those quarters.  DSIR\n# ranks documents by average n-gram fit and so under-buys these rare-in-pool conventions,\n# hence an explicit reserved slice of the budget for the documents that carry them.\nsurf = np.load(\"/workspace/work/surface.npy\")   # rows: [spaced-punct rate, markup rate]\npre, used_ex, used_nd = [], set(), set()\nlight = (nw >= 32) & (z[\"nonascii\"] <= 0.10) & (nw <= a.maxw)\nfor row, want in ((0, a.fill_sp), (1, a.fill_mk)):\n    if want <= 0:\n        continue\n    got = 0\n    for i in np.argsort(-surf[row]):\n        if not light[i]:\n            continue\n        e, n = int(ex[i]), int(nd[i])\n        if e in used_ex or n in used_nd:\n            continue\n        used_ex.add(e); used_nd.add(n)\n        pre.append(int(i)); got += int(tokc[i])\n        if got >= want:\n            break\nprefill = set(pre)\nded = [[i for i in l if i not in prefill] for l in ded]\n\n# ---------------------------------------------------------------- 5. balanced interleave\nout, ptr, tok = [int(ids[i]) for i in pre], [0] * NREG, [0] * NREG\nfill_tok = sum(int(tokc[i]) for i in pre)", "originalFile": "\"\"\"Curate /workspace/data/pool.jsonl -> /workspace/submission/selection.json (priority order).\n\nSTATED CRITERION\n================\nThe eval target is an equal-parts mixture of four registers (encyclopedic, high-quality\nweb prose, news, technical Q&A).  We therefore build one *reference distribution per\nregister* from the disclosed dev target and select, for each register, the pool documents\nthat are most likely under it relative to the pool's own background -- then interleave the\nfour ranked lists so that any prefix of the submission (in particular the 12M-token prefix\nthe trainer consumes) is register-balanced.\n\n  1. GATE   Drop documents failing surface-quality checks (Gopher/C4-style): too short,\n            implausible mean word length, duplicated lines, digit/symbol/non-ASCII heavy,\n            single-token domination, navigation boilerplate.  Junk web text is never\n            worth budget, whatever its topic.\n  2. SCORE  s_k(d) = per-token log-likelihood ratio log p_k / q over hashed word uni+bi-grams\n            (DSIR, Xie et al. 2023), p_k = register-k n-grams from multi_dev.npy quarter k,\n            q = pool background n-grams.  Shrunk toward 0 for short docs by n/(n+n0) so a\n            handful of lucky n-grams cannot outrank a long, consistently on-target document.\n  3. ASSIGN Each surviving doc goes to its best-matching register (argmax of the\n            per-register standardised score), so the four lists are disjoint.\n  4. DEDUP  Exact (normalised-text hash) and near-duplicate (min-of-word-8-gram-hashes\n            sketch) collapse; the highest-scoring representative survives.\n  5. ORDER  Greedy balanced interleave: repeatedly emit the next-best document of whichever\n            register currently has the fewest tokens emitted.  Equal-parts target => equal\n            token shares, and the balance holds at *every* prefix length.\n\nUsage:  PYTHONHASHSEED=0 python curate.py [--out selection.json] [--minw 128] [--n0 64]\nRequires work/scores.npz (built automatically by score_pool.py if absent);\nwork/tokcount.json (exact GPT-2 token counts for candidates) is used for the token\naccounting when present, otherwise len(chars)/4 is used as the estimate.\n\"\"\"\nimport argparse, json, os, subprocess, sys\nimport numpy as np\n\nSCORES = \"/workspace/work/scores.npz\"\nTOKCNT = \"/workspace/work/tokcount.json\"\nNREG = 4\nBUDGET = 12_000_000\n\nap = argparse.ArgumentParser()\nap.add_argument(\"--out\", default=\"/workspace/submission/selection.json\")\nap.add_argument(\"--minw\", type=int, default=128)      # min words per doc\nap.add_argument(\"--maxw\", type=int, default=12000)\nap.add_argument(\"--n0\", type=int, default=64)         # length-shrinkage constant\nap.add_argument(\"--emit_tokens\", type=int, default=48_000_000)  # how much to list (>> budget)\nap.add_argument(\"--balanced\", type=int, default=1)    # 1 = per-register quotas, 0 = global rank\nap.add_argument(\"--scores\", default=\"tok\", choices=[\"tok\", \"word\", \"both\"])\nap.add_argument(\"--share\", default=\"\")                # e.g. \"0.3,0.2,0.2,0.3\" register token shares\nap.add_argument(\"--gate\", type=int, default=1)        # 0 = length gate only (ablation)\nap.add_argument(\"--temp\", type=float, default=0.0)    # >0: Gumbel-perturbed order (diversity)\nap.add_argument(\"--fill_sp\", type=int, default=0)     # tokens reserved for spaced-punctuation docs\nap.add_argument(\"--fill_mk\", type=int, default=0)     # tokens reserved for markup/code-ish docs\na = ap.parse_args()\n\nif not os.path.exists(SCORES):\n    subprocess.check_call([sys.executable, \"/workspace/score_pool.py\"])\nz = np.load(SCORES)\nids, nw, nch = z[\"id\"], z[\"nw\"], z[\"nch\"]\nSW = np.stack([z[f\"s{k}\"] for k in range(NREG)])       # word-level register scores\nif a.scores in (\"tok\", \"both\"):\n    TOKSC = \"/workspace/work/scores_tok.npz\"\n    if not os.path.exists(TOKSC):\n        subprocess.check_call([sys.executable, \"/workspace/score_pool_tok.py\"])\n    zt = np.load(TOKSC)\n    assert (zt[\"id\"] == ids).all()\n    ST = np.stack([zt[f\"s{k}\"] for k in range(NREG)])  # GPT-2-token-level register scores\n\n# ---------------------------------------------------------------- 1. quality gate\ngate = (nw >= a.minw) & (nw <= a.maxw)\nif a.gate:\n    gate &= ((z[\"wlen\"] >= 3.0) & (z[\"wlen\"] <= 9.0)\n             & (z[\"dupl\"] <= 0.25) & (z[\"digit\"] <= 0.20) & (z[\"nonascii\"] <= 0.05)\n             & (z[\"toprep\"] <= 0.16) & (z[\"symb\"] <= 0.02) & (z[\"boiler\"] <= 2))\n\n# ---------------------------------------------------------------- 2. scores\ndef zsc(M):\n    return (M - M.mean(axis=1, keepdims=True)) / M.std(axis=1, keepdims=True)\n\nshrink = nw / (nw + a.n0)\nif a.scores == \"word\":\n    S = SW * shrink\nelif a.scores == \"tok\":\n    S = ST * shrink\nelse:                                    # average the two standardised views\n    S = 0.5 * (zsc(SW) + zsc(ST)) * shrink\nSz = zsc(S)\n\n# exact token counts when available, else 4 chars/token\ntokc = np.maximum(1, nch // 4).astype(np.int64)\nif os.path.exists(TOKCNT):\n    d = json.load(open(TOKCNT))\n    pos = {int(v): i for i, v in enumerate(ids)}\n    for k, v in d.items():\n        i = pos.get(int(k))\n        if i is not None:\n            tokc[i] = v\n\n# ---------------------------------------------------------------- 3. assign + 4. dedup\nbest = Sz.argmax(axis=0)\ncand = np.where(gate)[0]\nSord = Sz.copy()\nif a.temp > 0:                        # DSIR-style resampling: Gumbel noise on the score\n    g = np.random.default_rng(0).gumbel(size=Sord.shape)\n    Sord = Sord + a.temp * g\nlists = [[] for _ in range(NREG)]\nfor k in range(NREG):\n    sel = cand[best[cand] == k]\n    sel = sel[np.argsort(-Sord[k][sel])]                         # best-first\n    lists[k] = sel\n\nseen_ex, seen_nd = set(), set()\nex, nd = z[\"exhash\"], z[\"ndhash\"]\nded = [[] for _ in range(NREG)]\norder = np.concatenate([np.stack([np.full(len(l), k), l]) for k, l in enumerate(lists)], axis=1)\nrank = np.argsort(-np.array([Sz[k, i] for k, i in order.T]))      # dedup globally, best-first\nkeep = np.zeros(len(ids), dtype=bool)\nfor r in rank:\n    k, i = order[0, r], order[1, r]\n    e, n = int(ex[i]), int(nd[i])\n    if e in seen_ex or n in seen_nd:\n        continue\n    seen_ex.add(e); seen_nd.add(n); keep[i] = True\nfor k in range(NREG):\n    ded[k] = [i for i in lists[k] if keep[i]]\n\n# ---------------------------------------------------------------- 5. balanced interleave\nout, ptr, tok = [], [0] * NREG, [0] * NREG\nshare = [float(x) for x in a.share.split(\",\")] if a.share else [1.0 / NREG] * NREG\nshare = [s / sum(share) for s in share]\nif a.balanced:\n    total = 0\n    while total < a.emit_tokens:\n        live = [k for k in range(NREG) if ptr[k] < len(ded[k])]\n        if not live:\n            break\n        k = min(live, key=lambda k: tok[k] / share[k])\n        i = ded[k][ptr[k]]; ptr[k] += 1\n        out.append(int(ids[i])); tok[k] += int(tokc[i]); total += int(tokc[i])\nelse:  # ablation: one global ranking by best register score, no balancing\n    flat = np.concatenate(ded)\n    flat = flat[np.argsort(-Sz.max(axis=0)[flat])]\n    total = 0\n    for i in flat:\n        out.append(int(ids[i])); total += int(tokc[i])\n        tok[int(best[i])] += int(tokc[i])\n        if total >= a.emit_tokens:\n            break\n\nassert len(out) == len(set(out)), \"duplicate ids\"\nos.makedirs(os.path.dirname(a.out), exist_ok=True)\njson.dump(out, open(a.out, \"w\"))\n\n# report the composition of the prefix the trainer will actually use\ncum, comp, nd_ = 0, [0] * NREG, 0\npos = {int(v): i for i, v in enumerate(ids)}\nfor did in out:\n    i = pos[did]\n    comp[int(best[i])] += int(tokc[i]); cum += int(tokc[i]); nd_ += 1\n    if cum >= BUDGET:\n        break\nprint(f\"gate kept {gate.sum()}/{len(ids)}; dedup kept {keep.sum()}\")\nprint(f\"emitted {len(out)} ids (~{total/1e6:.1f}M tok); first {BUDGET/1e6:.0f}M tok = {nd_} docs\")\nprint(\"  register token mix in budget prefix:\", [round(c / max(1, cum), 3) for c in comp])\nprint(\"  wrote\", a.out)\n", "structuredPatch": [{"oldStart": 124, "oldLines": 8, "newStart": 124, "newLines": 36, "lines": [" for k in range(NREG):", "     ded[k] = [i for i in lists[k] if keep[i]]", " ", "+# ------------------------------------------- 4b. surface-convention gap fillers", "+# The target's encyclopedic quarter writes punctuation with a leading space (\" ,\", \" .\",", "+# \"@-@\") and its technical-Q&A quarter is HTML-marked-up (\"<p>\", \"</\", \"&gt;\").  Those are", "+# *token-level conventions*, not topics: a model that never sees them assigns them almost", "+# no probability, and they account for ~9% / ~7% of the tokens in those quarters.  DSIR", "+# ranks documents by average n-gram fit and so under-buys these rare-in-pool conventions,", "+# hence an explicit reserved slice of the budget for the documents that carry them.", "+surf = np.load(\"/workspace/work/surface.npy\")   # rows: [spaced-punct rate, markup rate]", "+pre, used_ex, used_nd = [], set(), set()", "+light = (nw >= 32) & (z[\"nonascii\"] <= 0.10) & (nw <= a.maxw)", "+for row, want in ((0, a.fill_sp), (1, a.fill_mk)):", "+    if want <= 0:", "+        continue", "+    got = 0", "+    for i in np.argsort(-surf[row]):", "+        if not light[i]:", "+            continue", "+        e, n = int(ex[i]), int(nd[i])", "+        if e in used_ex or n in used_nd:", "+            continue", "+        used_ex.add(e); used_nd.add(n)", "+        pre.append(int(i)); got += int(tokc[i])", "+        if got >= want:", "+            break", "+prefill = set(pre)", "+ded = [[i for i in l if i not in prefill] for l in ded]", "+", " # ---------------------------------------------------------------- 5. balanced interleave", "-out, ptr, tok = [], [0] * NREG, [0] * NREG", "+out, ptr, tok = [int(ids[i]) for i in pre], [0] * NREG, [0] * NREG", "+fill_tok = sum(int(tokc[i]) for i in pre)", " share = [float(x) for x in a.share.split(\",\")] if a.share else [1.0 / NREG] * NREG", " share = [s / sum(share) for s in share]", " if a.balanced:"]}], "userModified": false, "replaceAll": false}

160:    total = fill_tok
171:    total = fill_tok
190:print(f"  fillers: {len(pre)} docs / {fill_tok/1e6:.2f}M tok")
gate kept 151403/182016; dedup kept 138330
emitted 81345 ids (~48.0M tok); first 12M tok = 19960 docs
  register token mix in budget prefix: [0.239, 0.215, 0.211, 0.335]
  wrote work/s_fill.json

[stdout]
160:    total = fill_tok
171:    total = fill_tok
190:print(f"  fillers: {len(pre)} docs / {fill_tok/1e6:.2f}M tok")
gate kept 151403/182016; dedup kept 138330
emitted 81345 ids (~48.0M tok); first 12M tok = 19960 docs
  register token mix in budget prefix: [0.239, 0.215, 0.211, 0.335]
  wrote work/s_fill.json
[metadata] {"noOutputExpected": false}

fill1M: HELDOUT_PPL 304.725
  fill05M: HELDOUT_PPL 303.165
  fillsp2M: HELDOUT_PPL 312.681
3137

[stdout]
  fill1M: HELDOUT_PPL 304.725
  fill05M: HELDOUT_PPL 303.165
  fillsp2M: HELDOUT_PPL 312.681
3137
[metadata] {"noOutputExpected": false}

f025: HELDOUT_PPL 312.820
  f05_10: HELDOUT_PPL 303.449
  sh_r0: HELDOUT_PPL 308.351
  sh_news: HELDOUT_PPL 308.394
2887

[stdout]
  f025: HELDOUT_PPL 312.820
  f05_10: HELDOUT_PPL 303.449
  sh_r0: HELDOUT_PPL 308.351
  sh_news: HELDOUT_PPL 308.394
2887
[metadata] {"noOutputExpected": false}

fg: HELDOUT_PPL 305.487
  m64f: HELDOUT_PPL 304.312
  m16f: HELDOUT_PPL 309.051
  f0707: HELDOUT_PPL 303.747
2626

[stdout]
  fg: HELDOUT_PPL 305.487
  m64f: HELDOUT_PPL 304.312
  m16f: HELDOUT_PPL 309.051
  f0707: HELDOUT_PPL 303.747
2626
[metadata] {"noOutputExpected": false}

File created successfully at: /workspace/gapfill.py (file state is current in your context — no need to Read it back)

[metadata] {"type": "create", "filePath": "/workspace/gapfill.py", "content": "\"\"\"Refinement step: reallocate part of the budget to documents that supply target token\nmass the current selection under-covers.\n\nCriterion.  Let p_tgt be the equal-parts mixture unigram distribution of the target and\np_sel that of the current 12M-token selection prefix.  The excess cross-entropy a model\npays for a token the selection under-supplies is bounded by gap(t) = max(0, log p_tgt(t)\n- log p_sel(t)).  For a document d we score  g(d) = sum_t n_d(t) gap(t) / |d|  -- nats per\ntoken of *under-covered* target mass it would add.  We then swap the lowest-ranked tail of\nthe selection (REPL tokens) for the highest-g documents.  This generalises hand-picked\n\"surface convention\" fillers: whatever the selection is short of -- spaced punctuation,\nHTML markup, anything -- is what gets bought.\n\nUsage: python gapfill.py BASE_SEL OUT_SEL REPL_TOKENS\n\"\"\"\nimport json, pickle, sys\nimport numpy as np\n\nBASE, OUT, REPL = sys.argv[1], sys.argv[2], int(sys.argv[3])\nV, BUDGET = 50257, 12_000_000\ncache = pickle.load(open(\"/workspace/work/tokcache.pkl\", \"rb\"))\nz = np.load(\"/workspace/work/scores.npz\")\nids = z[\"id\"]\nnw = z[\"nw\"]\nlight = (nw >= 32) & (z[\"nonascii\"] <= 0.10)\nlightset = {int(i) for i, ok in zip(ids, light) if ok}\n\nbase = json.load(open(BASE))\nprefix, tot = [], 0\nfor i in base:\n    prefix.append(i); tot += len(cache[i]) + 1\n    if tot >= BUDGET:\n        break\n\nc_sel = np.zeros(V)\nfor i in prefix:\n    c_sel += np.bincount(cache[i].astype(np.int64), minlength=V)\nc_sel[50256] += len(prefix)\np_sel = (c_sel + 0.5) / (c_sel.sum() + 0.5 * V)\n\nd = np.load(\"/workspace/data/multi_dev.npy\"); q = len(d) // 4\np_tgt = np.zeros(V)\nfor k in range(4):\n    c = np.bincount(d[k * q:(k + 1) * q].astype(np.int64), minlength=V).astype(float)\n    p_tgt += 0.25 * c / c.sum()\ngap = np.maximum(0.0, np.log(p_tgt + 1e-12) - np.log(p_sel)).astype(np.float32)\nprint(f\"total under-covered target mass: {(p_tgt*gap).sum():.4f} nats/token\")\n\ninpre = set(prefix)\ng = np.full(len(ids), -1.0, dtype=np.float32)\nfor j, i in enumerate(ids):\n    ii = int(i)\n    if ii in inpre or ii not in lightset:\n        continue\n    t = cache[ii]\n    if len(t) < 32:\n        continue\n    g[j] = gap[t.astype(np.int64)].mean()\n\n# drop the tail of the base prefix (lowest priority = weakest register score) ...\ndrop, dtok = set(), 0\nfor i in reversed(prefix):\n    drop.add(i); dtok += len(cache[i]) + 1\n    if dtok >= REPL:\n        break\n# ... and buy the best gap-fillers instead\nadd, atok = [], 0\nfor j in np.argsort(-g):\n    if g[j] <= 0:\n        break\n    add.append(int(ids[j])); atok += len(cache[int(ids[j])]) + 1\n    if atok >= dtok:\n        break\nnewsel = add + [i for i in base if i not in drop]\nassert len(newsel) == len(set(newsel))\njson.dump(newsel, open(OUT, \"w\"))\nprint(f\"swapped {dtok/1e6:.2f}M tok of tail for {len(add)} gap-filler docs ({atok/1e6:.2f}M tok)\")\n", "structuredPatch": [], "originalFile": null, "userModified": false}

assert len(newsel) == len(set(newsel))
AssertionError
2576

[stdout]
    assert len(newsel) == len(set(newsel))
AssertionError
2576
[metadata] {"noOutputExpected": false}

total under-covered target mass: 0.4886 nats/token
swapped 1.20M tok of tail for 161 gap-filler docs (1.21M tok)
HELDOUT_PPL 331.634

[stdout]
total under-covered target mass: 0.4886 nats/token
swapped 1.20M tok of tail for 161 gap-filler docs (1.21M tok)
HELDOUT_PPL 331.634
[metadata] {"noOutputExpected": false}

File created successfully at: /workspace/surface_rates.py (file state is current in your context — no need to Read it back)

[metadata] {"type": "create", "filePath": "/workspace/surface_rates.py", "content": "\"\"\"Per-document rate of the two *surface conventions* the target uses but the raw pool\nmostly lacks, measured in GPT-2 token space:\n\n  row 0  spaced punctuation  -- tokens ' ,'(837) ' .'(764) ' @'(2488) ' ''(1058)\n         The encyclopedic quarter of the target is wikitext-style (\"Shackleton , in 1907 ,\n         ...  @-@ \") so ~9% of its tokens are of this form; pool median is 0.\n  row 1  HTML/markup         -- tokens '>'(29) '<'(27) '</'(3556) '&'(4775)\n         The technical-Q&A quarter is raw HTML (\"<p>\", \"<pre><code>\", \"&gt;\") so ~7% of\n         its tokens are of this form; pool median is 0.\n\nWrites work/surface.npy of shape (2, n_docs), aligned with work/scores.npz[\"id\"].\nRun: python surface_rates.py     (needs work/tokcache.pkl from score_pool_tok.py)\n\"\"\"\nimport pickle\nimport numpy as np\n\nSP = [837, 764, 2488, 1058]\nMK = [29, 27, 3556, 4775]\ncache = pickle.load(open(\"/workspace/work/tokcache.pkl\", \"rb\"))\nids = np.load(\"/workspace/work/scores.npz\")[\"id\"]\nout = np.zeros((2, len(ids)), dtype=np.float32)\nfor j, i in enumerate(ids):\n    t = cache[int(i)]\n    if len(t) == 0:\n        continue\n    c = np.bincount(t.astype(np.int64), minlength=50257)\n    out[0, j] = c[SP].sum() / len(t)\n    out[1, j] = c[MK].sum() / len(t)\nnp.save(\"/workspace/work/surface.npy\", out)\nprint(\"wrote work/surface.npy\", out.shape)\n", "structuredPatch": [], "originalFile": null, "userModified": false}