nano-gpt-data-curation
claude-opus-4-8
nano-gpt-data-curation — graded by the task's verifier against a measured baseline.
- backend
- modal
- variant
- claude claude-opus-4-8
- metric
- held-out multi-domain perplexity
- actions
- 67
- duration
- 1h 1m
- recorded
- 2026-08-04
run iddclm-nano-curation__u8hwsKZ
Everything claude-opus-4-8 did on this task, in order.
download JSON ↗claude-codeclaude-opus-4-867 turns35 narration67 action67 observation32 reasoning1 final
#26
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolBashInspect selected vs rejected docs (clean run)
args
{
"command": "cd /workspace; python3 -u -c \"\nimport json\nsel=json.load(open('submission/selection.json'))\nprint('n',len(sel),'unique',len(set(sel)),'range',min(sel),max(sel),flush=True)\nselset=set(sel)\ntexts={}\nwith open('data/pool.jsonl') as f:\n for line in f:\n r=json.loads(line); texts[r['id']]=r['text']\nprint('=== TOP-4 SELECTED ===',flush=True)\nfor i in sel[:4]:\n print('id',i,'::',texts[i][:280].replace(chr(10),' '),flush=True)\nprint('=== 4 REJECTED ===',flush=True)\nc=0\nfor i in range(182016):\n if i not in selset:\n print('id',i,'::',texts[i][:220].replace(chr(10),' '),flush=True); c+=1\n if c>=4: break\n\"",
"description": "Inspect selected vs rejected docs (clean run)"
}Bash result
n 31844 unique 31844 range 1 181948
=== TOP-4 SELECTED ===
id 131205 :: Index of /regional-patterns/assets/ Index of /regional-patterns/assets/ Name Last modified Size Description Parent Directory 30-Aug
id 135833 :: out<|endoftext|>Index of /wp-content/ Index of /wp-content/ Name Last modified Size Description Parent Directory 16-Nov-2018 02:36
id 158489 :: / Index of /wp-content/ Name Last modified Size Description Parent Directory 16-Nov-2018 02:36 - ai1wm-backups
id 162410 :: INGUBOX<|endoftext|>Index of / Index of / Name Last modified Size Description cgi-bin 21-Apr-2019 04:33 - Proudly S
=== 4 REJECTED ===
id 0 :: |Viewing Single Post From: Spoilers for the Week of February 11th| |Lil||Feb 1 2013, 09:58 AM| Don't care about Chloe/Taniel/Jen-Jen. Don't care about Sami, really, but hoping that we get some good "SAMANTHA GENE!!" Marl
id 2 :: A novel two-step immunotherapy approach has shown clinically beneficial responses in patients with advanced ovarian cancer. Following Lifestyle Recommendations Reduces Risk of Cancer Death People who follow the diet and
id 3 :: Free the Cans! Working Together to Reduce Waste In a blog about how people share, it’s worth the occasional reference to the bizarre ways that people DON’T SHARE. Is it safe to say we live in a society that places great
id 5 :: September 28, 2010 2010 Season - Bowman pulls down CCIW honor |Matt Bowman was named CCIW "Runner of the Week" after his fourth place finish at the Brissman-Lundeen Invitational in Rock Island, Illinois on September 24.|
[stdout]
n 31844 unique 31844 range 1 181948
=== TOP-4 SELECTED ===
id 131205 :: Index of /regional-patterns/assets/ Index of /regional-patterns/assets/ Name Last modified Size Description Parent Directory 30-Aug
id 135833 :: out<|endoftext|>Index of /wp-content/ Index of /wp-content/ Name Last modified Size Description Parent Directory 16-Nov-2018 02:36
id 158489 :: / Index of /wp-content/ Name Last modified Size Description Parent Directory 16-Nov-2018 02:36 - ai1wm-backups
id 162410 :: INGUBOX<|endoftext|>Index of / Index of / Name Last modified Size Description cgi-bin 21-Apr-2019 04:33 - Proudly S
=== 4 REJECTED ===
id 0 :: |Viewing Single Post From: Spoilers for the Week of February 11th| |Lil||Feb 1 2013, 09:58 AM| Don't care about Chloe/Taniel/Jen-Jen. Don't care about Sami, really, but hoping that we get some good "SAMANTHA GENE!!" Marl
id 2 :: A novel two-step immunotherapy approach has shown clinically beneficial responses in patients with advanced ovarian cancer. Following Lifestyle Recommendations Reduces Risk of Cancer Death People who follow the diet and
id 3 :: Free the Cans! Working Together to Reduce Waste In a blog about how people share, it’s worth the occasional reference to the bizarre ways that people DON’T SHARE. Is it safe to say we live in a society that places great
id 5 :: September 28, 2010 2010 Season - Bowman pulls down CCIW honor |Matt Bowman was named CCIW "Runner of the Week" after his fourth place finish at the Brissman-Lundeen Invitational in Rock Island, Illinois on September 24.|
[metadata] {"noOutputExpected": false}#27
narration
The top-scored docs are web-server "Index of /" autoindex pages — degenerate whitespace/boilerplate that games the mean-per-token score (the WikiText target inflates whitespace-token probability). This is a classic failure mode. Fix: (1) switch to presence-based scoring (mean log-odds over unique n-gram types, robust to repetition) and (2) add a letter-fraction gate to reject non-prose. Let me update the scoring section.
toolEdit
args
{
"replace_all": false,
"file_path": "/workspace/submission/curate.py",
"old_string": "EOS = 50256\nV = 50257 # GPT-2 vocab\nDB = 1 << 20 # hashed bigram buckets\nALPHA = 0.1 # additive smoothing\nMIN_TOK = 100 # length floor: drop fragments\nLAMBDA = 1.0 # weight of bigram term relative to unigram term\nCOVER_TOKENS = 30_000_000 # emit ids covering ~2.5x the 12M budget\n\nt_start = time.time()",
"new_string": "EOS = 50256\nV = 50257 # GPT-2 vocab\nDB = 1 << 20 # hashed bigram buckets\nALPHA = 0.1 # additive smoothing\nMIN_TOK = 100 # length floor: drop fragments\nMIN_ALPHA = 0.55 # letter-fraction floor: reject non-prose (dir listings, tables)\nLAMBDA = 0.5 # weight of bigram term relative to unigram term\nCOVER_TOKENS = 30_000_000 # emit ids covering ~2.5x the 12M budget\n\nimport re as _re\n_NONALPHA = _re.compile(r\"[^A-Za-z]\")\n\nt_start = time.time()"
}Edit result
The file /workspace/submission/curate.py has been updated successfully. (file state is current in your context — no need to Read it back)
[metadata] {"filePath": "/workspace/submission/curate.py", "oldString": "EOS = 50256\nV = 50257 # GPT-2 vocab\nDB = 1 << 20 # hashed bigram buckets\nALPHA = 0.1 # additive smoothing\nMIN_TOK = 100 # length floor: drop fragments\nLAMBDA = 1.0 # weight of bigram term relative to unigram term\nCOVER_TOKENS = 30_000_000 # emit ids covering ~2.5x the 12M budget\n\nt_start = time.time()", "newString": "EOS = 50256\nV = 50257 # GPT-2 vocab\nDB = 1 << 20 # hashed bigram buckets\nALPHA = 0.1 # additive smoothing\nMIN_TOK = 100 # length floor: drop fragments\nMIN_ALPHA = 0.55 # letter-fraction floor: reject non-prose (dir listings, tables)\nLAMBDA = 0.5 # weight of bigram term relative to unigram term\nCOVER_TOKENS = 30_000_000 # emit ids covering ~2.5x the 12M budget\n\nimport re as _re\n_NONALPHA = _re.compile(r\"[^A-Za-z]\")\n\nt_start = time.time()", "originalFile": "\"\"\"Curate a pretraining subset from a raw web pool by DSIR-style importance ranking.\n\nCriterion (stated, reproducible):\n Score each pool document by its mean per-token log importance weight\n s(doc) = (1/N) * sum_t [ log p_target(gram_t) - log p_pool(gram_t) ]\n under a bag-of-n-grams model (unigrams + hashed bigrams). p_target is estimated\n from the disclosed broad multi-domain HQ target (data/multi_dev.npy, GPT-2\n tokens); p_pool is estimated from the whole pool (the background distribution).\n Documents whose token distribution looks more like the HQ target than like the\n average web page score higher. We keep documents above a small length floor,\n drop exact duplicates, and emit ids in descending score order (best first)\n until we cover several times the 12M-token training budget.\n\nThis is Data Selection via Importance Resampling (Xie et al. 2023), greedy\ntop-k variant, with the importance target = the disclosed evaluation domain.\n\nNo labels, no internet: the only supervision is the disclosed target token stream.\n\"\"\"\nimport json, time, numpy as np\nfrom transformers import AutoTokenizer\n\nPOOL = \"/workspace/data/pool.jsonl\"\nDEV = \"/workspace/data/multi_dev.npy\"\nOUT = \"/workspace/submission/selection.json\"\n\nEOS = 50256\nV = 50257 # GPT-2 vocab\nDB = 1 << 20 # hashed bigram buckets\nALPHA = 0.1 # additive smoothing\nMIN_TOK = 100 # length floor: drop fragments\nLAMBDA = 1.0 # weight of bigram term relative to unigram term\nCOVER_TOKENS = 30_000_000 # emit ids covering ~2.5x the 12M budget\n\nt_start = time.time()\ntok = AutoTokenizer.from_pretrained(\"gpt2\")\n\n# ------------------------------------------------------------------ load pool\nids, texts = [], []\nwith open(POOL) as f:\n for line in f:\n r = json.loads(line)\n ids.append(r[\"id\"]); texts.append(r[\"text\"])\nids = np.asarray(ids, dtype=np.int64)\nndoc = len(ids)\nprint(f\"[{time.time()-t_start:.0f}s] loaded pool: {ndoc} docs\", flush=True)\n\n# ------------------------------------------------------ tokenize pool (batched)\ndocs = []\nB = 4000\nfor i in range(0, len(texts), B):\n enc = tok(texts[i:i+B], add_special_tokens=False).input_ids\n docs.extend(np.asarray(x, dtype=np.int32) for x in enc)\nlens = np.fromiter((len(d) for d in docs), dtype=np.int64, count=ndoc)\noffsets = np.zeros(ndoc + 1, dtype=np.int64)\nnp.cumsum(lens, out=offsets[1:])\nalltok = np.concatenate([d for d in docs]).astype(np.int64) # one contiguous stream\nprint(f\"[{time.time()-t_start:.0f}s] tokenized pool: {lens.sum()/1e6:.1f}M tokens\", flush=True)\n\ndef bigram_ids(a):\n \"\"\"Hashed bigram bucket ids for an int64 token array -> length len(a)-1.\"\"\"\n return ((a[:-1] * np.int64(1000003) + a[1:]) & np.int64(DB - 1))\n\n# ---------------------------------------------------- target (positive) counts\ndev = np.load(DEV).astype(np.int64)\npos_stream = dev[dev != EOS] # drop doc separators\npos_uni = np.bincount(pos_stream, minlength=V).astype(np.float64)\npos_big = np.bincount(bigram_ids(dev), minlength=DB).astype(np.float64) # EOS-cross negligible\n\n# ------------------------------------------------- pool (background) counts\n# Single bincount over the concatenated stream. Cross-doc boundary bigrams\n# (~ndoc of them out of ~193M) are negligible noise in the background estimate.\nneg_uni = np.bincount(alltok, minlength=V).astype(np.float64)\nall_big = bigram_ids(alltok) # length len(alltok)-1\nneg_big = np.bincount(all_big, minlength=DB).astype(np.float64)\nprint(f\"[{time.time()-t_start:.0f}s] built count tables\", flush=True)\n\n# --------------------------------------------------- log importance weights\nuni_w = (np.log(pos_uni + ALPHA) - np.log(pos_uni.sum() + ALPHA * V)) \\\n - (np.log(neg_uni + ALPHA) - np.log(neg_uni.sum() + ALPHA * V))\nbig_w = (np.log(pos_big + ALPHA) - np.log(pos_big.sum() + ALPHA * DB)) \\\n - (np.log(neg_big + ALPHA) - np.log(neg_big.sum() + ALPHA * DB))\n\n# ------------------------------------------------- vectorized per-doc scoring\n# Per-token unigram weight, summed per doc via reduceat over offsets.\nuni_tokw = uni_w[alltok]\ndoc_uni = np.add.reduceat(uni_tokw, offsets[:-1])\ndoc_uni[lens == 0] = 0.0\n\n# Bigram weight per position in the concatenated stream; zero-out the positions\n# that straddle a document boundary so no doc gets credit for a cross-doc bigram.\nbig_posw = big_w[all_big] # length len(alltok)-1\n# a bigram at global position p spans docs iff p+1 is a document start (in offsets)\nbnd = offsets[1:-1] # internal doc-start positions\nstraddle = bnd - 1 # bigram index that straddles\nstraddle = straddle[(straddle >= 0) & (straddle < big_posw.shape[0])]\nbig_posw[straddle] = 0.0\n# per-doc bigram sum: doc d owns bigram positions [offsets[d], offsets[d+1]-1)\nbig_starts = offsets[:-1].copy()\ndoc_big = np.zeros(ndoc, dtype=np.float64)\nvalid = lens > 1\n# cumulative sum trick for ranges\ncsum = np.concatenate([[0.0], np.cumsum(big_posw)])\n# doc d bigram positions run [offsets[d], offsets[d+1]-1) in big_posw indexing\nlo = offsets[:-1]\nhi = np.maximum(offsets[1:] - 1, offsets[:-1])\ndoc_big = csum[hi] - csum[lo]\n\nscores = np.where(lens >= MIN_TOK,\n (doc_uni + LAMBDA * doc_big) / np.maximum(lens, 1),\n -1e30)\nprint(f\"[{time.time()-t_start:.0f}s] scored docs\", flush=True)\n\n# ------------------------------------------- rank, dedup, cover the budget\norder = np.argsort(-scores)\nselection, seen, tot = [], set(), 0\nfor j in order:\n if scores[j] <= -1e29: # failed length floor\n break\n h = hash(texts[j]) # exact-duplicate guard\n if h in seen:\n continue\n seen.add(h)\n selection.append(int(ids[j]))\n tot += int(lens[j])\n if tot >= COVER_TOKENS:\n break\n\njson.dump(selection, open(OUT, \"w\"))\nvalid_scores = scores[scores > -1e29]\nprint(f\"[{time.time()-t_start:.0f}s] wrote {len(selection)} ids covering ~{tot/1e6:.1f}M tokens -> {OUT}\", flush=True)\nprint(f\"score p10={np.percentile(valid_scores,10):.3f} p50={np.percentile(valid_scores,50):.3f} \"\n f\"p90={np.percentile(valid_scores,90):.3f} top={scores[order[0]]:.3f} \"\n f\"kept_min={scores[order[len(selection)-1]]:.3f}\", flush=True)\n", "structuredPatch": [{"oldStart": 28, "oldLines": 9, "newStart": 28, "newLines": 13, "lines": [" DB = 1 << 20 # hashed bigram buckets", " ALPHA = 0.1 # additive smoothing", " MIN_TOK = 100 # length floor: drop fragments", "-LAMBDA = 1.0 # weight of bigram term relative to unigram term", "+MIN_ALPHA = 0.55 # letter-fraction floor: reject non-prose (dir listings, tables)", "+LAMBDA = 0.5 # weight of bigram term relative to unigram term", " COVER_TOKENS = 30_000_000 # emit ids covering ~2.5x the 12M budget", " ", "+import re as _re", "+_NONALPHA = _re.compile(r\"[^A-Za-z]\")", "+", " t_start = time.time()", " tok = AutoTokenizer.from_pretrained(\"gpt2\")", " "]}], "userModified": false, "replaceAll": false}#28
narration
Now replace the vectorized count-based scoring with presence-based scoring plus the letter-fraction gate:
toolEdit
args
{
"replace_all": false,
"file_path": "/workspace/submission/curate.py",
"old_string": "# ------------------------------------------------- vectorized per-doc scoring\n# Per-token unigram weight, summed per doc via reduceat over offsets.\nuni_tokw = uni_w[alltok]\ndoc_uni = np.add.reduceat(uni_tokw, offsets[:-1])\ndoc_uni[lens == 0] = 0.0\n\n# Bigram weight per position in the concatenated stream; zero-out the positions\n# that straddle a document boundary so no doc gets credit for a cross-doc bigram.\nbig_posw = big_w[all_big] # length len(alltok)-1\n# a bigram at global position p spans docs iff p+1 is a document start (in offsets)\nbnd = offsets[1:-1] # internal doc-start positions\nstraddle = bnd - 1 # bigram index that straddles\nstraddle = straddle[(straddle >= 0) & (straddle < big_posw.shape[0])]\nbig_posw[straddle] = 0.0\n# per-doc bigram sum: doc d owns bigram positions [offsets[d], offsets[d+1]-1)\nbig_starts = offsets[:-1].copy()\ndoc_big = np.zeros(ndoc, dtype=np.float64)\nvalid = lens > 1\n# cumulative sum trick for ranges\ncsum = np.concatenate([[0.0], np.cumsum(big_posw)])\n# doc d bigram positions run [offsets[d], offsets[d+1]-1) in big_posw indexing\nlo = offsets[:-1]\nhi = np.maximum(offsets[1:] - 1, offsets[:-1])\ndoc_big = csum[hi] - csum[lo]\n\nscores = np.where(lens >= MIN_TOK,\n (doc_uni + LAMBDA * doc_big) / np.maximum(lens, 1),\n -1e30)\nprint(f\"[{time.time()-t_start:.0f}s] scored docs\", flush=True)",
"new_string": "# ---------------------------- presence-based per-doc scoring (repetition-robust)\n# Score = mean log importance weight over the *set* of unigram (and bigram) types\n# the document contains. Counting types instead of tokens stops repetitive\n# boilerplate (whitespace runs, directory listings) from gaming a per-token mean.\n# A letter-fraction gate additionally rejects non-prose documents.\nscores = np.full(ndoc, -1e30, dtype=np.float64)\nfor i in range(ndoc):\n d = docs[i]; n = d.shape[0]\n if n < MIN_TOK:\n continue\n t = texts[i]\n if len(t) == 0 or len(_NONALPHA.sub(\"\", t)) / len(t) < MIN_ALPHA:\n continue # not enough letters -> not prose\n di = d.astype(np.int64)\n su = uni_w[np.unique(di)].mean()\n ub = np.unique((di[:-1] * np.int64(1000003) + di[1:]) & np.int64(DB - 1))\n sb = big_w[ub].mean()\n scores[i] = su + LAMBDA * sb\n if (i & 0x3FFF) == 0:\n print(f\"[{time.time()-t_start:.0f}s] scoring {i}/{ndoc}\", flush=True)\nprint(f\"[{time.time()-t_start:.0f}s] scored docs\", flush=True)"
}Edit result
The file /workspace/submission/curate.py has been updated successfully. (file state is current in your context — no need to Read it back)
[metadata] {"filePath": "/workspace/submission/curate.py", "oldString": "# ------------------------------------------------- vectorized per-doc scoring\n# Per-token unigram weight, summed per doc via reduceat over offsets.\nuni_tokw = uni_w[alltok]\ndoc_uni = np.add.reduceat(uni_tokw, offsets[:-1])\ndoc_uni[lens == 0] = 0.0\n\n# Bigram weight per position in the concatenated stream; zero-out the positions\n# that straddle a document boundary so no doc gets credit for a cross-doc bigram.\nbig_posw = big_w[all_big] # length len(alltok)-1\n# a bigram at global position p spans docs iff p+1 is a document start (in offsets)\nbnd = offsets[1:-1] # internal doc-start positions\nstraddle = bnd - 1 # bigram index that straddles\nstraddle = straddle[(straddle >= 0) & (straddle < big_posw.shape[0])]\nbig_posw[straddle] = 0.0\n# per-doc bigram sum: doc d owns bigram positions [offsets[d], offsets[d+1]-1)\nbig_starts = offsets[:-1].copy()\ndoc_big = np.zeros(ndoc, dtype=np.float64)\nvalid = lens > 1\n# cumulative sum trick for ranges\ncsum = np.concatenate([[0.0], np.cumsum(big_posw)])\n# doc d bigram positions run [offsets[d], offsets[d+1]-1) in big_posw indexing\nlo = offsets[:-1]\nhi = np.maximum(offsets[1:] - 1, offsets[:-1])\ndoc_big = csum[hi] - csum[lo]\n\nscores = np.where(lens >= MIN_TOK,\n (doc_uni + LAMBDA * doc_big) / np.maximum(lens, 1),\n -1e30)\nprint(f\"[{time.time()-t_start:.0f}s] scored docs\", flush=True)", "newString": "# ---------------------------- presence-based per-doc scoring (repetition-robust)\n# Score = mean log importance weight over the *set* of unigram (and bigram) types\n# the document contains. Counting types instead of tokens stops repetitive\n# boilerplate (whitespace runs, directory listings) from gaming a per-token mean.\n# A letter-fraction gate additionally rejects non-prose documents.\nscores = np.full(ndoc, -1e30, dtype=np.float64)\nfor i in range(ndoc):\n d = docs[i]; n = d.shape[0]\n if n < MIN_TOK:\n continue\n t = texts[i]\n if len(t) == 0 or len(_NONALPHA.sub(\"\", t)) / len(t) < MIN_ALPHA:\n continue # not enough letters -> not prose\n di = d.astype(np.int64)\n su = uni_w[np.unique(di)].mean()\n ub = np.unique((di[:-1] * np.int64(1000003) + di[1:]) & np.int64(DB - 1))\n sb = big_w[ub].mean()\n scores[i] = su + LAMBDA * sb\n if (i & 0x3FFF) == 0:\n print(f\"[{time.time()-t_start:.0f}s] scoring {i}/{ndoc}\", flush=True)\nprint(f\"[{time.time()-t_start:.0f}s] scored docs\", flush=True)", "originalFile": "\"\"\"Curate a pretraining subset from a raw web pool by DSIR-style importance ranking.\n\nCriterion (stated, reproducible):\n Score each pool document by its mean per-token log importance weight\n s(doc) = (1/N) * sum_t [ log p_target(gram_t) - log p_pool(gram_t) ]\n under a bag-of-n-grams model (unigrams + hashed bigrams). p_target is estimated\n from the disclosed broad multi-domain HQ target (data/multi_dev.npy, GPT-2\n tokens); p_pool is estimated from the whole pool (the background distribution).\n Documents whose token distribution looks more like the HQ target than like the\n average web page score higher. We keep documents above a small length floor,\n drop exact duplicates, and emit ids in descending score order (best first)\n until we cover several times the 12M-token training budget.\n\nThis is Data Selection via Importance Resampling (Xie et al. 2023), greedy\ntop-k variant, with the importance target = the disclosed evaluation domain.\n\nNo labels, no internet: the only supervision is the disclosed target token stream.\n\"\"\"\nimport json, time, numpy as np\nfrom transformers import AutoTokenizer\n\nPOOL = \"/workspace/data/pool.jsonl\"\nDEV = \"/workspace/data/multi_dev.npy\"\nOUT = \"/workspace/submission/selection.json\"\n\nEOS = 50256\nV = 50257 # GPT-2 vocab\nDB = 1 << 20 # hashed bigram buckets\nALPHA = 0.1 # additive smoothing\nMIN_TOK = 100 # length floor: drop fragments\nMIN_ALPHA = 0.55 # letter-fraction floor: reject non-prose (dir listings, tables)\nLAMBDA = 0.5 # weight of bigram term relative to unigram term\nCOVER_TOKENS = 30_000_000 # emit ids covering ~2.5x the 12M budget\n\nimport re as _re\n_NONALPHA = _re.compile(r\"[^A-Za-z]\")\n\nt_start = time.time()\ntok = AutoTokenizer.from_pretrained(\"gpt2\")\n\n# ------------------------------------------------------------------ load pool\nids, texts = [], []\nwith open(POOL) as f:\n for line in f:\n r = json.loads(line)\n ids.append(r[\"id\"]); texts.append(r[\"text\"])\nids = np.asarray(ids, dtype=np.int64)\nndoc = len(ids)\nprint(f\"[{time.time()-t_start:.0f}s] loaded pool: {ndoc} docs\", flush=True)\n\n# ------------------------------------------------------ tokenize pool (batched)\ndocs = []\nB = 4000\nfor i in range(0, len(texts), B):\n enc = tok(texts[i:i+B], add_special_tokens=False).input_ids\n docs.extend(np.asarray(x, dtype=np.int32) for x in enc)\nlens = np.fromiter((len(d) for d in docs), dtype=np.int64, count=ndoc)\noffsets = np.zeros(ndoc + 1, dtype=np.int64)\nnp.cumsum(lens, out=offsets[1:])\nalltok = np.concatenate([d for d in docs]).astype(np.int64) # one contiguous stream\nprint(f\"[{time.time()-t_start:.0f}s] tokenized pool: {lens.sum()/1e6:.1f}M tokens\", flush=True)\n\ndef bigram_ids(a):\n \"\"\"Hashed bigram bucket ids for an int64 token array -> length len(a)-1.\"\"\"\n return ((a[:-1] * np.int64(1000003) + a[1:]) & np.int64(DB - 1))\n\n# ---------------------------------------------------- target (positive) counts\ndev = np.load(DEV).astype(np.int64)\npos_stream = dev[dev != EOS] # drop doc separators\npos_uni = np.bincount(pos_stream, minlength=V).astype(np.float64)\npos_big = np.bincount(bigram_ids(dev), minlength=DB).astype(np.float64) # EOS-cross negligible\n\n# ------------------------------------------------- pool (background) counts\n# Single bincount over the concatenated stream. Cross-doc boundary bigrams\n# (~ndoc of them out of ~193M) are negligible noise in the background estimate.\nneg_uni = np.bincount(alltok, minlength=V).astype(np.float64)\nall_big = bigram_ids(alltok) # length len(alltok)-1\nneg_big = np.bincount(all_big, minlength=DB).astype(np.float64)\nprint(f\"[{time.time()-t_start:.0f}s] built count tables\", flush=True)\n\n# --------------------------------------------------- log importance weights\nuni_w = (np.log(pos_uni + ALPHA) - np.log(pos_uni.sum() + ALPHA * V)) \\\n - (np.log(neg_uni + ALPHA) - np.log(neg_uni.sum() + ALPHA * V))\nbig_w = (np.log(pos_big + ALPHA) - np.log(pos_big.sum() + ALPHA * DB)) \\\n - (np.log(neg_big + ALPHA) - np.log(neg_big.sum() + ALPHA * DB))\n\n# ------------------------------------------------- vectorized per-doc scoring\n# Per-token unigram weight, summed per doc via reduceat over offsets.\nuni_tokw = uni_w[alltok]\ndoc_uni = np.add.reduceat(uni_tokw, offsets[:-1])\ndoc_uni[lens == 0] = 0.0\n\n# Bigram weight per position in the concatenated stream; zero-out the positions\n# that straddle a document boundary so no doc gets credit for a cross-doc bigram.\nbig_posw = big_w[all_big] # length len(alltok)-1\n# a bigram at global position p spans docs iff p+1 is a document start (in offsets)\nbnd = offsets[1:-1] # internal doc-start positions\nstraddle = bnd - 1 # bigram index that straddles\nstraddle = straddle[(straddle >= 0) & (straddle < big_posw.shape[0])]\nbig_posw[straddle] = 0.0\n# per-doc bigram sum: doc d owns bigram positions [offsets[d], offsets[d+1]-1)\nbig_starts = offsets[:-1].copy()\ndoc_big = np.zeros(ndoc, dtype=np.float64)\nvalid = lens > 1\n# cumulative sum trick for ranges\ncsum = np.concatenate([[0.0], np.cumsum(big_posw)])\n# doc d bigram positions run [offsets[d], offsets[d+1]-1) in big_posw indexing\nlo = offsets[:-1]\nhi = np.maximum(offsets[1:] - 1, offsets[:-1])\ndoc_big = csum[hi] - csum[lo]\n\nscores = np.where(lens >= MIN_TOK,\n (doc_uni + LAMBDA * doc_big) / np.maximum(lens, 1),\n -1e30)\nprint(f\"[{time.time()-t_start:.0f}s] scored docs\", flush=True)\n\n# ------------------------------------------- rank, dedup, cover the budget\norder = np.argsort(-scores)\nselection, seen, tot = [], set(), 0\nfor j in order:\n if scores[j] <= -1e29: # failed length floor\n break\n h = hash(texts[j]) # exact-duplicate guard\n if h in seen:\n continue\n seen.add(h)\n selection.append(int(ids[j]))\n tot += int(lens[j])\n if tot >= COVER_TOKENS:\n break\n\njson.dump(selection, open(OUT, \"w\"))\nvalid_scores = scores[scores > -1e29]\nprint(f\"[{time.time()-t_start:.0f}s] wrote {len(selection)} ids covering ~{tot/1e6:.1f}M tokens -> {OUT}\", flush=True)\nprint(f\"score p10={np.percentile(valid_scores,10):.3f} p50={np.percentile(valid_scores,50):.3f} \"\n f\"p90={np.percentile(valid_scores,90):.3f} top={scores[order[0]]:.3f} \"\n f\"kept_min={scores[order[len(selection)-1]]:.3f}\", flush=True)\n", "structuredPatch": [{"oldStart": 84, "oldLines": 34, "newStart": 84, "newLines": 26, "lines": [" big_w = (np.log(pos_big + ALPHA) - np.log(pos_big.sum() + ALPHA * DB)) \\", " - (np.log(neg_big + ALPHA) - np.log(neg_big.sum() + ALPHA * DB))", " ", "-# ------------------------------------------------- vectorized per-doc scoring", "-# Per-token unigram weight, summed per doc via reduceat over offsets.", "-uni_tokw = uni_w[alltok]", "-doc_uni = np.add.reduceat(uni_tokw, offsets[:-1])", "-doc_uni[lens == 0] = 0.0", "-", "-# Bigram weight per position in the concatenated stream; zero-out the positions", "-# that straddle a document boundary so no doc gets credit for a cross-doc bigram.", "-big_posw = big_w[all_big] # length len(alltok)-1", "-# a bigram at global position p spans docs iff p+1 is a document start (in offsets)", "-bnd = offsets[1:-1] # internal doc-start positions", "-straddle = bnd - 1 # bigram index that straddles", "-straddle = straddle[(straddle >= 0) & (straddle < big_posw.shape[0])]", "-big_posw[straddle] = 0.0", "-# per-doc bigram sum: doc d owns bigram positions [offsets[d], offsets[d+1]-1)", "-big_starts = offsets[:-1].copy()", "-doc_big = np.zeros(ndoc, dtype=np.float64)", "-valid = lens > 1", "-# cumulative sum trick for ranges", "-csum = np.concatenate([[0.0], np.cumsum(big_posw)])", "-# doc d bigram positions run [offsets[d], offsets[d+1]-1) in big_posw indexing", "-lo = offsets[:-1]", "-hi = np.maximum(offsets[1:] - 1, offsets[:-1])", "-doc_big = csum[hi] - csum[lo]", "-", "-scores = np.where(lens >= MIN_TOK,", "- (doc_uni + LAMBDA * doc_big) / np.maximum(lens, 1),", "- -1e30)", "+# ---------------------------- presence-based per-doc scoring (repetition-robust)", "+# Score = mean log importance weight over the *set* of unigram (and bigram) types", "+# the document contains. Counting types instead of tokens stops repetitive", "+# boilerplate (whitespace runs, directory listings) from gaming a per-token mean.", "+# A letter-fraction gate additionally rejects non-prose documents.", "+scores = np.full(ndoc, -1e30, dtype=np.float64)", "+for i in range(ndoc):", "+ d = docs[i]; n = d.shape[0]", "+ if n < MIN_TOK:", "+ continue", "+ t = texts[i]", "+ if len(t) == 0 or len(_NONALPHA.sub(\"\", t)) / len(t) < MIN_ALPHA:", "+ continue # not enough letters -> not prose", "+ di = d.astype(np.int64)", "+ su = uni_w[np.unique(di)].mean()", "+ ub = np.unique((di[:-1] * np.int64(1000003) + di[1:]) & np.int64(DB - 1))", "+ sb = big_w[ub].mean()", "+ scores[i] = su + LAMBDA * sb", "+ if (i & 0x3FFF) == 0:", "+ print(f\"[{time.time()-t_start:.0f}s] scoring {i}/{ndoc}\", flush=True)", " print(f\"[{time.time()-t_start:.0f}s] scored docs\", flush=True)", " ", " # ------------------------------------------- rank, dedup, cover the budget"]}], "userModified": false, "replaceAll": false}#29
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolBashRun curate v2 with presence scoring + gate
args
{
"command": "cd /workspace; export TOKENIZERS_PARALLELISM=true; cat /workspace/.timer/remaining_secs; echo \"--- running curate.py v2 ---\"; time python3 submission/curate.py 2>&1 | grep -v \"sequence length\"",
"description": "Run curate v2 with presence scoring + gate"
}Bash result
Command did not complete within its 120s timeout and was moved to the background (ID: bmyob2zgd). Output is being written to: /tmp/claude-0/-workspace/4bb4d0fa-198e-4747-8902-b67961f6b67a/tasks/bmyob2zgd.output. You will be notified when it completes. To check interim output, use Read on that file path.
Session cwd remains /workspace; directory changes made by the backgrounded command do not apply to subsequent commands.
[metadata] {"noOutputExpected": false, "backgroundTaskId": "bmyob2zgd", "timedOutAfterMs": 120000, "backgroundCwdHint": "Session cwd remains /workspace; directory changes made by the backgrounded command do not apply to subsequent commands."}#30
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolBashCheck curate v2 progress
args
{
"command": "sleep 60; tail -5 /tmp/claude-0/-workspace/4bb4d0fa-198e-4747-8902-b67961f6b67a/tasks/bmyob2zgd.output",
"description": "Check curate v2 progress"
}Bash result
score p10=-2.097 p50=-0.783 p90=-0.263 top=0.512 kept_min=-0.412
real 2m26.856s
user 14m23.140s
sys 0m3.240s
[stdout]
score p10=-2.097 p50=-0.783 p90=-0.263 top=0.512 kept_min=-0.412
real 2m26.856s
user 14m23.140s
sys 0m3.240s
[metadata] {"noOutputExpected": false}#31
narration
Presence-based scoring done. Let me verify the top selections are now clean prose, and check the full output.
toolBashInspect top selections after fix
args
{
"command": "cat /tmp/claude-0/-workspace/4bb4d0fa-198e-4747-8902-b67961f6b67a/tasks/bmyob2zgd.output | grep -E \"wrote|score|scored docs\" | tail -5; echo \"==== inspect top ====\"; cd /workspace; python3 -u -c \"\nimport json\nsel=json.load(open('submission/selection.json'))\nprint('n',len(sel),'unique',len(set(sel)),'range',min(sel),max(sel),flush=True)\ntexts={}\nwith open('data/pool.jsonl') as f:\n for line in f:\n r=json.loads(line); texts[r['id']]=r['text']\nfor rank,i in enumerate(sel[:6]):\n print(f'--- rank {rank} id {i} ---',flush=True)\n print(texts[i][:260].replace(chr(10),' '),flush=True)\n\"",
"description": "Inspect top selections after fix"
}Bash result
[141s] scored docs
[141s] wrote 37428 ids covering ~30.0M tokens -> /workspace/submission/selection.json
score p10=-2.097 p50=-0.783 p90=-0.263 top=0.512 kept_min=-0.412
==== inspect top ====
n 37428 unique 37428 range 4 181924
--- rank 0 id 66736 ---
The Director General of Military Operations for the Pakistani Army spoke with his Indian counterpart via telephone on Wednesday – and denied New Delhi’s claim that Pakistani troops killed their two soldiers in a cross border raid. Pakistani military authoritie
--- rank 1 id 8287 ---
China has told North Korean companies operating in its territory to close down as it implements United Nations sanctions against the reclusive state. The companies will be shut by early January. Joint Chinese and North Korean ventures will also be forced to cl
--- rank 2 id 93562 ---
White House has dismissed North Korea's threat to pull out of the truce that ended the Korean War. U.S. officials say it is just another example of escalating rhetoric from Pyongyang. White House spokesman Ari Fleischer said the United States has heard it all
--- rank 3 id 105412 ---
redistributed.)<|endoftext|>Accompanied by a high-power delegation of military leadership, US Defence Secretary Chu-ck Hagel’s three-day visit to the capital is aimed at “nurturing” ties with India and pushing some of the pending deals. During this visit, Hag
--- rank 4 id 28954 ---
U.S. intelligence officials will meet with congressional lawmakers about North Korea's alleged nuclear cooperation with Syria Thursday on Capitol Hill. High-ranking U.S. officials are telling reporters the lawmakers will be shown video of North Koreans working
--- rank 5 id 75295 ---
It said the exclusive courts for child victims have to be set up in every district having over 100 cases under the POCSO Act. The bench said the Centre would have to appoint trained, sensitised prosecutors and support persons to deal with the POCSO cases and a
[stdout]
[141s] scored docs
[141s] wrote 37428 ids covering ~30.0M tokens -> /workspace/submission/selection.json
score p10=-2.097 p50=-0.783 p90=-0.263 top=0.512 kept_min=-0.412
==== inspect top ====
n 37428 unique 37428 range 4 181924
--- rank 0 id 66736 ---
The Director General of Military Operations for the Pakistani Army spoke with his Indian counterpart via telephone on Wednesday – and denied New Delhi’s claim that Pakistani troops killed their two soldiers in a cross border raid. Pakistani military authoritie
--- rank 1 id 8287 ---
China has told North Korean companies operating in its territory to close down as it implements United Nations sanctions against the reclusive state. The companies will be shut by early January. Joint Chinese and North Korean ventures will also be forced to cl
--- rank 2 id 93562 ---
White House has dismissed North Korea's threat to pull out of the truce that ended the Korean War. U.S. officials say it is just another example of escalating rhetoric from Pyongyang. White House spokesman Ari Fleischer said the United States has heard it all
--- rank 3 id 105412 ---
redistributed.)<|endoftext|>Accompanied by a high-power delegation of military leadership, US Defence Secretary Chu-ck Hagel’s three-day visit to the capital is aimed at “nurturing” ties with India and pushing some of the pending deals. During this visit, Hag
--- rank 4 id 28954 ---
U.S. intelligence officials will meet with congressional lawmakers about North Korea's alleged nuclear cooperation with Syria Thursday on Capitol Hill. High-ranking U.S. officials are telling reporters the lawmakers will be shown video of North Koreans working
--- rank 5 id 75295 ---
It said the exclusive courts for child victims have to be set up in every district having over 100 cases under the POCSO Act. The bench said the Centre would have to appoint trained, sensitised prosecutors and support persons to deal with the POCSO cases and a
[metadata] {"noOutputExpected": false}#32
narration
Top selections are now clean news/geopolitical prose (the "news" register) — the fix worked. 37,428 ids covering 30M tokens. Let me update tasks and validate with dev_eval, plus build a random baseline for comparison.
toolTaskUpdate
args
{
"taskId": "1",
"status": "completed"
}TaskUpdate result
Updated task #1 status
[metadata] {"success": true, "taskId": "1", "updatedFields": ["status"], "statusChange": {"from": "in_progress", "to": "completed"}}#33
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolTaskUpdate
args
{
"taskId": "2",
"status": "in_progress"
}TaskUpdate result
Updated task #2 status
[metadata] {"success": true, "taskId": "2", "updatedFields": ["status"], "statusChange": {"from": "pending", "to": "in_progress"}}#34
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolBashCreate random baseline selection
args
{
"command": "cd /workspace; python3 -c \"\nimport json, numpy as np\n# deterministic random baseline: shuffle all ids with fixed seed\nrng=np.random.default_rng(0)\nids=list(range(182016)); rng.shuffle(ids)\njson.dump([int(x) for x in ids[:60000]], open('/tmp/random_sel.json','w'))\nprint('random baseline written', 60000, 'ids')\n\"; cat /workspace/.timer/remaining_secs",
"description": "Create random baseline selection"
}Bash result
random baseline written 60000 ids
4148
[stdout]
random baseline written 60000 ids
4148
[metadata] {"noOutputExpected": false}#35
narration
Now run dev_eval on my curated selection (the key validation number):
toolBashRun dev_eval on curated selection
args
{
"command": "cd /workspace; export TOKENIZERS_PARALLELISM=false; echo \"=== dev_eval on CURATED selection ===\"; time bash dev_eval.sh /workspace/submission/selection.json 2>&1 | grep -vE \"sequence length|Token indices\"",
"description": "Run dev_eval on curated selection"
}Bash result
=== dev_eval on CURATED selection ===
packed 12000000 tokens from 37428 ids -> /tmp/dev_train.npy
HELDOUT_PPL 441.012 loss 6.0891 train_s 56 /tmp/dev_train.npy
(dev PPL on multi_dev; official uses a hidden disjoint sample of the same HQ domain)
real 1m46.208s
user 1m39.770s
sys 0m3.310s
[stdout]
=== dev_eval on CURATED selection ===
packed 12000000 tokens from 37428 ids -> /tmp/dev_train.npy
HELDOUT_PPL 441.012 loss 6.0891 train_s 56 /tmp/dev_train.npy
(dev PPL on multi_dev; official uses a hidden disjoint sample of the same HQ domain)
real 1m46.208s
user 1m39.770s
sys 0m3.310s
[metadata] {"noOutputExpected": false}#36
narration
Curated selection: PPL 441.0. Now the random baseline for comparison:
toolBashRun dev_eval on random baseline
args
{
"command": "cd /workspace; export TOKENIZERS_PARALLELISM=false; echo \"=== dev_eval on RANDOM baseline ===\"; time bash dev_eval.sh /tmp/random_sel.json 2>&1 | grep -vE \"sequence length|Token indices\"",
"description": "Run dev_eval on random baseline"
}Bash result
=== dev_eval on RANDOM baseline ===
packed 12000000 tokens from 60000 ids -> /tmp/dev_train.npy
HELDOUT_PPL 485.420 loss 6.1850 train_s 56 /tmp/dev_train.npy
(dev PPL on multi_dev; official uses a hidden disjoint sample of the same HQ domain)
real 1m45.556s
user 1m39.570s
sys 0m2.830s
[stdout]
=== dev_eval on RANDOM baseline ===
packed 12000000 tokens from 60000 ids -> /tmp/dev_train.npy
HELDOUT_PPL 485.420 loss 6.1850 train_s 56 /tmp/dev_train.npy
(dev PPL on multi_dev; official uses a hidden disjoint sample of the same HQ domain)
real 1m45.556s
user 1m39.570s
sys 0m2.830s
[metadata] {"noOutputExpected": false}#37
narration
Curated 441.0 vs random 485.4 — a real ~9% reduction. But the top docs looked news-heavy; the target is equal parts 4 registers, so register imbalance may be capping gains. Let me diagnose the register mix of my selection before deciding whether to rebalance.
toolBashDiagnose register composition of selection vs pool
args
{
"command": "cd /workspace; cat /workspace/.timer/remaining_secs; python3 -u -c \"\nimport json,re\nsel=set(json.load(open('submission/selection.json')))\nseltop=json.load(open('submission/selection.json'))[:12000] # roughly the budget-filling portion\ntexts={}\nwith open('data/pool.jsonl') as f:\n for line in f:\n r=json.loads(line); texts[r['id']]=r['text']\ndef reg(t):\n tl=t.lower()\n if ('<p>' in t or '</p>' in t or '<code>' in t or '</code>' in t or 'def ' in t or 'function ' in t or '{' in t and '}' in t and ';' in t):\n return 'code_qa'\n if re.search(r'\\b(reuters|said|told|according to|on (monday|tuesday|wednesday|thursday|friday|saturday|sunday))\\b', tl):\n return 'news'\n return 'prose/other'\nfrom collections import Counter\nc=Counter(reg(texts[i]) for i in seltop)\nprint('register mix of top-12k selected (budget-filling):', dict(c))\n# pool overall sample\nimport random; random.seed(0)\npoolsample=random.sample(range(182016),4000)\ncp=Counter(reg(texts[i]) for i in poolsample)\nprint('register mix of random pool sample:', dict(cp))\n\"",
"description": "Diagnose register composition of selection vs pool"
}Bash result
3888
register mix of top-12k selected (budget-filling): {'news': 7503, 'prose/other': 4268, 'code_qa': 229}
register mix of random pool sample: {'prose/other': 2880, 'news': 956, 'code_qa': 164}
[stdout]
3888
register mix of top-12k selected (budget-filling): {'news': 7503, 'prose/other': 4268, 'code_qa': 229}
register mix of random pool sample: {'prose/other': 2880, 'news': 956, 'code_qa': 164}
[metadata] {"noOutputExpected": false}#38
narration
My selection is 63% news, 2% code/Q&A — badly imbalanced vs the target's equal parts (25% each). The eval averages over all 4 registers equally, so my under-represented registers (code/Q&A, encyclopedic) drag the average up. Register-balanced DSIR should help materially. Let me first check the target's own register split to confirm I can estimate per-register weights.
toolBashCheck target register split sizes
args
{
"command": "cd /workspace; python3 -u -c \"\nimport numpy as np, re\nfrom transformers import AutoTokenizer\ntok=AutoTokenizer.from_pretrained('gpt2')\ndev=np.load('data/multi_dev.npy').astype(np.int64)\nEOS=50256\n# split into docs\nidx=np.where(dev==EOS)[0]\nbounds=[0]+ (idx+1).tolist()+[len(dev)]\ndocs=[dev[bounds[k]:bounds[k+1]] for k in range(len(bounds)-1)]\ndocs=[d[d!=EOS] for d in docs if (d!=EOS).sum()>20]\ndef reg(t):\n if '@-@' in t or '@,@' in t or '@.@' in t: return 'wiki'\n if ('<p>' in t or '<code>' in t or '</' in t or 'def ' in t or 'import ' in t or ('{' in t and '}' in t and ';' in t)): return 'code_qa'\n tl=t.lower()\n if re.search(r'\\b(reuters|said|told|according to|on (monday|tuesday|wednesday|thursday|friday|saturday|sunday))\\b',tl): return 'news'\n return 'web'\nfrom collections import Counter\nc=Counter(); toks=Counter()\nfor d in docs:\n t=tok.decode(d.tolist()); r=reg(t); c[r]+=1; toks[r]+=len(d)\nprint('target doc counts by register:', dict(c))\nprint('target token counts by register:', dict(toks))\nprint('total target docs', len(docs))\n\"",
"description": "Check target register split sizes"
}Bash result
target doc counts by register: {'wiki': 966, 'web': 828, 'news': 544, 'code_qa': 366}
target token counts by register: {'wiki': 163490, 'web': 206521, 'news': 354044, 'code_qa': 273221}
total target docs 2704
[stdout]
target doc counts by register: {'wiki': 966, 'web': 828, 'news': 544, 'code_qa': 366}
target token counts by register: {'wiki': 163490, 'web': 206521, 'news': 354044, 'code_qa': 273221}
total target docs 2704
[metadata] {"noOutputExpected": false}#39
narration
Each register has enough target data (366–966 docs) to estimate per-register weights. I'll rewrite curate.py as register-balanced DSIR: build a separate importance target per register, route each pool doc to its best-fit register, then interleave by equal token quota so the 12M budget is ~25% each. Let me write it.
toolWrite
args
{
"file_path": "/workspace/submission/curate.py",
"content": "\"\"\"Curate a pretraining subset from a raw web pool by register-balanced DSIR.\n\nThe disclosed evaluation target is a BROAD, multi-domain HQ English mix in equal\nparts: encyclopedic (Wikipedia), general web prose, news, and technical Q&A.\nHeld-out perplexity averages over all four registers, so a training set that\nover-weights one register (e.g. news, which is easiest to match) leaves the\nothers under-trained and drags the average up. We therefore select in two steps:\n\n (1) DSIR importance scoring, PER REGISTER. We split the disclosed target\n (data/multi_dev.npy, GPT-2 tokens) into four register buckets by surface\n signatures, and for each register r estimate a bag-of-n-grams importance\n weight\n w_r(gram) = log p_target_r(gram) - log p_pool(gram)\n (unigrams + hashed bigrams; p_pool from the whole pool as background).\n Each pool document is scored against every register using a PRESENCE-based\n mean over its unique n-gram types (repetition-robust), behind a length\n floor and a letter-fraction gate that reject fragments and non-prose\n (directory listings, tables). A document is routed to its best-fit\n register (argmax score).\n\n (2) Balanced quota fill. Within each register we rank routed documents by\n score, then interleave the four ranked lists by EQUAL token quota, so the\n first 12M tokens the trainer consumes are ~25% from each register --\n matching the equal-parts evaluation mixture.\n\nThis is Data Selection via Importance Resampling (Xie et al. 2023) with a\nper-domain target, plus classic quality gating. No labels, no internet: the only\nsupervision is the disclosed target token stream.\n\"\"\"\nimport json, time, re, numpy as np\nfrom transformers import AutoTokenizer\n\nPOOL = \"/workspace/data/pool.jsonl\"\nDEV = \"/workspace/data/multi_dev.npy\"\nOUT = \"/workspace/submission/selection.json\"\n\nEOS = 50256\nV = 50257 # GPT-2 vocab\nDB = 1 << 20 # hashed bigram buckets\nALPHA = 0.1 # additive smoothing\nMIN_TOK = 100 # length floor: drop fragments\nMIN_ALPHA = 0.55 # letter-fraction floor: reject non-prose\nLAMBDA = 0.5 # weight of bigram term relative to unigram term\nCOVER_TOKENS = 30_000_000 # emit ids covering ~2.5x the 12M budget\nREGISTERS = [\"wiki\", \"web\", \"news\", \"code_qa\"]\nQUOTA = {r: 0.25 for r in REGISTERS} # equal-parts target mixture\n\n_NONALPHA = re.compile(r\"[^A-Za-z]\")\n_NEWS = re.compile(r\"\\b(reuters|said|told|according to|\"\n r\"on (monday|tuesday|wednesday|thursday|friday|saturday|sunday))\\b\")\n\ndef register_of(t):\n \"\"\"Surface-signature register label for a piece of text.\"\"\"\n if \"@-@\" in t or \"@,@\" in t or \"@.@\" in t:\n return \"wiki\"\n if (\"<p>\" in t or \"<code>\" in t or \"</\" in t or \"def \" in t or \"import \" in t\n or (\"{\" in t and \"}\" in t and \";\" in t)):\n return \"code_qa\"\n if _NEWS.search(t.lower()):\n return \"news\"\n return \"web\"\n\ndef bigrams(a):\n return ((a[:-1] * np.int64(1000003) + a[1:]) & np.int64(DB - 1))\n\nt0 = time.time()\ntok = AutoTokenizer.from_pretrained(\"gpt2\")\n\n# ------------------------------------------------------------------ load pool\nids, texts = [], []\nwith open(POOL) as f:\n for line in f:\n r = json.loads(line)\n ids.append(r[\"id\"]); texts.append(r[\"text\"])\nids = np.asarray(ids, dtype=np.int64)\nndoc = len(ids)\nprint(f\"[{time.time()-t0:.0f}s] loaded pool: {ndoc} docs\", flush=True)\n\n# ------------------------------------------------------ tokenize pool (batched)\ndocs = []\nfor i in range(0, ndoc, 4000):\n enc = tok(texts[i:i+4000], add_special_tokens=False).input_ids\n docs.extend(np.asarray(x, dtype=np.int32) for x in enc)\nlens = np.fromiter((len(d) for d in docs), dtype=np.int64, count=ndoc)\nalltok = np.concatenate([d for d in docs]).astype(np.int64)\nprint(f\"[{time.time()-t0:.0f}s] tokenized pool: {lens.sum()/1e6:.1f}M tokens\", flush=True)\n\n# --------------------------------- pool background counts (shared denominator)\nneg_uni = np.bincount(alltok, minlength=V).astype(np.float64)\nneg_big = np.bincount(bigrams(alltok), minlength=DB).astype(np.float64)\nneg_uni_ln = np.log(neg_uni + ALPHA) - np.log(neg_uni.sum() + ALPHA * V)\nneg_big_ln = np.log(neg_big + ALPHA) - np.log(neg_big.sum() + ALPHA * DB)\ndel alltok\nprint(f\"[{time.time()-t0:.0f}s] built pool background\", flush=True)\n\n# ---------------------------- per-register target counts from disclosed target\ndev = np.load(DEV).astype(np.int64)\nsplits = np.where(dev == EOS)[0]\nb = [0] + (splits + 1).tolist() + [len(dev)]\nuni_w = {r: None for r in REGISTERS}\nbig_w = {r: None for r in REGISTERS}\npos_uni = {r: np.zeros(V, np.float64) for r in REGISTERS}\npos_big = {r: np.zeros(DB, np.float64) for r in REGISTERS}\nfor k in range(len(b) - 1):\n d = dev[b[k]:b[k+1]]\n d = d[d != EOS]\n if len(d) < 20:\n continue\n r = register_of(tok.decode(d.tolist()))\n pos_uni[r] += np.bincount(d, minlength=V)\n if len(d) > 1:\n pos_big[r] += np.bincount(bigrams(d), minlength=DB)\nfor r in REGISTERS:\n uni_w[r] = (np.log(pos_uni[r] + ALPHA) - np.log(pos_uni[r].sum() + ALPHA * V)) - neg_uni_ln\n big_w[r] = (np.log(pos_big[r] + ALPHA) - np.log(pos_big[r].sum() + ALPHA * DB)) - neg_big_ln\nprint(f\"[{time.time()-t0:.0f}s] built {len(REGISTERS)} register targets \"\n f\"(tok: {{{', '.join(f'{r}:{int(pos_uni[r].sum())}' for r in REGISTERS)}}})\", flush=True)\n\n# ------------------------ presence-based scoring of every doc vs every register\nbest_reg = np.full(ndoc, -1, dtype=np.int64)\nbest_score = np.full(ndoc, -1e30, dtype=np.float64)\nfor i in range(ndoc):\n d = docs[i]; n = d.shape[0]\n if n < MIN_TOK:\n continue\n t = texts[i]\n if len(t) == 0 or len(_NONALPHA.sub(\"\", t)) / len(t) < MIN_ALPHA:\n continue\n di = d.astype(np.int64)\n u = np.unique(di)\n ub = np.unique(bigrams(di)) if n > 1 else np.array([0], dtype=np.int64)\n bi, bs = -1, -1e30\n for ri, r in enumerate(REGISTERS):\n s = uni_w[r][u].mean() + LAMBDA * big_w[r][ub].mean()\n if s > bs:\n bs, bi = s, ri\n best_reg[i] = bi; best_score[i] = bs\n if (i & 0x3FFF) == 0:\n print(f\"[{time.time()-t0:.0f}s] scoring {i}/{ndoc}\", flush=True)\nprint(f\"[{time.time()-t0:.0f}s] scored docs\", flush=True)\n\n# ------------------------------------- rank within register, interleave by quota\nranked = {}\nfor ri, r in enumerate(REGISTERS):\n idx = np.where(best_reg == ri)[0]\n idx = idx[np.argsort(-best_score[idx])]\n ranked[r] = idx.tolist()\n print(f\" {r}: {len(idx)} docs routed\", flush=True)\n\nptr = {r: 0 for r in REGISTERS}\ncum = {r: 0 for r in REGISTERS}\nselection, seen, tot = [], set(), 0\nwhile tot < COVER_TOKENS:\n # pick the register that is furthest below its token quota and still has docs\n cand = [r for r in REGISTERS if ptr[r] < len(ranked[r])]\n if not cand:\n break\n r = min(cand, key=lambda r: cum[r] / QUOTA[r])\n j = ranked[r][ptr[r]]; ptr[r] += 1\n h = hash(texts[j])\n if h in seen:\n continue\n seen.add(h)\n selection.append(int(ids[j]))\n cum[r] += int(lens[j]); tot += int(lens[j])\n\njson.dump(selection, open(OUT, \"w\"))\nmix = {r: round(cum[r] / max(1, tot), 3) for r in REGISTERS}\nprint(f\"[{time.time()-t0:.0f}s] wrote {len(selection)} ids covering ~{tot/1e6:.1f}M tokens\", flush=True)\nprint(f\"register token mix of emitted list: {mix}\", flush=True)\n"
}Write result
The file /workspace/submission/curate.py has been updated successfully. (file state is current in your context — no need to Read it back)
[metadata] {"type": "update", "filePath": "/workspace/submission/curate.py", "content": "\"\"\"Curate a pretraining subset from a raw web pool by register-balanced DSIR.\n\nThe disclosed evaluation target is a BROAD, multi-domain HQ English mix in equal\nparts: encyclopedic (Wikipedia), general web prose, news, and technical Q&A.\nHeld-out perplexity averages over all four registers, so a training set that\nover-weights one register (e.g. news, which is easiest to match) leaves the\nothers under-trained and drags the average up. We therefore select in two steps:\n\n (1) DSIR importance scoring, PER REGISTER. We split the disclosed target\n (data/multi_dev.npy, GPT-2 tokens) into four register buckets by surface\n signatures, and for each register r estimate a bag-of-n-grams importance\n weight\n w_r(gram) = log p_target_r(gram) - log p_pool(gram)\n (unigrams + hashed bigrams; p_pool from the whole pool as background).\n Each pool document is scored against every register using a PRESENCE-based\n mean over its unique n-gram types (repetition-robust), behind a length\n floor and a letter-fraction gate that reject fragments and non-prose\n (directory listings, tables). A document is routed to its best-fit\n register (argmax score).\n\n (2) Balanced quota fill. Within each register we rank routed documents by\n score, then interleave the four ranked lists by EQUAL token quota, so the\n first 12M tokens the trainer consumes are ~25% from each register --\n matching the equal-parts evaluation mixture.\n\nThis is Data Selection via Importance Resampling (Xie et al. 2023) with a\nper-domain target, plus classic quality gating. No labels, no internet: the only\nsupervision is the disclosed target token stream.\n\"\"\"\nimport json, time, re, numpy as np\nfrom transformers import AutoTokenizer\n\nPOOL = \"/workspace/data/pool.jsonl\"\nDEV = \"/workspace/data/multi_dev.npy\"\nOUT = \"/workspace/submission/selection.json\"\n\nEOS = 50256\nV = 50257 # GPT-2 vocab\nDB = 1 << 20 # hashed bigram buckets\nALPHA = 0.1 # additive smoothing\nMIN_TOK = 100 # length floor: drop fragments\nMIN_ALPHA = 0.55 # letter-fraction floor: reject non-prose\nLAMBDA = 0.5 # weight of bigram term relative to unigram term\nCOVER_TOKENS = 30_000_000 # emit ids covering ~2.5x the 12M budget\nREGISTERS = [\"wiki\", \"web\", \"news\", \"code_qa\"]\nQUOTA = {r: 0.25 for r in REGISTERS} # equal-parts target mixture\n\n_NONALPHA = re.compile(r\"[^A-Za-z]\")\n_NEWS = re.compile(r\"\\b(reuters|said|told|according to|\"\n r\"on (monday|tuesday|wednesday|thursday|friday|saturday|sunday))\\b\")\n\ndef register_of(t):\n \"\"\"Surface-signature register label for a piece of text.\"\"\"\n if \"@-@\" in t or \"@,@\" in t or \"@.@\" in t:\n return \"wiki\"\n if (\"<p>\" in t or \"<code>\" in t or \"</\" in t or \"def \" in t or \"import \" in t\n or (\"{\" in t and \"}\" in t and \";\" in t)):\n return \"code_qa\"\n if _NEWS.search(t.lower()):\n return \"news\"\n return \"web\"\n\ndef bigrams(a):\n return ((a[:-1] * np.int64(1000003) + a[1:]) & np.int64(DB - 1))\n\nt0 = time.time()\ntok = AutoTokenizer.from_pretrained(\"gpt2\")\n\n# ------------------------------------------------------------------ load pool\nids, texts = [], []\nwith open(POOL) as f:\n for line in f:\n r = json.loads(line)\n ids.append(r[\"id\"]); texts.append(r[\"text\"])\nids = np.asarray(ids, dtype=np.int64)\nndoc = len(ids)\nprint(f\"[{time.time()-t0:.0f}s] loaded pool: {ndoc} docs\", flush=True)\n\n# ------------------------------------------------------ tokenize pool (batched)\ndocs = []\nfor i in range(0, ndoc, 4000):\n enc = tok(texts[i:i+4000], add_special_tokens=False).input_ids\n docs.extend(np.asarray(x, dtype=np.int32) for x in enc)\nlens = np.fromiter((len(d) for d in docs), dtype=np.int64, count=ndoc)\nalltok = np.concatenate([d for d in docs]).astype(np.int64)\nprint(f\"[{time.time()-t0:.0f}s] tokenized pool: {lens.sum()/1e6:.1f}M tokens\", flush=True)\n\n# --------------------------------- pool background counts (shared denominator)\nneg_uni = np.bincount(alltok, minlength=V).astype(np.float64)\nneg_big = np.bincount(bigrams(alltok), minlength=DB).astype(np.float64)\nneg_uni_ln = np.log(neg_uni + ALPHA) - np.log(neg_uni.sum() + ALPHA * V)\nneg_big_ln = np.log(neg_big + ALPHA) - np.log(neg_big.sum() + ALPHA * DB)\ndel alltok\nprint(f\"[{time.time()-t0:.0f}s] built pool background\", flush=True)\n\n# ---------------------------- per-register target counts from disclosed target\ndev = np.load(DEV).astype(np.int64)\nsplits = np.where(dev == EOS)[0]\nb = [0] + (splits + 1).tolist() + [len(dev)]\nuni_w = {r: None for r in REGISTERS}\nbig_w = {r: None for r in REGISTERS}\npos_uni = {r: np.zeros(V, np.float64) for r in REGISTERS}\npos_big = {r: np.zeros(DB, np.float64) for r in REGISTERS}\nfor k in range(len(b) - 1):\n d = dev[b[k]:b[k+1]]\n d = d[d != EOS]\n if len(d) < 20:\n continue\n r = register_of(tok.decode(d.tolist()))\n pos_uni[r] += np.bincount(d, minlength=V)\n if len(d) > 1:\n pos_big[r] += np.bincount(bigrams(d), minlength=DB)\nfor r in REGISTERS:\n uni_w[r] = (np.log(pos_uni[r] + ALPHA) - np.log(pos_uni[r].sum() + ALPHA * V)) - neg_uni_ln\n big_w[r] = (np.log(pos_big[r] + ALPHA) - np.log(pos_big[r].sum() + ALPHA * DB)) - neg_big_ln\nprint(f\"[{time.time()-t0:.0f}s] built {len(REGISTERS)} register targets \"\n f\"(tok: {{{', '.join(f'{r}:{int(pos_uni[r].sum())}' for r in REGISTERS)}}})\", flush=True)\n\n# ------------------------ presence-based scoring of every doc vs every register\nbest_reg = np.full(ndoc, -1, dtype=np.int64)\nbest_score = np.full(ndoc, -1e30, dtype=np.float64)\nfor i in range(ndoc):\n d = docs[i]; n = d.shape[0]\n if n < MIN_TOK:\n continue\n t = texts[i]\n if len(t) == 0 or len(_NONALPHA.sub(\"\", t)) / len(t) < MIN_ALPHA:\n continue\n di = d.astype(np.int64)\n u = np.unique(di)\n ub = np.unique(bigrams(di)) if n > 1 else np.array([0], dtype=np.int64)\n bi, bs = -1, -1e30\n for ri, r in enumerate(REGISTERS):\n s = uni_w[r][u].mean() + LAMBDA * big_w[r][ub].mean()\n if s > bs:\n bs, bi = s, ri\n best_reg[i] = bi; best_score[i] = bs\n if (i & 0x3FFF) == 0:\n print(f\"[{time.time()-t0:.0f}s] scoring {i}/{ndoc}\", flush=True)\nprint(f\"[{time.time()-t0:.0f}s] scored docs\", flush=True)\n\n# ------------------------------------- rank within register, interleave by quota\nranked = {}\nfor ri, r in enumerate(REGISTERS):\n idx = np.where(best_reg == ri)[0]\n idx = idx[np.argsort(-best_score[idx])]\n ranked[r] = idx.tolist()\n print(f\" {r}: {len(idx)} docs routed\", flush=True)\n\nptr = {r: 0 for r in REGISTERS}\ncum = {r: 0 for r in REGISTERS}\nselection, seen, tot = [], set(), 0\nwhile tot < COVER_TOKENS:\n # pick the register that is furthest below its token quota and still has docs\n cand = [r for r in REGISTERS if ptr[r] < len(ranked[r])]\n if not cand:\n break\n r = min(cand, key=lambda r: cum[r] / QUOTA[r])\n j = ranked[r][ptr[r]]; ptr[r] += 1\n h = hash(texts[j])\n if h in seen:\n continue\n seen.add(h)\n selection.append(int(ids[j]))\n cum[r] += int(lens[j]); tot += int(lens[j])\n\njson.dump(selection, open(OUT, \"w\"))\nmix = {r: round(cum[r] / max(1, tot), 3) for r in REGISTERS}\nprint(f\"[{time.time()-t0:.0f}s] wrote {len(selection)} ids covering ~{tot/1e6:.1f}M tokens\", flush=True)\nprint(f\"register token mix of emitted list: {mix}\", flush=True)\n", "structuredPatch": [{"oldStart": 1, "oldLines": 22, "newStart": 1, "newLines": 33, "lines": ["-\"\"\"Curate a pretraining subset from a raw web pool by DSIR-style importance ranking.", "+\"\"\"Curate a pretraining subset from a raw web pool by register-balanced DSIR.", " ", "-Criterion (stated, reproducible):", "- Score each pool document by its mean per-token log importance weight", "- s(doc) = (1/N) * sum_t [ log p_target(gram_t) - log p_pool(gram_t) ]", "- under a bag-of-n-grams model (unigrams + hashed bigrams). p_target is estimated", "- from the disclosed broad multi-domain HQ target (data/multi_dev.npy, GPT-2", "- tokens); p_pool is estimated from the whole pool (the background distribution).", "- Documents whose token distribution looks more like the HQ target than like the", "- average web page score higher. We keep documents above a small length floor,", "- drop exact duplicates, and emit ids in descending score order (best first)", "- until we cover several times the 12M-token training budget.", "+The disclosed evaluation target is a BROAD, multi-domain HQ English mix in equal", "+parts: encyclopedic (Wikipedia), general web prose, news, and technical Q&A.", "+Held-out perplexity averages over all four registers, so a training set that", "+over-weights one register (e.g. news, which is easiest to match) leaves the", "+others under-trained and drags the average up. We therefore select in two steps:", " ", "-This is Data Selection via Importance Resampling (Xie et al. 2023), greedy", "-top-k variant, with the importance target = the disclosed evaluation domain.", "+ (1) DSIR importance scoring, PER REGISTER. We split the disclosed target", "+ (data/multi_dev.npy, GPT-2 tokens) into four register buckets by surface", "+ signatures, and for each register r estimate a bag-of-n-grams importance", "+ weight", "+ w_r(gram) = log p_target_r(gram) - log p_pool(gram)", "+ (unigrams + hashed bigrams; p_pool from the whole pool as background).", "+ Each pool document is scored against every register using a PRESENCE-based", "+ mean over its unique n-gram types (repetition-robust), behind a length", "+ floor and a letter-fraction gate that reject fragments and non-prose", "+ (directory listings, tables). A document is routed to its best-fit", "+ register (argmax score).", " ", "-No labels, no internet: the only supervision is the disclosed target token stream.", "+ (2) Balanced quota fill. Within each register we rank routed documents by", "+ score, then interleave the four ranked lists by EQUAL token quota, so the", "+ first 12M tokens the trainer consumes are ~25% from each register --", "+ matching the equal-parts evaluation mixture.", "+", "+This is Data Selection via Importance Resampling (Xie et al. 2023) with a", "+per-domain target, plus classic quality gating. No labels, no internet: the only", "+supervision is the disclosed target token stream.", " \"\"\"", "-import json, time, numpy as np", "+import json, time, re, numpy as np", " from transformers import AutoTokenizer", " ", " POOL = \"/workspace/data/pool.jsonl\""]}, {"oldStart": 27, "oldLines": 15, "newStart": 38, "newLines": 32, "lines": [" V = 50257 # GPT-2 vocab", " DB = 1 << 20 # hashed bigram buckets", " ALPHA = 0.1 # additive smoothing", "-MIN_TOK = 100 # length floor: drop fragments", "-MIN_ALPHA = 0.55 # letter-fraction floor: reject non-prose (dir listings, tables)", "-LAMBDA = 0.5 # weight of bigram term relative to unigram term", "-COVER_TOKENS = 30_000_000 # emit ids covering ~2.5x the 12M budget", "+MIN_TOK = 100 # length floor: drop fragments", "+MIN_ALPHA = 0.55 # letter-fraction floor: reject non-prose", "+LAMBDA = 0.5 # weight of bigram term relative to unigram term", "+COVER_TOKENS = 30_000_000 # emit ids covering ~2.5x the 12M budget", "+REGISTERS = [\"wiki\", \"web\", \"news\", \"code_qa\"]", "+QUOTA = {r: 0.25 for r in REGISTERS} # equal-parts target mixture", " ", "-import re as _re", "-_NONALPHA = _re.compile(r\"[^A-Za-z]\")", "+_NONALPHA = re.compile(r\"[^A-Za-z]\")", "+_NEWS = re.compile(r\"\\b(reuters|said|told|according to|\"", "+ r\"on (monday|tuesday|wednesday|thursday|friday|saturday|sunday))\\b\")", " ", "-t_start = time.time()", "+def register_of(t):", "+ \"\"\"Surface-signature register label for a piece of text.\"\"\"", "+ if \"@-@\" in t or \"@,@\" in t or \"@.@\" in t:", "+ return \"wiki\"", "+ if (\"<p>\" in t or \"<code>\" in t or \"</\" in t or \"def \" in t or \"import \" in t", "+ or (\"{\" in t and \"}\" in t and \";\" in t)):", "+ return \"code_qa\"", "+ if _NEWS.search(t.lower()):", "+ return \"news\"", "+ return \"web\"", "+", "+def bigrams(a):", "+ return ((a[:-1] * np.int64(1000003) + a[1:]) & np.int64(DB - 1))", "+", "+t0 = time.time()", " tok = AutoTokenizer.from_pretrained(\"gpt2\")", " ", " # ------------------------------------------------------------------ load pool"]}, {"oldStart": 46, "oldLines": 84, "newStart": 74, "newLines": 97, "lines": [" ids.append(r[\"id\"]); texts.append(r[\"text\"])", " ids = np.asarray(ids, dtype=np.int64)", " ndoc = len(ids)", "-print(f\"[{time.time()-t_start:.0f}s] loaded pool: {ndoc} docs\", flush=True)", "+print(f\"[{time.time()-t0:.0f}s] loaded pool: {ndoc} docs\", flush=True)", " ", " # ------------------------------------------------------ tokenize pool (batched)", " docs = []", "-B = 4000", "-for i in range(0, len(texts), B):", "- enc = tok(texts[i:i+B], add_special_tokens=False).input_ids", "+for i in range(0, ndoc, 4000):", "+ enc = tok(texts[i:i+4000], add_special_tokens=False).input_ids", " docs.extend(np.asarray(x, dtype=np.int32) for x in enc)", " lens = np.fromiter((len(d) for d in docs), dtype=np.int64, count=ndoc)", "-offsets = np.zeros(ndoc + 1, dtype=np.int64)", "-np.cumsum(lens, out=offsets[1:])", "-alltok = np.concatenate([d for d in docs]).astype(np.int64) # one contiguous stream", "-print(f\"[{time.time()-t_start:.0f}s] tokenized pool: {lens.sum()/1e6:.1f}M tokens\", flush=True)", "+alltok = np.concatenate([d for d in docs]).astype(np.int64)", "+print(f\"[{time.time()-t0:.0f}s] tokenized pool: {lens.sum()/1e6:.1f}M tokens\", flush=True)", " ", "-def bigram_ids(a):", "- \"\"\"Hashed bigram bucket ids for an int64 token array -> length len(a)-1.\"\"\"", "- return ((a[:-1] * np.int64(1000003) + a[1:]) & np.int64(DB - 1))", "+# --------------------------------- pool background counts (shared denominator)", "+neg_uni = np.bincount(alltok, minlength=V).astype(np.float64)", "+neg_big = np.bincount(bigrams(alltok), minlength=DB).astype(np.float64)", "+neg_uni_ln = np.log(neg_uni + ALPHA) - np.log(neg_uni.sum() + ALPHA * V)", "+neg_big_ln = np.log(neg_big + ALPHA) - np.log(neg_big.sum() + ALPHA * DB)", "+del alltok", "+print(f\"[{time.time()-t0:.0f}s] built pool background\", flush=True)", " ", "-# ---------------------------------------------------- target (positive) counts", "+# ---------------------------- per-register target counts from disclosed target", " dev = np.load(DEV).astype(np.int64)", "-pos_stream = dev[dev != EOS] # drop doc separators", "-pos_uni = np.bincount(pos_stream, minlength=V).astype(np.float64)", "-pos_big = np.bincount(bigram_ids(dev), minlength=DB).astype(np.float64) # EOS-cross negligible", "+splits = np.where(dev == EOS)[0]", "+b = [0] + (splits + 1).tolist() + [len(dev)]", "+uni_w = {r: None for r in REGISTERS}", "+big_w = {r: None for r in REGISTERS}", "+pos_uni = {r: np.zeros(V, np.float64) for r in REGISTERS}", "+pos_big = {r: np.zeros(DB, np.float64) for r in REGISTERS}", "+for k in range(len(b) - 1):", "+ d = dev[b[k]:b[k+1]]", "+ d = d[d != EOS]", "+ if len(d) < 20:", "+ continue", "+ r = register_of(tok.decode(d.tolist()))", "+ pos_uni[r] += np.bincount(d, minlength=V)", "+ if len(d) > 1:", "+ pos_big[r] += np.bincount(bigrams(d), minlength=DB)", "+for r in REGISTERS:", "+ uni_w[r] = (np.log(pos_uni[r] + ALPHA) - np.log(pos_uni[r].sum() + ALPHA * V)) - neg_uni_ln", "+ big_w[r] = (np.log(pos_big[r] + ALPHA) - np.log(pos_big[r].sum() + ALPHA * DB)) - neg_big_ln", "+print(f\"[{time.time()-t0:.0f}s] built {len(REGISTERS)} register targets \"", "+ f\"(tok: {{{', '.join(f'{r}:{int(pos_uni[r].sum())}' for r in REGISTERS)}}})\", flush=True)", " ", "-# ------------------------------------------------- pool (background) counts", "-# Single bincount over the concatenated stream. Cross-doc boundary bigrams", "-# (~ndoc of them out of ~193M) are negligible noise in the background estimate.", "-neg_uni = np.bincount(alltok, minlength=V).astype(np.float64)", "-all_big = bigram_ids(alltok) # length len(alltok)-1", "-neg_big = np.bincount(all_big, minlength=DB).astype(np.float64)", "-print(f\"[{time.time()-t_start:.0f}s] built count tables\", flush=True)", "-", "-# --------------------------------------------------- log importance weights", "-uni_w = (np.log(pos_uni + ALPHA) - np.log(pos_uni.sum() + ALPHA * V)) \\", "- - (np.log(neg_uni + ALPHA) - np.log(neg_uni.sum() + ALPHA * V))", "-big_w = (np.log(pos_big + ALPHA) - np.log(pos_big.sum() + ALPHA * DB)) \\", "- - (np.log(neg_big + ALPHA) - np.log(neg_big.sum() + ALPHA * DB))", "-", "-# ---------------------------- presence-based per-doc scoring (repetition-robust)", "-# Score = mean log importance weight over the *set* of unigram (and bigram) types", "-# the document contains. Counting types instead of tokens stops repetitive", "-# boilerplate (whitespace runs, directory listings) from gaming a per-token mean.", "-# A letter-fraction gate additionally rejects non-prose documents.", "-scores = np.full(ndoc, -1e30, dtype=np.float64)", "+# ------------------------ presence-based scoring of every doc vs every register", "+best_reg = np.full(ndoc, -1, dtype=np.int64)", "+best_score = np.full(ndoc, -1e30, dtype=np.float64)", " for i in range(ndoc):", " d = docs[i]; n = d.shape[0]", " if n < MIN_TOK:", " continue", " t = texts[i]", " if len(t) == 0 or len(_NONALPHA.sub(\"\", t)) / len(t) < MIN_ALPHA:", "- continue # not enough letters -> not prose", "+ continue", " di = d.astype(np.int64)", "- su = uni_w[np.unique(di)].mean()", "- ub = np.unique((di[:-1] * np.int64(1000003) + di[1:]) & np.int64(DB - 1))", "- sb = big_w[ub].mean()", "- scores[i] = su + LAMBDA * sb", "+ u = np.unique(di)", "+ ub = np.unique(bigrams(di)) if n > 1 else np.array([0], dtype=np.int64)", "+ bi, bs = -1, -1e30", "+ for ri, r in enumerate(REGISTERS):", "+ s = uni_w[r][u].mean() + LAMBDA * big_w[r][ub].mean()", "+ if s > bs:", "+ bs, bi = s, ri", "+ best_reg[i] = bi; best_score[i] = bs", " if (i & 0x3FFF) == 0:", "- print(f\"[{time.time()-t_start:.0f}s] scoring {i}/{ndoc}\", flush=True)", "-print(f\"[{time.time()-t_start:.0f}s] scored docs\", flush=True)", "+ print(f\"[{time.time()-t0:.0f}s] scoring {i}/{ndoc}\", flush=True)", "+print(f\"[{time.time()-t0:.0f}s] scored docs\", flush=True)", " ", "-# ------------------------------------------- rank, dedup, cover the budget", "-order = np.argsort(-scores)", "+# ------------------------------------- rank within register, interleave by quota", "+ranked = {}", "+for ri, r in enumerate(REGISTERS):", "+ idx = np.where(best_reg == ri)[0]", "+ idx = idx[np.argsort(-best_score[idx])]", "+ ranked[r] = idx.tolist()", "+ print(f\" {r}: {len(idx)} docs routed\", flush=True)", "+", "+ptr = {r: 0 for r in REGISTERS}", "+cum = {r: 0 for r in REGISTERS}", " selection, seen, tot = [], set(), 0", "-for j in order:", "- if scores[j] <= -1e29: # failed length floor", "+while tot < COVER_TOKENS:", "+ # pick the register that is furthest below its token quota and still has docs", "+ cand = [r for r in REGISTERS if ptr[r] < len(ranked[r])]", "+ if not cand:", " break", "- h = hash(texts[j]) # exact-duplicate guard", "+ r = min(cand, key=lambda r: cum[r] / QUOTA[r])", "+ j = ranked[r][ptr[r]]; ptr[r] += 1", "+ h = hash(texts[j])", " if h in seen:", " continue", " seen.add(h)", " selection.append(int(ids[j]))", "- tot += int(lens[j])", "- if tot >= COVER_TOKENS:", "- break", "+ cum[r] += int(lens[j]); tot += int(lens[j])", " ", " json.dump(selection, open(OUT, \"w\"))", "-valid_scores = scores[scores > -1e29]", "-print(f\"[{time.time()-t_start:.0f}s] wrote {len(selection)} ids covering ~{tot/1e6:.1f}M tokens -> {OUT}\", flush=True)", "-print(f\"score p10={np.percentile(valid_scores,10):.3f} p50={np.percentile(valid_scores,50):.3f} \"", "- f\"p90={np.percentile(valid_scores,90):.3f} top={scores[order[0]]:.3f} \"", "- f\"kept_min={scores[order[len(selection)-1]]:.3f}\", flush=True)", "+mix = {r: round(cum[r] / max(1, tot), 3) for r in REGISTERS}", "+print(f\"[{time.time()-t0:.0f}s] wrote {len(selection)} ids covering ~{tot/1e6:.1f}M tokens\", flush=True)", "+print(f\"register token mix of emitted list: {mix}\", flush=True)"]}], "originalFile": "\"\"\"Curate a pretraining subset from a raw web pool by DSIR-style importance ranking.\n\nCriterion (stated, reproducible):\n Score each pool document by its mean per-token log importance weight\n s(doc) = (1/N) * sum_t [ log p_target(gram_t) - log p_pool(gram_t) ]\n under a bag-of-n-grams model (unigrams + hashed bigrams). p_target is estimated\n from the disclosed broad multi-domain HQ target (data/multi_dev.npy, GPT-2\n tokens); p_pool is estimated from the whole pool (the background distribution).\n Documents whose token distribution looks more like the HQ target than like the\n average web page score higher. We keep documents above a small length floor,\n drop exact duplicates, and emit ids in descending score order (best first)\n until we cover several times the 12M-token training budget.\n\nThis is Data Selection via Importance Resampling (Xie et al. 2023), greedy\ntop-k variant, with the importance target = the disclosed evaluation domain.\n\nNo labels, no internet: the only supervision is the disclosed target token stream.\n\"\"\"\nimport json, time, numpy as np\nfrom transformers import AutoTokenizer\n\nPOOL = \"/workspace/data/pool.jsonl\"\nDEV = \"/workspace/data/multi_dev.npy\"\nOUT = \"/workspace/submission/selection.json\"\n\nEOS = 50256\nV = 50257 # GPT-2 vocab\nDB = 1 << 20 # hashed bigram buckets\nALPHA = 0.1 # additive smoothing\nMIN_TOK = 100 # length floor: drop fragments\nMIN_ALPHA = 0.55 # letter-fraction floor: reject non-prose (dir listings, tables)\nLAMBDA = 0.5 # weight of bigram term relative to unigram term\nCOVER_TOKENS = 30_000_000 # emit ids covering ~2.5x the 12M budget\n\nimport re as _re\n_NONALPHA = _re.compile(r\"[^A-Za-z]\")\n\nt_start = time.time()\ntok = AutoTokenizer.from_pretrained(\"gpt2\")\n\n# ------------------------------------------------------------------ load pool\nids, texts = [], []\nwith open(POOL) as f:\n for line in f:\n r = json.loads(line)\n ids.append(r[\"id\"]); texts.append(r[\"text\"])\nids = np.asarray(ids, dtype=np.int64)\nndoc = len(ids)\nprint(f\"[{time.time()-t_start:.0f}s] loaded pool: {ndoc} docs\", flush=True)\n\n# ------------------------------------------------------ tokenize pool (batched)\ndocs = []\nB = 4000\nfor i in range(0, len(texts), B):\n enc = tok(texts[i:i+B], add_special_tokens=False).input_ids\n docs.extend(np.asarray(x, dtype=np.int32) for x in enc)\nlens = np.fromiter((len(d) for d in docs), dtype=np.int64, count=ndoc)\noffsets = np.zeros(ndoc + 1, dtype=np.int64)\nnp.cumsum(lens, out=offsets[1:])\nalltok = np.concatenate([d for d in docs]).astype(np.int64) # one contiguous stream\nprint(f\"[{time.time()-t_start:.0f}s] tokenized pool: {lens.sum()/1e6:.1f}M tokens\", flush=True)\n\ndef bigram_ids(a):\n \"\"\"Hashed bigram bucket ids for an int64 token array -> length len(a)-1.\"\"\"\n return ((a[:-1] * np.int64(1000003) + a[1:]) & np.int64(DB - 1))\n\n# ---------------------------------------------------- target (positive) counts\ndev = np.load(DEV).astype(np.int64)\npos_stream = dev[dev != EOS] # drop doc separators\npos_uni = np.bincount(pos_stream, minlength=V).astype(np.float64)\npos_big = np.bincount(bigram_ids(dev), minlength=DB).astype(np.float64) # EOS-cross negligible\n\n# ------------------------------------------------- pool (background) counts\n# Single bincount over the concatenated stream. Cross-doc boundary bigrams\n# (~ndoc of them out of ~193M) are negligible noise in the background estimate.\nneg_uni = np.bincount(alltok, minlength=V).astype(np.float64)\nall_big = bigram_ids(alltok) # length len(alltok)-1\nneg_big = np.bincount(all_big, minlength=DB).astype(np.float64)\nprint(f\"[{time.time()-t_start:.0f}s] built count tables\", flush=True)\n\n# --------------------------------------------------- log importance weights\nuni_w = (np.log(pos_uni + ALPHA) - np.log(pos_uni.sum() + ALPHA * V)) \\\n - (np.log(neg_uni + ALPHA) - np.log(neg_uni.sum() + ALPHA * V))\nbig_w = (np.log(pos_big + ALPHA) - np.log(pos_big.sum() + ALPHA * DB)) \\\n - (np.log(neg_big + ALPHA) - np.log(neg_big.sum() + ALPHA * DB))\n\n# ---------------------------- presence-based per-doc scoring (repetition-robust)\n# Score = mean log importance weight over the *set* of unigram (and bigram) types\n# the document contains. Counting types instead of tokens stops repetitive\n# boilerplate (whitespace runs, directory listings) from gaming a per-token mean.\n# A letter-fraction gate additionally rejects non-prose documents.\nscores = np.full(ndoc, -1e30, dtype=np.float64)\nfor i in range(ndoc):\n d = docs[i]; n = d.shape[0]\n if n < MIN_TOK:\n continue\n t = texts[i]\n if len(t) == 0 or len(_NONALPHA.sub(\"\", t)) / len(t) < MIN_ALPHA:\n continue # not enough letters -> not prose\n di = d.astype(np.int64)\n su = uni_w[np.unique(di)].mean()\n ub = np.unique((di[:-1] * np.int64(1000003) + di[1:]) & np.int64(DB - 1))\n sb = big_w[ub].mean()\n scores[i] = su + LAMBDA * sb\n if (i & 0x3FFF) == 0:\n print(f\"[{time.time()-t_start:.0f}s] scoring {i}/{ndoc}\", flush=True)\nprint(f\"[{time.time()-t_start:.0f}s] scored docs\", flush=True)\n\n# ------------------------------------------- rank, dedup, cover the budget\norder = np.argsort(-scores)\nselection, seen, tot = [], set(), 0\nfor j in order:\n if scores[j] <= -1e29: # failed length floor\n break\n h = hash(texts[j]) # exact-duplicate guard\n if h in seen:\n continue\n seen.add(h)\n selection.append(int(ids[j]))\n tot += int(lens[j])\n if tot >= COVER_TOKENS:\n break\n\njson.dump(selection, open(OUT, \"w\"))\nvalid_scores = scores[scores > -1e29]\nprint(f\"[{time.time()-t_start:.0f}s] wrote {len(selection)} ids covering ~{tot/1e6:.1f}M tokens -> {OUT}\", flush=True)\nprint(f\"score p10={np.percentile(valid_scores,10):.3f} p50={np.percentile(valid_scores,50):.3f} \"\n f\"p90={np.percentile(valid_scores,90):.3f} top={scores[order[0]]:.3f} \"\n f\"kept_min={scores[order[len(selection)-1]]:.3f}\", flush=True)\n", "userModified": false}#40
narration
Let me back up the current (unbalanced, 441 PPL) selection before overwriting, then run the balanced version.
toolBashRun register-balanced curate
args
{
"command": "cd /workspace; cp submission/selection.json /tmp/selection_unbalanced_441.json; cat /workspace/.timer/remaining_secs; export TOKENIZERS_PARALLELISM=true; echo \"--- running balanced curate.py ---\"; time python3 submission/curate.py 2>&1 | grep -vE \"sequence length|Token indices\"",
"description": "Run register-balanced curate"
}Bash result
Command did not complete within its 120s timeout and was moved to the background (ID: bn9px52je). Output is being written to: /tmp/claude-0/-workspace/4bb4d0fa-198e-4747-8902-b67961f6b67a/tasks/bn9px52je.output. You will be notified when it completes. To check interim output, use Read on that file path.
Session cwd remains /workspace; directory changes made by the backgrounded command do not apply to subsequent commands.
[metadata] {"noOutputExpected": false, "backgroundTaskId": "bn9px52je", "timedOutAfterMs": 120000, "backgroundCwdHint": "Session cwd remains /workspace; directory changes made by the backgrounded command do not apply to subsequent commands."}#41
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolBashCheck balanced curate progress
args
{
"command": "sleep 75; tail -12 /tmp/claude-0/-workspace/4bb4d0fa-198e-4747-8902-b67961f6b67a/tasks/bn9px52je.output",
"description": "Check balanced curate progress"
}Bash result
[152s] scoring 180224/182016
[152s] scored docs
wiki: 105 docs routed
web: 40067 docs routed
news: 124211 docs routed
code_qa: 6767 docs routed
[153s] wrote 29979 ids covering ~30.0M tokens
register token mix of emitted list: {'wiki': 0.015, 'web': 0.328, 'news': 0.328, 'code_qa': 0.329}
real 2m38.087s
user 15m6.010s
sys 0m2.920s
[stdout]
[152s] scoring 180224/182016
[152s] scored docs
wiki: 105 docs routed
web: 40067 docs routed
news: 124211 docs routed
code_qa: 6767 docs routed
[153s] wrote 29979 ids covering ~30.0M tokens
register token mix of emitted list: {'wiki': 0.015, 'web': 0.328, 'news': 0.328, 'code_qa': 0.329}
real 2m38.087s
user 15m6.010s
sys 0m2.920s
[metadata] {"noOutputExpected": false}#42
narration
The balanced version starved wiki (1.5%) — argmax routing fails for wiki because the target's WikiText formatting (@-@, spaced punctuation) doesn't exist in the raw pool, so no pool doc scores highest on wiki. Two fixes: (1) normalize the target text (strip WikiText formatting) so wiki matches on content vocabulary, and (2) fill each register's quota from its own ranked list (with global dedup) so every quota fills. Let me check time and patch.
toolBashCheck remaining time
args
{
"command": "cat /workspace/.timer/remaining_secs",
"description": "Check remaining time"
}Bash result
3467
[stdout]
3467
[metadata] {"noOutputExpected": false}#43
narration
~58 min left. Let me patch: normalize the wiki target (strip WikiText formatting) and fill each quota from its own ranked list with global dedup.
toolEdit
args
{
"replace_all": false,
"file_path": "/workspace/submission/curate.py",
"old_string": "def register_of(t):",
"new_string": "def normalize_target(t):\n \"\"\"Strip WikiText-103 surface formatting so the wiki target matches on CONTENT\n vocabulary, not on markup absent from the raw pool (which would otherwise make\n every pool doc look un-wiki-like and starve the encyclopedic quota).\"\"\"\n t = t.replace(\"@-@\", \"-\").replace(\"@,@\", \",\").replace(\"@.@\", \".\")\n t = re.sub(r\"\\s+([,.;:!?)])\", r\"\\1\", t) # drop space before punctuation\n t = re.sub(r\"([(\\[])\\s+\", r\"\\1\", t) # drop space after open bracket\n t = re.sub(r\"\\s+\", \" \", t) # collapse whitespace\n return t\n\ndef register_of(t):"
}Edit result
The file /workspace/submission/curate.py has been updated successfully. (file state is current in your context — no need to Read it back)
[metadata] {"filePath": "/workspace/submission/curate.py", "oldString": "def register_of(t):", "newString": "def normalize_target(t):\n \"\"\"Strip WikiText-103 surface formatting so the wiki target matches on CONTENT\n vocabulary, not on markup absent from the raw pool (which would otherwise make\n every pool doc look un-wiki-like and starve the encyclopedic quota).\"\"\"\n t = t.replace(\"@-@\", \"-\").replace(\"@,@\", \",\").replace(\"@.@\", \".\")\n t = re.sub(r\"\\s+([,.;:!?)])\", r\"\\1\", t) # drop space before punctuation\n t = re.sub(r\"([(\\[])\\s+\", r\"\\1\", t) # drop space after open bracket\n t = re.sub(r\"\\s+\", \" \", t) # collapse whitespace\n return t\n\ndef register_of(t):", "originalFile": "\"\"\"Curate a pretraining subset from a raw web pool by register-balanced DSIR.\n\nThe disclosed evaluation target is a BROAD, multi-domain HQ English mix in equal\nparts: encyclopedic (Wikipedia), general web prose, news, and technical Q&A.\nHeld-out perplexity averages over all four registers, so a training set that\nover-weights one register (e.g. news, which is easiest to match) leaves the\nothers under-trained and drags the average up. We therefore select in two steps:\n\n (1) DSIR importance scoring, PER REGISTER. We split the disclosed target\n (data/multi_dev.npy, GPT-2 tokens) into four register buckets by surface\n signatures, and for each register r estimate a bag-of-n-grams importance\n weight\n w_r(gram) = log p_target_r(gram) - log p_pool(gram)\n (unigrams + hashed bigrams; p_pool from the whole pool as background).\n Each pool document is scored against every register using a PRESENCE-based\n mean over its unique n-gram types (repetition-robust), behind a length\n floor and a letter-fraction gate that reject fragments and non-prose\n (directory listings, tables). A document is routed to its best-fit\n register (argmax score).\n\n (2) Balanced quota fill. Within each register we rank routed documents by\n score, then interleave the four ranked lists by EQUAL token quota, so the\n first 12M tokens the trainer consumes are ~25% from each register --\n matching the equal-parts evaluation mixture.\n\nThis is Data Selection via Importance Resampling (Xie et al. 2023) with a\nper-domain target, plus classic quality gating. No labels, no internet: the only\nsupervision is the disclosed target token stream.\n\"\"\"\nimport json, time, re, numpy as np\nfrom transformers import AutoTokenizer\n\nPOOL = \"/workspace/data/pool.jsonl\"\nDEV = \"/workspace/data/multi_dev.npy\"\nOUT = \"/workspace/submission/selection.json\"\n\nEOS = 50256\nV = 50257 # GPT-2 vocab\nDB = 1 << 20 # hashed bigram buckets\nALPHA = 0.1 # additive smoothing\nMIN_TOK = 100 # length floor: drop fragments\nMIN_ALPHA = 0.55 # letter-fraction floor: reject non-prose\nLAMBDA = 0.5 # weight of bigram term relative to unigram term\nCOVER_TOKENS = 30_000_000 # emit ids covering ~2.5x the 12M budget\nREGISTERS = [\"wiki\", \"web\", \"news\", \"code_qa\"]\nQUOTA = {r: 0.25 for r in REGISTERS} # equal-parts target mixture\n\n_NONALPHA = re.compile(r\"[^A-Za-z]\")\n_NEWS = re.compile(r\"\\b(reuters|said|told|according to|\"\n r\"on (monday|tuesday|wednesday|thursday|friday|saturday|sunday))\\b\")\n\ndef register_of(t):\n \"\"\"Surface-signature register label for a piece of text.\"\"\"\n if \"@-@\" in t or \"@,@\" in t or \"@.@\" in t:\n return \"wiki\"\n if (\"<p>\" in t or \"<code>\" in t or \"</\" in t or \"def \" in t or \"import \" in t\n or (\"{\" in t and \"}\" in t and \";\" in t)):\n return \"code_qa\"\n if _NEWS.search(t.lower()):\n return \"news\"\n return \"web\"\n\ndef bigrams(a):\n return ((a[:-1] * np.int64(1000003) + a[1:]) & np.int64(DB - 1))\n\nt0 = time.time()\ntok = AutoTokenizer.from_pretrained(\"gpt2\")\n\n# ------------------------------------------------------------------ load pool\nids, texts = [], []\nwith open(POOL) as f:\n for line in f:\n r = json.loads(line)\n ids.append(r[\"id\"]); texts.append(r[\"text\"])\nids = np.asarray(ids, dtype=np.int64)\nndoc = len(ids)\nprint(f\"[{time.time()-t0:.0f}s] loaded pool: {ndoc} docs\", flush=True)\n\n# ------------------------------------------------------ tokenize pool (batched)\ndocs = []\nfor i in range(0, ndoc, 4000):\n enc = tok(texts[i:i+4000], add_special_tokens=False).input_ids\n docs.extend(np.asarray(x, dtype=np.int32) for x in enc)\nlens = np.fromiter((len(d) for d in docs), dtype=np.int64, count=ndoc)\nalltok = np.concatenate([d for d in docs]).astype(np.int64)\nprint(f\"[{time.time()-t0:.0f}s] tokenized pool: {lens.sum()/1e6:.1f}M tokens\", flush=True)\n\n# --------------------------------- pool background counts (shared denominator)\nneg_uni = np.bincount(alltok, minlength=V).astype(np.float64)\nneg_big = np.bincount(bigrams(alltok), minlength=DB).astype(np.float64)\nneg_uni_ln = np.log(neg_uni + ALPHA) - np.log(neg_uni.sum() + ALPHA * V)\nneg_big_ln = np.log(neg_big + ALPHA) - np.log(neg_big.sum() + ALPHA * DB)\ndel alltok\nprint(f\"[{time.time()-t0:.0f}s] built pool background\", flush=True)\n\n# ---------------------------- per-register target counts from disclosed target\ndev = np.load(DEV).astype(np.int64)\nsplits = np.where(dev == EOS)[0]\nb = [0] + (splits + 1).tolist() + [len(dev)]\nuni_w = {r: None for r in REGISTERS}\nbig_w = {r: None for r in REGISTERS}\npos_uni = {r: np.zeros(V, np.float64) for r in REGISTERS}\npos_big = {r: np.zeros(DB, np.float64) for r in REGISTERS}\nfor k in range(len(b) - 1):\n d = dev[b[k]:b[k+1]]\n d = d[d != EOS]\n if len(d) < 20:\n continue\n r = register_of(tok.decode(d.tolist()))\n pos_uni[r] += np.bincount(d, minlength=V)\n if len(d) > 1:\n pos_big[r] += np.bincount(bigrams(d), minlength=DB)\nfor r in REGISTERS:\n uni_w[r] = (np.log(pos_uni[r] + ALPHA) - np.log(pos_uni[r].sum() + ALPHA * V)) - neg_uni_ln\n big_w[r] = (np.log(pos_big[r] + ALPHA) - np.log(pos_big[r].sum() + ALPHA * DB)) - neg_big_ln\nprint(f\"[{time.time()-t0:.0f}s] built {len(REGISTERS)} register targets \"\n f\"(tok: {{{', '.join(f'{r}:{int(pos_uni[r].sum())}' for r in REGISTERS)}}})\", flush=True)\n\n# ------------------------ presence-based scoring of every doc vs every register\nbest_reg = np.full(ndoc, -1, dtype=np.int64)\nbest_score = np.full(ndoc, -1e30, dtype=np.float64)\nfor i in range(ndoc):\n d = docs[i]; n = d.shape[0]\n if n < MIN_TOK:\n continue\n t = texts[i]\n if len(t) == 0 or len(_NONALPHA.sub(\"\", t)) / len(t) < MIN_ALPHA:\n continue\n di = d.astype(np.int64)\n u = np.unique(di)\n ub = np.unique(bigrams(di)) if n > 1 else np.array([0], dtype=np.int64)\n bi, bs = -1, -1e30\n for ri, r in enumerate(REGISTERS):\n s = uni_w[r][u].mean() + LAMBDA * big_w[r][ub].mean()\n if s > bs:\n bs, bi = s, ri\n best_reg[i] = bi; best_score[i] = bs\n if (i & 0x3FFF) == 0:\n print(f\"[{time.time()-t0:.0f}s] scoring {i}/{ndoc}\", flush=True)\nprint(f\"[{time.time()-t0:.0f}s] scored docs\", flush=True)\n\n# ------------------------------------- rank within register, interleave by quota\nranked = {}\nfor ri, r in enumerate(REGISTERS):\n idx = np.where(best_reg == ri)[0]\n idx = idx[np.argsort(-best_score[idx])]\n ranked[r] = idx.tolist()\n print(f\" {r}: {len(idx)} docs routed\", flush=True)\n\nptr = {r: 0 for r in REGISTERS}\ncum = {r: 0 for r in REGISTERS}\nselection, seen, tot = [], set(), 0\nwhile tot < COVER_TOKENS:\n # pick the register that is furthest below its token quota and still has docs\n cand = [r for r in REGISTERS if ptr[r] < len(ranked[r])]\n if not cand:\n break\n r = min(cand, key=lambda r: cum[r] / QUOTA[r])\n j = ranked[r][ptr[r]]; ptr[r] += 1\n h = hash(texts[j])\n if h in seen:\n continue\n seen.add(h)\n selection.append(int(ids[j]))\n cum[r] += int(lens[j]); tot += int(lens[j])\n\njson.dump(selection, open(OUT, \"w\"))\nmix = {r: round(cum[r] / max(1, tot), 3) for r in REGISTERS}\nprint(f\"[{time.time()-t0:.0f}s] wrote {len(selection)} ids covering ~{tot/1e6:.1f}M tokens\", flush=True)\nprint(f\"register token mix of emitted list: {mix}\", flush=True)\n", "structuredPatch": [{"oldStart": 49, "oldLines": 6, "newStart": 49, "newLines": 16, "lines": [" _NEWS = re.compile(r\"\\b(reuters|said|told|according to|\"", " r\"on (monday|tuesday|wednesday|thursday|friday|saturday|sunday))\\b\")", " ", "+def normalize_target(t):", "+ \"\"\"Strip WikiText-103 surface formatting so the wiki target matches on CONTENT", "+ vocabulary, not on markup absent from the raw pool (which would otherwise make", "+ every pool doc look un-wiki-like and starve the encyclopedic quota).\"\"\"", "+ t = t.replace(\"@-@\", \"-\").replace(\"@,@\", \",\").replace(\"@.@\", \".\")", "+ t = re.sub(r\"\\s+([,.;:!?)])\", r\"\\1\", t) # drop space before punctuation", "+ t = re.sub(r\"([(\\[])\\s+\", r\"\\1\", t) # drop space after open bracket", "+ t = re.sub(r\"\\s+\", \" \", t) # collapse whitespace", "+ return t", "+", " def register_of(t):", " \"\"\"Surface-signature register label for a piece of text.\"\"\"", " if \"@-@\" in t or \"@,@\" in t or \"@.@\" in t:"]}], "userModified": false, "replaceAll": false}#44
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolEdit
args
{
"replace_all": false,
"file_path": "/workspace/submission/curate.py",
"old_string": " r = register_of(tok.decode(d.tolist()))\n pos_uni[r] += np.bincount(d, minlength=V)\n if len(d) > 1:\n pos_big[r] += np.bincount(bigrams(d), minlength=DB)",
"new_string": " raw = tok.decode(d.tolist())\n r = register_of(raw)\n if r == \"wiki\": # de-format encyclopedic text\n d = np.asarray(tok(normalize_target(raw), add_special_tokens=False).input_ids,\n dtype=np.int64)\n pos_uni[r] += np.bincount(d, minlength=V)\n if len(d) > 1:\n pos_big[r] += np.bincount(bigrams(d), minlength=DB)"
}Edit result
The file /workspace/submission/curate.py has been updated successfully. (file state is current in your context — no need to Read it back)
[metadata] {"filePath": "/workspace/submission/curate.py", "oldString": " r = register_of(tok.decode(d.tolist()))\n pos_uni[r] += np.bincount(d, minlength=V)\n if len(d) > 1:\n pos_big[r] += np.bincount(bigrams(d), minlength=DB)", "newString": " raw = tok.decode(d.tolist())\n r = register_of(raw)\n if r == \"wiki\": # de-format encyclopedic text\n d = np.asarray(tok(normalize_target(raw), add_special_tokens=False).input_ids,\n dtype=np.int64)\n pos_uni[r] += np.bincount(d, minlength=V)\n if len(d) > 1:\n pos_big[r] += np.bincount(bigrams(d), minlength=DB)", "originalFile": "\"\"\"Curate a pretraining subset from a raw web pool by register-balanced DSIR.\n\nThe disclosed evaluation target is a BROAD, multi-domain HQ English mix in equal\nparts: encyclopedic (Wikipedia), general web prose, news, and technical Q&A.\nHeld-out perplexity averages over all four registers, so a training set that\nover-weights one register (e.g. news, which is easiest to match) leaves the\nothers under-trained and drags the average up. We therefore select in two steps:\n\n (1) DSIR importance scoring, PER REGISTER. We split the disclosed target\n (data/multi_dev.npy, GPT-2 tokens) into four register buckets by surface\n signatures, and for each register r estimate a bag-of-n-grams importance\n weight\n w_r(gram) = log p_target_r(gram) - log p_pool(gram)\n (unigrams + hashed bigrams; p_pool from the whole pool as background).\n Each pool document is scored against every register using a PRESENCE-based\n mean over its unique n-gram types (repetition-robust), behind a length\n floor and a letter-fraction gate that reject fragments and non-prose\n (directory listings, tables). A document is routed to its best-fit\n register (argmax score).\n\n (2) Balanced quota fill. Within each register we rank routed documents by\n score, then interleave the four ranked lists by EQUAL token quota, so the\n first 12M tokens the trainer consumes are ~25% from each register --\n matching the equal-parts evaluation mixture.\n\nThis is Data Selection via Importance Resampling (Xie et al. 2023) with a\nper-domain target, plus classic quality gating. No labels, no internet: the only\nsupervision is the disclosed target token stream.\n\"\"\"\nimport json, time, re, numpy as np\nfrom transformers import AutoTokenizer\n\nPOOL = \"/workspace/data/pool.jsonl\"\nDEV = \"/workspace/data/multi_dev.npy\"\nOUT = \"/workspace/submission/selection.json\"\n\nEOS = 50256\nV = 50257 # GPT-2 vocab\nDB = 1 << 20 # hashed bigram buckets\nALPHA = 0.1 # additive smoothing\nMIN_TOK = 100 # length floor: drop fragments\nMIN_ALPHA = 0.55 # letter-fraction floor: reject non-prose\nLAMBDA = 0.5 # weight of bigram term relative to unigram term\nCOVER_TOKENS = 30_000_000 # emit ids covering ~2.5x the 12M budget\nREGISTERS = [\"wiki\", \"web\", \"news\", \"code_qa\"]\nQUOTA = {r: 0.25 for r in REGISTERS} # equal-parts target mixture\n\n_NONALPHA = re.compile(r\"[^A-Za-z]\")\n_NEWS = re.compile(r\"\\b(reuters|said|told|according to|\"\n r\"on (monday|tuesday|wednesday|thursday|friday|saturday|sunday))\\b\")\n\ndef normalize_target(t):\n \"\"\"Strip WikiText-103 surface formatting so the wiki target matches on CONTENT\n vocabulary, not on markup absent from the raw pool (which would otherwise make\n every pool doc look un-wiki-like and starve the encyclopedic quota).\"\"\"\n t = t.replace(\"@-@\", \"-\").replace(\"@,@\", \",\").replace(\"@.@\", \".\")\n t = re.sub(r\"\\s+([,.;:!?)])\", r\"\\1\", t) # drop space before punctuation\n t = re.sub(r\"([(\\[])\\s+\", r\"\\1\", t) # drop space after open bracket\n t = re.sub(r\"\\s+\", \" \", t) # collapse whitespace\n return t\n\ndef register_of(t):\n \"\"\"Surface-signature register label for a piece of text.\"\"\"\n if \"@-@\" in t or \"@,@\" in t or \"@.@\" in t:\n return \"wiki\"\n if (\"<p>\" in t or \"<code>\" in t or \"</\" in t or \"def \" in t or \"import \" in t\n or (\"{\" in t and \"}\" in t and \";\" in t)):\n return \"code_qa\"\n if _NEWS.search(t.lower()):\n return \"news\"\n return \"web\"\n\ndef bigrams(a):\n return ((a[:-1] * np.int64(1000003) + a[1:]) & np.int64(DB - 1))\n\nt0 = time.time()\ntok = AutoTokenizer.from_pretrained(\"gpt2\")\n\n# ------------------------------------------------------------------ load pool\nids, texts = [], []\nwith open(POOL) as f:\n for line in f:\n r = json.loads(line)\n ids.append(r[\"id\"]); texts.append(r[\"text\"])\nids = np.asarray(ids, dtype=np.int64)\nndoc = len(ids)\nprint(f\"[{time.time()-t0:.0f}s] loaded pool: {ndoc} docs\", flush=True)\n\n# ------------------------------------------------------ tokenize pool (batched)\ndocs = []\nfor i in range(0, ndoc, 4000):\n enc = tok(texts[i:i+4000], add_special_tokens=False).input_ids\n docs.extend(np.asarray(x, dtype=np.int32) for x in enc)\nlens = np.fromiter((len(d) for d in docs), dtype=np.int64, count=ndoc)\nalltok = np.concatenate([d for d in docs]).astype(np.int64)\nprint(f\"[{time.time()-t0:.0f}s] tokenized pool: {lens.sum()/1e6:.1f}M tokens\", flush=True)\n\n# --------------------------------- pool background counts (shared denominator)\nneg_uni = np.bincount(alltok, minlength=V).astype(np.float64)\nneg_big = np.bincount(bigrams(alltok), minlength=DB).astype(np.float64)\nneg_uni_ln = np.log(neg_uni + ALPHA) - np.log(neg_uni.sum() + ALPHA * V)\nneg_big_ln = np.log(neg_big + ALPHA) - np.log(neg_big.sum() + ALPHA * DB)\ndel alltok\nprint(f\"[{time.time()-t0:.0f}s] built pool background\", flush=True)\n\n# ---------------------------- per-register target counts from disclosed target\ndev = np.load(DEV).astype(np.int64)\nsplits = np.where(dev == EOS)[0]\nb = [0] + (splits + 1).tolist() + [len(dev)]\nuni_w = {r: None for r in REGISTERS}\nbig_w = {r: None for r in REGISTERS}\npos_uni = {r: np.zeros(V, np.float64) for r in REGISTERS}\npos_big = {r: np.zeros(DB, np.float64) for r in REGISTERS}\nfor k in range(len(b) - 1):\n d = dev[b[k]:b[k+1]]\n d = d[d != EOS]\n if len(d) < 20:\n continue\n r = register_of(tok.decode(d.tolist()))\n pos_uni[r] += np.bincount(d, minlength=V)\n if len(d) > 1:\n pos_big[r] += np.bincount(bigrams(d), minlength=DB)\nfor r in REGISTERS:\n uni_w[r] = (np.log(pos_uni[r] + ALPHA) - np.log(pos_uni[r].sum() + ALPHA * V)) - neg_uni_ln\n big_w[r] = (np.log(pos_big[r] + ALPHA) - np.log(pos_big[r].sum() + ALPHA * DB)) - neg_big_ln\nprint(f\"[{time.time()-t0:.0f}s] built {len(REGISTERS)} register targets \"\n f\"(tok: {{{', '.join(f'{r}:{int(pos_uni[r].sum())}' for r in REGISTERS)}}})\", flush=True)\n\n# ------------------------ presence-based scoring of every doc vs every register\nbest_reg = np.full(ndoc, -1, dtype=np.int64)\nbest_score = np.full(ndoc, -1e30, dtype=np.float64)\nfor i in range(ndoc):\n d = docs[i]; n = d.shape[0]\n if n < MIN_TOK:\n continue\n t = texts[i]\n if len(t) == 0 or len(_NONALPHA.sub(\"\", t)) / len(t) < MIN_ALPHA:\n continue\n di = d.astype(np.int64)\n u = np.unique(di)\n ub = np.unique(bigrams(di)) if n > 1 else np.array([0], dtype=np.int64)\n bi, bs = -1, -1e30\n for ri, r in enumerate(REGISTERS):\n s = uni_w[r][u].mean() + LAMBDA * big_w[r][ub].mean()\n if s > bs:\n bs, bi = s, ri\n best_reg[i] = bi; best_score[i] = bs\n if (i & 0x3FFF) == 0:\n print(f\"[{time.time()-t0:.0f}s] scoring {i}/{ndoc}\", flush=True)\nprint(f\"[{time.time()-t0:.0f}s] scored docs\", flush=True)\n\n# ------------------------------------- rank within register, interleave by quota\nranked = {}\nfor ri, r in enumerate(REGISTERS):\n idx = np.where(best_reg == ri)[0]\n idx = idx[np.argsort(-best_score[idx])]\n ranked[r] = idx.tolist()\n print(f\" {r}: {len(idx)} docs routed\", flush=True)\n\nptr = {r: 0 for r in REGISTERS}\ncum = {r: 0 for r in REGISTERS}\nselection, seen, tot = [], set(), 0\nwhile tot < COVER_TOKENS:\n # pick the register that is furthest below its token quota and still has docs\n cand = [r for r in REGISTERS if ptr[r] < len(ranked[r])]\n if not cand:\n break\n r = min(cand, key=lambda r: cum[r] / QUOTA[r])\n j = ranked[r][ptr[r]]; ptr[r] += 1\n h = hash(texts[j])\n if h in seen:\n continue\n seen.add(h)\n selection.append(int(ids[j]))\n cum[r] += int(lens[j]); tot += int(lens[j])\n\njson.dump(selection, open(OUT, \"w\"))\nmix = {r: round(cum[r] / max(1, tot), 3) for r in REGISTERS}\nprint(f\"[{time.time()-t0:.0f}s] wrote {len(selection)} ids covering ~{tot/1e6:.1f}M tokens\", flush=True)\nprint(f\"register token mix of emitted list: {mix}\", flush=True)\n", "structuredPatch": [{"oldStart": 116, "oldLines": 7, "newStart": 116, "newLines": 11, "lines": [" d = d[d != EOS]", " if len(d) < 20:", " continue", "- r = register_of(tok.decode(d.tolist()))", "+ raw = tok.decode(d.tolist())", "+ r = register_of(raw)", "+ if r == \"wiki\": # de-format encyclopedic text", "+ d = np.asarray(tok(normalize_target(raw), add_special_tokens=False).input_ids,", "+ dtype=np.int64)", " pos_uni[r] += np.bincount(d, minlength=V)", " if len(d) > 1:", " pos_big[r] += np.bincount(bigrams(d), minlength=DB)"]}], "userModified": false, "replaceAll": false}#45
narration
Now replace the scoring loop and fill algorithm to use a full per-register score matrix with independent ranked lists:
toolEdit
args
{
"replace_all": false,
"file_path": "/workspace/submission/curate.py",
"old_string": "# ------------------------ presence-based scoring of every doc vs every register\nbest_reg = np.full(ndoc, -1, dtype=np.int64)\nbest_score = np.full(ndoc, -1e30, dtype=np.float64)\nfor i in range(ndoc):\n d = docs[i]; n = d.shape[0]\n if n < MIN_TOK:\n continue\n t = texts[i]\n if len(t) == 0 or len(_NONALPHA.sub(\"\", t)) / len(t) < MIN_ALPHA:\n continue\n di = d.astype(np.int64)\n u = np.unique(di)\n ub = np.unique(bigrams(di)) if n > 1 else np.array([0], dtype=np.int64)\n bi, bs = -1, -1e30\n for ri, r in enumerate(REGISTERS):\n s = uni_w[r][u].mean() + LAMBDA * big_w[r][ub].mean()\n if s > bs:\n bs, bi = s, ri\n best_reg[i] = bi; best_score[i] = bs\n if (i & 0x3FFF) == 0:\n print(f\"[{time.time()-t0:.0f}s] scoring {i}/{ndoc}\", flush=True)\nprint(f\"[{time.time()-t0:.0f}s] scored docs\", flush=True)\n\n# ------------------------------------- rank within register, interleave by quota\nranked = {}\nfor ri, r in enumerate(REGISTERS):\n idx = np.where(best_reg == ri)[0]\n idx = idx[np.argsort(-best_score[idx])]\n ranked[r] = idx.tolist()\n print(f\" {r}: {len(idx)} docs routed\", flush=True)\n\nptr = {r: 0 for r in REGISTERS}\ncum = {r: 0 for r in REGISTERS}\nselection, seen, tot = [], set(), 0\nwhile tot < COVER_TOKENS:\n # pick the register that is furthest below its token quota and still has docs\n cand = [r for r in REGISTERS if ptr[r] < len(ranked[r])]\n if not cand:\n break\n r = min(cand, key=lambda r: cum[r] / QUOTA[r])\n j = ranked[r][ptr[r]]; ptr[r] += 1\n h = hash(texts[j])\n if h in seen:\n continue\n seen.add(h)\n selection.append(int(ids[j]))\n cum[r] += int(lens[j]); tot += int(lens[j])\n\njson.dump(selection, open(OUT, \"w\"))\nmix = {r: round(cum[r] / max(1, tot), 3) for r in REGISTERS}\nprint(f\"[{time.time()-t0:.0f}s] wrote {len(selection)} ids covering ~{tot/1e6:.1f}M tokens\", flush=True)\nprint(f\"register token mix of emitted list: {mix}\", flush=True)",
"new_string": "# ------------------------ presence-based scoring of every doc vs every register\nNR = len(REGISTERS)\nscoremat = np.full((ndoc, NR), -1e30, dtype=np.float64)\ngate_ok = np.zeros(ndoc, dtype=bool)\nfor i in range(ndoc):\n d = docs[i]; n = d.shape[0]\n if n < MIN_TOK:\n continue\n t = texts[i]\n if len(t) == 0 or len(_NONALPHA.sub(\"\", t)) / len(t) < MIN_ALPHA:\n continue\n di = d.astype(np.int64)\n u = np.unique(di)\n ub = np.unique(bigrams(di)) if n > 1 else np.array([0], dtype=np.int64)\n for ri, r in enumerate(REGISTERS):\n scoremat[i, ri] = uni_w[r][u].mean() + LAMBDA * big_w[r][ub].mean()\n gate_ok[i] = True\n if (i & 0x3FFF) == 0:\n print(f\"[{time.time()-t0:.0f}s] scoring {i}/{ndoc}\", flush=True)\nprint(f\"[{time.time()-t0:.0f}s] scored docs ({gate_ok.sum()} passed gates)\", flush=True)\n\n# ------------- independent per-register ranked lists, interleaved by equal quota\nok = np.where(gate_ok)[0]\nranked = {r: ok[np.argsort(-scoremat[ok, ri])] for ri, r in enumerate(REGISTERS)}\n\nptr = {r: 0 for r in REGISTERS}\ncum = {r: 0 for r in REGISTERS}\nselection, seen, chosen, tot = [], set(), set(), 0\nwhile tot < COVER_TOKENS:\n cand = [r for r in REGISTERS if ptr[r] < len(ranked[r])]\n if not cand:\n break\n r = min(cand, key=lambda r: cum[r] / QUOTA[r])\n # advance past docs already claimed by another register\n while ptr[r] < len(ranked[r]) and ranked[r][ptr[r]] in chosen:\n ptr[r] += 1\n if ptr[r] >= len(ranked[r]):\n continue\n j = int(ranked[r][ptr[r]]); ptr[r] += 1\n chosen.add(j)\n h = hash(texts[j])\n if h in seen: # exact-duplicate guard\n continue\n seen.add(h)\n selection.append(int(ids[j]))\n cum[r] += int(lens[j]); tot += int(lens[j])\n\njson.dump(selection, open(OUT, \"w\"))\nmix = {r: round(cum[r] / max(1, tot), 3) for r in REGISTERS}\nprint(f\"[{time.time()-t0:.0f}s] wrote {len(selection)} ids covering ~{tot/1e6:.1f}M tokens\", flush=True)\nprint(f\"register token mix of emitted list: {mix}\", flush=True)"
}Edit result
The file /workspace/submission/curate.py has been updated successfully. (file state is current in your context — no need to Read it back)
[metadata] {"filePath": "/workspace/submission/curate.py", "oldString": "# ------------------------ presence-based scoring of every doc vs every register\nbest_reg = np.full(ndoc, -1, dtype=np.int64)\nbest_score = np.full(ndoc, -1e30, dtype=np.float64)\nfor i in range(ndoc):\n d = docs[i]; n = d.shape[0]\n if n < MIN_TOK:\n continue\n t = texts[i]\n if len(t) == 0 or len(_NONALPHA.sub(\"\", t)) / len(t) < MIN_ALPHA:\n continue\n di = d.astype(np.int64)\n u = np.unique(di)\n ub = np.unique(bigrams(di)) if n > 1 else np.array([0], dtype=np.int64)\n bi, bs = -1, -1e30\n for ri, r in enumerate(REGISTERS):\n s = uni_w[r][u].mean() + LAMBDA * big_w[r][ub].mean()\n if s > bs:\n bs, bi = s, ri\n best_reg[i] = bi; best_score[i] = bs\n if (i & 0x3FFF) == 0:\n print(f\"[{time.time()-t0:.0f}s] scoring {i}/{ndoc}\", flush=True)\nprint(f\"[{time.time()-t0:.0f}s] scored docs\", flush=True)\n\n# ------------------------------------- rank within register, interleave by quota\nranked = {}\nfor ri, r in enumerate(REGISTERS):\n idx = np.where(best_reg == ri)[0]\n idx = idx[np.argsort(-best_score[idx])]\n ranked[r] = idx.tolist()\n print(f\" {r}: {len(idx)} docs routed\", flush=True)\n\nptr = {r: 0 for r in REGISTERS}\ncum = {r: 0 for r in REGISTERS}\nselection, seen, tot = [], set(), 0\nwhile tot < COVER_TOKENS:\n # pick the register that is furthest below its token quota and still has docs\n cand = [r for r in REGISTERS if ptr[r] < len(ranked[r])]\n if not cand:\n break\n r = min(cand, key=lambda r: cum[r] / QUOTA[r])\n j = ranked[r][ptr[r]]; ptr[r] += 1\n h = hash(texts[j])\n if h in seen:\n continue\n seen.add(h)\n selection.append(int(ids[j]))\n cum[r] += int(lens[j]); tot += int(lens[j])\n\njson.dump(selection, open(OUT, \"w\"))\nmix = {r: round(cum[r] / max(1, tot), 3) for r in REGISTERS}\nprint(f\"[{time.time()-t0:.0f}s] wrote {len(selection)} ids covering ~{tot/1e6:.1f}M tokens\", flush=True)\nprint(f\"register token mix of emitted list: {mix}\", flush=True)", "newString": "# ------------------------ presence-based scoring of every doc vs every register\nNR = len(REGISTERS)\nscoremat = np.full((ndoc, NR), -1e30, dtype=np.float64)\ngate_ok = np.zeros(ndoc, dtype=bool)\nfor i in range(ndoc):\n d = docs[i]; n = d.shape[0]\n if n < MIN_TOK:\n continue\n t = texts[i]\n if len(t) == 0 or len(_NONALPHA.sub(\"\", t)) / len(t) < MIN_ALPHA:\n continue\n di = d.astype(np.int64)\n u = np.unique(di)\n ub = np.unique(bigrams(di)) if n > 1 else np.array([0], dtype=np.int64)\n for ri, r in enumerate(REGISTERS):\n scoremat[i, ri] = uni_w[r][u].mean() + LAMBDA * big_w[r][ub].mean()\n gate_ok[i] = True\n if (i & 0x3FFF) == 0:\n print(f\"[{time.time()-t0:.0f}s] scoring {i}/{ndoc}\", flush=True)\nprint(f\"[{time.time()-t0:.0f}s] scored docs ({gate_ok.sum()} passed gates)\", flush=True)\n\n# ------------- independent per-register ranked lists, interleaved by equal quota\nok = np.where(gate_ok)[0]\nranked = {r: ok[np.argsort(-scoremat[ok, ri])] for ri, r in enumerate(REGISTERS)}\n\nptr = {r: 0 for r in REGISTERS}\ncum = {r: 0 for r in REGISTERS}\nselection, seen, chosen, tot = [], set(), set(), 0\nwhile tot < COVER_TOKENS:\n cand = [r for r in REGISTERS if ptr[r] < len(ranked[r])]\n if not cand:\n break\n r = min(cand, key=lambda r: cum[r] / QUOTA[r])\n # advance past docs already claimed by another register\n while ptr[r] < len(ranked[r]) and ranked[r][ptr[r]] in chosen:\n ptr[r] += 1\n if ptr[r] >= len(ranked[r]):\n continue\n j = int(ranked[r][ptr[r]]); ptr[r] += 1\n chosen.add(j)\n h = hash(texts[j])\n if h in seen: # exact-duplicate guard\n continue\n seen.add(h)\n selection.append(int(ids[j]))\n cum[r] += int(lens[j]); tot += int(lens[j])\n\njson.dump(selection, open(OUT, \"w\"))\nmix = {r: round(cum[r] / max(1, tot), 3) for r in REGISTERS}\nprint(f\"[{time.time()-t0:.0f}s] wrote {len(selection)} ids covering ~{tot/1e6:.1f}M tokens\", flush=True)\nprint(f\"register token mix of emitted list: {mix}\", flush=True)", "originalFile": "\"\"\"Curate a pretraining subset from a raw web pool by register-balanced DSIR.\n\nThe disclosed evaluation target is a BROAD, multi-domain HQ English mix in equal\nparts: encyclopedic (Wikipedia), general web prose, news, and technical Q&A.\nHeld-out perplexity averages over all four registers, so a training set that\nover-weights one register (e.g. news, which is easiest to match) leaves the\nothers under-trained and drags the average up. We therefore select in two steps:\n\n (1) DSIR importance scoring, PER REGISTER. We split the disclosed target\n (data/multi_dev.npy, GPT-2 tokens) into four register buckets by surface\n signatures, and for each register r estimate a bag-of-n-grams importance\n weight\n w_r(gram) = log p_target_r(gram) - log p_pool(gram)\n (unigrams + hashed bigrams; p_pool from the whole pool as background).\n Each pool document is scored against every register using a PRESENCE-based\n mean over its unique n-gram types (repetition-robust), behind a length\n floor and a letter-fraction gate that reject fragments and non-prose\n (directory listings, tables). A document is routed to its best-fit\n register (argmax score).\n\n (2) Balanced quota fill. Within each register we rank routed documents by\n score, then interleave the four ranked lists by EQUAL token quota, so the\n first 12M tokens the trainer consumes are ~25% from each register --\n matching the equal-parts evaluation mixture.\n\nThis is Data Selection via Importance Resampling (Xie et al. 2023) with a\nper-domain target, plus classic quality gating. No labels, no internet: the only\nsupervision is the disclosed target token stream.\n\"\"\"\nimport json, time, re, numpy as np\nfrom transformers import AutoTokenizer\n\nPOOL = \"/workspace/data/pool.jsonl\"\nDEV = \"/workspace/data/multi_dev.npy\"\nOUT = \"/workspace/submission/selection.json\"\n\nEOS = 50256\nV = 50257 # GPT-2 vocab\nDB = 1 << 20 # hashed bigram buckets\nALPHA = 0.1 # additive smoothing\nMIN_TOK = 100 # length floor: drop fragments\nMIN_ALPHA = 0.55 # letter-fraction floor: reject non-prose\nLAMBDA = 0.5 # weight of bigram term relative to unigram term\nCOVER_TOKENS = 30_000_000 # emit ids covering ~2.5x the 12M budget\nREGISTERS = [\"wiki\", \"web\", \"news\", \"code_qa\"]\nQUOTA = {r: 0.25 for r in REGISTERS} # equal-parts target mixture\n\n_NONALPHA = re.compile(r\"[^A-Za-z]\")\n_NEWS = re.compile(r\"\\b(reuters|said|told|according to|\"\n r\"on (monday|tuesday|wednesday|thursday|friday|saturday|sunday))\\b\")\n\ndef normalize_target(t):\n \"\"\"Strip WikiText-103 surface formatting so the wiki target matches on CONTENT\n vocabulary, not on markup absent from the raw pool (which would otherwise make\n every pool doc look un-wiki-like and starve the encyclopedic quota).\"\"\"\n t = t.replace(\"@-@\", \"-\").replace(\"@,@\", \",\").replace(\"@.@\", \".\")\n t = re.sub(r\"\\s+([,.;:!?)])\", r\"\\1\", t) # drop space before punctuation\n t = re.sub(r\"([(\\[])\\s+\", r\"\\1\", t) # drop space after open bracket\n t = re.sub(r\"\\s+\", \" \", t) # collapse whitespace\n return t\n\ndef register_of(t):\n \"\"\"Surface-signature register label for a piece of text.\"\"\"\n if \"@-@\" in t or \"@,@\" in t or \"@.@\" in t:\n return \"wiki\"\n if (\"<p>\" in t or \"<code>\" in t or \"</\" in t or \"def \" in t or \"import \" in t\n or (\"{\" in t and \"}\" in t and \";\" in t)):\n return \"code_qa\"\n if _NEWS.search(t.lower()):\n return \"news\"\n return \"web\"\n\ndef bigrams(a):\n return ((a[:-1] * np.int64(1000003) + a[1:]) & np.int64(DB - 1))\n\nt0 = time.time()\ntok = AutoTokenizer.from_pretrained(\"gpt2\")\n\n# ------------------------------------------------------------------ load pool\nids, texts = [], []\nwith open(POOL) as f:\n for line in f:\n r = json.loads(line)\n ids.append(r[\"id\"]); texts.append(r[\"text\"])\nids = np.asarray(ids, dtype=np.int64)\nndoc = len(ids)\nprint(f\"[{time.time()-t0:.0f}s] loaded pool: {ndoc} docs\", flush=True)\n\n# ------------------------------------------------------ tokenize pool (batched)\ndocs = []\nfor i in range(0, ndoc, 4000):\n enc = tok(texts[i:i+4000], add_special_tokens=False).input_ids\n docs.extend(np.asarray(x, dtype=np.int32) for x in enc)\nlens = np.fromiter((len(d) for d in docs), dtype=np.int64, count=ndoc)\nalltok = np.concatenate([d for d in docs]).astype(np.int64)\nprint(f\"[{time.time()-t0:.0f}s] tokenized pool: {lens.sum()/1e6:.1f}M tokens\", flush=True)\n\n# --------------------------------- pool background counts (shared denominator)\nneg_uni = np.bincount(alltok, minlength=V).astype(np.float64)\nneg_big = np.bincount(bigrams(alltok), minlength=DB).astype(np.float64)\nneg_uni_ln = np.log(neg_uni + ALPHA) - np.log(neg_uni.sum() + ALPHA * V)\nneg_big_ln = np.log(neg_big + ALPHA) - np.log(neg_big.sum() + ALPHA * DB)\ndel alltok\nprint(f\"[{time.time()-t0:.0f}s] built pool background\", flush=True)\n\n# ---------------------------- per-register target counts from disclosed target\ndev = np.load(DEV).astype(np.int64)\nsplits = np.where(dev == EOS)[0]\nb = [0] + (splits + 1).tolist() + [len(dev)]\nuni_w = {r: None for r in REGISTERS}\nbig_w = {r: None for r in REGISTERS}\npos_uni = {r: np.zeros(V, np.float64) for r in REGISTERS}\npos_big = {r: np.zeros(DB, np.float64) for r in REGISTERS}\nfor k in range(len(b) - 1):\n d = dev[b[k]:b[k+1]]\n d = d[d != EOS]\n if len(d) < 20:\n continue\n raw = tok.decode(d.tolist())\n r = register_of(raw)\n if r == \"wiki\": # de-format encyclopedic text\n d = np.asarray(tok(normalize_target(raw), add_special_tokens=False).input_ids,\n dtype=np.int64)\n pos_uni[r] += np.bincount(d, minlength=V)\n if len(d) > 1:\n pos_big[r] += np.bincount(bigrams(d), minlength=DB)\nfor r in REGISTERS:\n uni_w[r] = (np.log(pos_uni[r] + ALPHA) - np.log(pos_uni[r].sum() + ALPHA * V)) - neg_uni_ln\n big_w[r] = (np.log(pos_big[r] + ALPHA) - np.log(pos_big[r].sum() + ALPHA * DB)) - neg_big_ln\nprint(f\"[{time.time()-t0:.0f}s] built {len(REGISTERS)} register targets \"\n f\"(tok: {{{', '.join(f'{r}:{int(pos_uni[r].sum())}' for r in REGISTERS)}}})\", flush=True)\n\n# ------------------------ presence-based scoring of every doc vs every register\nbest_reg = np.full(ndoc, -1, dtype=np.int64)\nbest_score = np.full(ndoc, -1e30, dtype=np.float64)\nfor i in range(ndoc):\n d = docs[i]; n = d.shape[0]\n if n < MIN_TOK:\n continue\n t = texts[i]\n if len(t) == 0 or len(_NONALPHA.sub(\"\", t)) / len(t) < MIN_ALPHA:\n continue\n di = d.astype(np.int64)\n u = np.unique(di)\n ub = np.unique(bigrams(di)) if n > 1 else np.array([0], dtype=np.int64)\n bi, bs = -1, -1e30\n for ri, r in enumerate(REGISTERS):\n s = uni_w[r][u].mean() + LAMBDA * big_w[r][ub].mean()\n if s > bs:\n bs, bi = s, ri\n best_reg[i] = bi; best_score[i] = bs\n if (i & 0x3FFF) == 0:\n print(f\"[{time.time()-t0:.0f}s] scoring {i}/{ndoc}\", flush=True)\nprint(f\"[{time.time()-t0:.0f}s] scored docs\", flush=True)\n\n# ------------------------------------- rank within register, interleave by quota\nranked = {}\nfor ri, r in enumerate(REGISTERS):\n idx = np.where(best_reg == ri)[0]\n idx = idx[np.argsort(-best_score[idx])]\n ranked[r] = idx.tolist()\n print(f\" {r}: {len(idx)} docs routed\", flush=True)\n\nptr = {r: 0 for r in REGISTERS}\ncum = {r: 0 for r in REGISTERS}\nselection, seen, tot = [], set(), 0\nwhile tot < COVER_TOKENS:\n # pick the register that is furthest below its token quota and still has docs\n cand = [r for r in REGISTERS if ptr[r] < len(ranked[r])]\n if not cand:\n break\n r = min(cand, key=lambda r: cum[r] / QUOTA[r])\n j = ranked[r][ptr[r]]; ptr[r] += 1\n h = hash(texts[j])\n if h in seen:\n continue\n seen.add(h)\n selection.append(int(ids[j]))\n cum[r] += int(lens[j]); tot += int(lens[j])\n\njson.dump(selection, open(OUT, \"w\"))\nmix = {r: round(cum[r] / max(1, tot), 3) for r in REGISTERS}\nprint(f\"[{time.time()-t0:.0f}s] wrote {len(selection)} ids covering ~{tot/1e6:.1f}M tokens\", flush=True)\nprint(f\"register token mix of emitted list: {mix}\", flush=True)\n", "structuredPatch": [{"oldStart": 131, "oldLines": 8, "newStart": 131, "newLines": 9, "lines": [" f\"(tok: {{{', '.join(f'{r}:{int(pos_uni[r].sum())}' for r in REGISTERS)}}})\", flush=True)", " ", " # ------------------------ presence-based scoring of every doc vs every register", "-best_reg = np.full(ndoc, -1, dtype=np.int64)", "-best_score = np.full(ndoc, -1e30, dtype=np.float64)", "+NR = len(REGISTERS)", "+scoremat = np.full((ndoc, NR), -1e30, dtype=np.float64)", "+gate_ok = np.zeros(ndoc, dtype=bool)", " for i in range(ndoc):", " d = docs[i]; n = d.shape[0]", " if n < MIN_TOK:"]}, {"oldStart": 143, "oldLines": 36, "newStart": 144, "newLines": 34, "lines": [" di = d.astype(np.int64)", " u = np.unique(di)", " ub = np.unique(bigrams(di)) if n > 1 else np.array([0], dtype=np.int64)", "- bi, bs = -1, -1e30", " for ri, r in enumerate(REGISTERS):", "- s = uni_w[r][u].mean() + LAMBDA * big_w[r][ub].mean()", "- if s > bs:", "- bs, bi = s, ri", "- best_reg[i] = bi; best_score[i] = bs", "+ scoremat[i, ri] = uni_w[r][u].mean() + LAMBDA * big_w[r][ub].mean()", "+ gate_ok[i] = True", " if (i & 0x3FFF) == 0:", " print(f\"[{time.time()-t0:.0f}s] scoring {i}/{ndoc}\", flush=True)", "-print(f\"[{time.time()-t0:.0f}s] scored docs\", flush=True)", "+print(f\"[{time.time()-t0:.0f}s] scored docs ({gate_ok.sum()} passed gates)\", flush=True)", " ", "-# ------------------------------------- rank within register, interleave by quota", "-ranked = {}", "-for ri, r in enumerate(REGISTERS):", "- idx = np.where(best_reg == ri)[0]", "- idx = idx[np.argsort(-best_score[idx])]", "- ranked[r] = idx.tolist()", "- print(f\" {r}: {len(idx)} docs routed\", flush=True)", "+# ------------- independent per-register ranked lists, interleaved by equal quota", "+ok = np.where(gate_ok)[0]", "+ranked = {r: ok[np.argsort(-scoremat[ok, ri])] for ri, r in enumerate(REGISTERS)}", " ", " ptr = {r: 0 for r in REGISTERS}", " cum = {r: 0 for r in REGISTERS}", "-selection, seen, tot = [], set(), 0", "+selection, seen, chosen, tot = [], set(), set(), 0", " while tot < COVER_TOKENS:", "- # pick the register that is furthest below its token quota and still has docs", " cand = [r for r in REGISTERS if ptr[r] < len(ranked[r])]", " if not cand:", " break", " r = min(cand, key=lambda r: cum[r] / QUOTA[r])", "- j = ranked[r][ptr[r]]; ptr[r] += 1", "+ # advance past docs already claimed by another register", "+ while ptr[r] < len(ranked[r]) and ranked[r][ptr[r]] in chosen:", "+ ptr[r] += 1", "+ if ptr[r] >= len(ranked[r]):", "+ continue", "+ j = int(ranked[r][ptr[r]]); ptr[r] += 1", "+ chosen.add(j)", " h = hash(texts[j])", "- if h in seen:", "+ if h in seen: # exact-duplicate guard", " continue", " seen.add(h)", " selection.append(int(ids[j]))"]}], "userModified": false, "replaceAll": false}#46
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolBashRun balanced+normalized curate
args
{
"command": "cd /workspace; export TOKENIZERS_PARALLELISM=true; echo \"--- running balanced+normalized curate.py ---\"; time python3 submission/curate.py 2>&1 | grep -vE \"sequence length|Token indices\"",
"description": "Run balanced+normalized curate"
}Bash result
Command did not complete within its 120s timeout and was moved to the background (ID: bniprzajj). Output is being written to: /tmp/claude-0/-workspace/4bb4d0fa-198e-4747-8902-b67961f6b67a/tasks/bniprzajj.output. You will be notified when it completes. To check interim output, use Read on that file path.
Session cwd remains /workspace; directory changes made by the backgrounded command do not apply to subsequent commands.
[metadata] {"noOutputExpected": false, "backgroundTaskId": "bniprzajj", "timedOutAfterMs": 120000, "backgroundCwdHint": "Session cwd remains /workspace; directory changes made by the backgrounded command do not apply to subsequent commands."}#47
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolBashCheck balanced+normalized curate progress
args
{
"command": "sleep 80; tail -8 /tmp/claude-0/-workspace/4bb4d0fa-198e-4747-8902-b67961f6b67a/tasks/bniprzajj.output",
"description": "Check balanced+normalized curate progress"
}Bash result
[155s] scoring 180224/182016
[156s] scored docs (171150 passed gates)
[156s] wrote 35442 ids covering ~30.0M tokens
register token mix of emitted list: {'wiki': 0.25, 'web': 0.25, 'news': 0.25, 'code_qa': 0.25}
real 2m41.046s
user 15m40.170s
sys 0m3.080s
[stdout]
[155s] scoring 180224/182016
[156s] scored docs (171150 passed gates)
[156s] wrote 35442 ids covering ~30.0M tokens
register token mix of emitted list: {'wiki': 0.25, 'web': 0.25, 'news': 0.25, 'code_qa': 0.25}
real 2m41.046s
user 15m40.170s
sys 0m3.080s
[metadata] {"noOutputExpected": false}#48
narration
Now exactly 25% each register. Let me validate it with dev_eval — the real test.
toolBashEval balanced+normalized selection
args
{
"command": "cd /workspace; export TOKENIZERS_PARALLELISM=false; cat /workspace/.timer/remaining_secs; echo \"=== dev_eval on BALANCED+NORMALIZED selection ===\"; time bash dev_eval.sh /workspace/submission/selection.json 2>&1 | grep -vE \"sequence length|Token indices\"",
"description": "Eval balanced+normalized selection"
}Bash result
3187
=== dev_eval on BALANCED+NORMALIZED selection ===
packed 12000000 tokens from 35442 ids -> /tmp/dev_train.npy
HELDOUT_PPL 331.696 loss 5.8042 train_s 56 /tmp/dev_train.npy
(dev PPL on multi_dev; official uses a hidden disjoint sample of the same HQ domain)
real 1m43.846s
user 1m38.270s
sys 0m2.730s
[stdout]
3187
=== dev_eval on BALANCED+NORMALIZED selection ===
packed 12000000 tokens from 35442 ids -> /tmp/dev_train.npy
HELDOUT_PPL 331.696 loss 5.8042 train_s 56 /tmp/dev_train.npy
(dev PPL on multi_dev; official uses a hidden disjoint sample of the same HQ domain)
real 1m43.846s
user 1m38.270s
sys 0m2.730s
[metadata] {"noOutputExpected": false}#49
narration
331.7 — big jump from 441 (unbalanced) and 485 (random). Register balancing works. Let me back this up.
Key insight for further gains: the held-out PPL is token-weighted over the concatenated target (non-overlapping windows), so the eval effectively weights registers by their token share, not equally. My earlier measurement showed the target's token shares are ~wiki 16%, web 21%, news 35%, code 27% — not 25% each. Matching the training mixture to the target's token proportions should reduce PPL further. Let me back up and make the quota data-driven from the target's own token shares.
toolBashBack up current best selection
args
{
"command": "cp /workspace/submission/selection.json /tmp/selection_balanced_equal_332.json; echo \"backed up 331.7 selection\"",
"description": "Back up current best selection"
}Bash result
backed up 331.7 selection
[stdout]
backed up 331.7 selection
[metadata] {"noOutputExpected": false}#50
reasoning (redacted)
encrypted chain-of-thought — recorded, not exposed
toolEdit
args
{
"replace_all": false,
"file_path": "/workspace/submission/curate.py",
"old_string": "REGISTERS = [\"wiki\", \"web\", \"news\", \"code_qa\"]\nQUOTA = {r: 0.25 for r in REGISTERS} # equal-parts target mixture",
"new_string": "REGISTERS = [\"wiki\", \"web\", \"news\", \"code_qa\"]\n# QUOTA (target training mixture) is set below to the target's own per-register\n# TOKEN proportions: held-out PPL is window-weighted over the concatenated target,\n# so the eval effectively weights each register by its token share -- matching that\n# share in training minimizes the token-weighted average loss."
}Edit result
The file /workspace/submission/curate.py has been updated successfully. (file state is current in your context — no need to Read it back)
[metadata] {"filePath": "/workspace/submission/curate.py", "oldString": "REGISTERS = [\"wiki\", \"web\", \"news\", \"code_qa\"]\nQUOTA = {r: 0.25 for r in REGISTERS} # equal-parts target mixture", "newString": "REGISTERS = [\"wiki\", \"web\", \"news\", \"code_qa\"]\n# QUOTA (target training mixture) is set below to the target's own per-register\n# TOKEN proportions: held-out PPL is window-weighted over the concatenated target,\n# so the eval effectively weights each register by its token share -- matching that\n# share in training minimizes the token-weighted average loss.", "originalFile": "\"\"\"Curate a pretraining subset from a raw web pool by register-balanced DSIR.\n\nThe disclosed evaluation target is a BROAD, multi-domain HQ English mix in equal\nparts: encyclopedic (Wikipedia), general web prose, news, and technical Q&A.\nHeld-out perplexity averages over all four registers, so a training set that\nover-weights one register (e.g. news, which is easiest to match) leaves the\nothers under-trained and drags the average up. We therefore select in two steps:\n\n (1) DSIR importance scoring, PER REGISTER. We split the disclosed target\n (data/multi_dev.npy, GPT-2 tokens) into four register buckets by surface\n signatures, and for each register r estimate a bag-of-n-grams importance\n weight\n w_r(gram) = log p_target_r(gram) - log p_pool(gram)\n (unigrams + hashed bigrams; p_pool from the whole pool as background).\n Each pool document is scored against every register using a PRESENCE-based\n mean over its unique n-gram types (repetition-robust), behind a length\n floor and a letter-fraction gate that reject fragments and non-prose\n (directory listings, tables). A document is routed to its best-fit\n register (argmax score).\n\n (2) Balanced quota fill. Within each register we rank routed documents by\n score, then interleave the four ranked lists by EQUAL token quota, so the\n first 12M tokens the trainer consumes are ~25% from each register --\n matching the equal-parts evaluation mixture.\n\nThis is Data Selection via Importance Resampling (Xie et al. 2023) with a\nper-domain target, plus classic quality gating. No labels, no internet: the only\nsupervision is the disclosed target token stream.\n\"\"\"\nimport json, time, re, numpy as np\nfrom transformers import AutoTokenizer\n\nPOOL = \"/workspace/data/pool.jsonl\"\nDEV = \"/workspace/data/multi_dev.npy\"\nOUT = \"/workspace/submission/selection.json\"\n\nEOS = 50256\nV = 50257 # GPT-2 vocab\nDB = 1 << 20 # hashed bigram buckets\nALPHA = 0.1 # additive smoothing\nMIN_TOK = 100 # length floor: drop fragments\nMIN_ALPHA = 0.55 # letter-fraction floor: reject non-prose\nLAMBDA = 0.5 # weight of bigram term relative to unigram term\nCOVER_TOKENS = 30_000_000 # emit ids covering ~2.5x the 12M budget\nREGISTERS = [\"wiki\", \"web\", \"news\", \"code_qa\"]\nQUOTA = {r: 0.25 for r in REGISTERS} # equal-parts target mixture\n\n_NONALPHA = re.compile(r\"[^A-Za-z]\")\n_NEWS = re.compile(r\"\\b(reuters|said|told|according to|\"\n r\"on (monday|tuesday|wednesday|thursday|friday|saturday|sunday))\\b\")\n\ndef normalize_target(t):\n \"\"\"Strip WikiText-103 surface formatting so the wiki target matches on CONTENT\n vocabulary, not on markup absent from the raw pool (which would otherwise make\n every pool doc look un-wiki-like and starve the encyclopedic quota).\"\"\"\n t = t.replace(\"@-@\", \"-\").replace(\"@,@\", \",\").replace(\"@.@\", \".\")\n t = re.sub(r\"\\s+([,.;:!?)])\", r\"\\1\", t) # drop space before punctuation\n t = re.sub(r\"([(\\[])\\s+\", r\"\\1\", t) # drop space after open bracket\n t = re.sub(r\"\\s+\", \" \", t) # collapse whitespace\n return t\n\ndef register_of(t):\n \"\"\"Surface-signature register label for a piece of text.\"\"\"\n if \"@-@\" in t or \"@,@\" in t or \"@.@\" in t:\n return \"wiki\"\n if (\"<p>\" in t or \"<code>\" in t or \"</\" in t or \"def \" in t or \"import \" in t\n or (\"{\" in t and \"}\" in t and \";\" in t)):\n return \"code_qa\"\n if _NEWS.search(t.lower()):\n return \"news\"\n return \"web\"\n\ndef bigrams(a):\n return ((a[:-1] * np.int64(1000003) + a[1:]) & np.int64(DB - 1))\n\nt0 = time.time()\ntok = AutoTokenizer.from_pretrained(\"gpt2\")\n\n# ------------------------------------------------------------------ load pool\nids, texts = [], []\nwith open(POOL) as f:\n for line in f:\n r = json.loads(line)\n ids.append(r[\"id\"]); texts.append(r[\"text\"])\nids = np.asarray(ids, dtype=np.int64)\nndoc = len(ids)\nprint(f\"[{time.time()-t0:.0f}s] loaded pool: {ndoc} docs\", flush=True)\n\n# ------------------------------------------------------ tokenize pool (batched)\ndocs = []\nfor i in range(0, ndoc, 4000):\n enc = tok(texts[i:i+4000], add_special_tokens=False).input_ids\n docs.extend(np.asarray(x, dtype=np.int32) for x in enc)\nlens = np.fromiter((len(d) for d in docs), dtype=np.int64, count=ndoc)\nalltok = np.concatenate([d for d in docs]).astype(np.int64)\nprint(f\"[{time.time()-t0:.0f}s] tokenized pool: {lens.sum()/1e6:.1f}M tokens\", flush=True)\n\n# --------------------------------- pool background counts (shared denominator)\nneg_uni = np.bincount(alltok, minlength=V).astype(np.float64)\nneg_big = np.bincount(bigrams(alltok), minlength=DB).astype(np.float64)\nneg_uni_ln = np.log(neg_uni + ALPHA) - np.log(neg_uni.sum() + ALPHA * V)\nneg_big_ln = np.log(neg_big + ALPHA) - np.log(neg_big.sum() + ALPHA * DB)\ndel alltok\nprint(f\"[{time.time()-t0:.0f}s] built pool background\", flush=True)\n\n# ---------------------------- per-register target counts from disclosed target\ndev = np.load(DEV).astype(np.int64)\nsplits = np.where(dev == EOS)[0]\nb = [0] + (splits + 1).tolist() + [len(dev)]\nuni_w = {r: None for r in REGISTERS}\nbig_w = {r: None for r in REGISTERS}\npos_uni = {r: np.zeros(V, np.float64) for r in REGISTERS}\npos_big = {r: np.zeros(DB, np.float64) for r in REGISTERS}\nfor k in range(len(b) - 1):\n d = dev[b[k]:b[k+1]]\n d = d[d != EOS]\n if len(d) < 20:\n continue\n raw = tok.decode(d.tolist())\n r = register_of(raw)\n if r == \"wiki\": # de-format encyclopedic text\n d = np.asarray(tok(normalize_target(raw), add_special_tokens=False).input_ids,\n dtype=np.int64)\n pos_uni[r] += np.bincount(d, minlength=V)\n if len(d) > 1:\n pos_big[r] += np.bincount(bigrams(d), minlength=DB)\nfor r in REGISTERS:\n uni_w[r] = (np.log(pos_uni[r] + ALPHA) - np.log(pos_uni[r].sum() + ALPHA * V)) - neg_uni_ln\n big_w[r] = (np.log(pos_big[r] + ALPHA) - np.log(pos_big[r].sum() + ALPHA * DB)) - neg_big_ln\nprint(f\"[{time.time()-t0:.0f}s] built {len(REGISTERS)} register targets \"\n f\"(tok: {{{', '.join(f'{r}:{int(pos_uni[r].sum())}' for r in REGISTERS)}}})\", flush=True)\n\n# ------------------------ presence-based scoring of every doc vs every register\nNR = len(REGISTERS)\nscoremat = np.full((ndoc, NR), -1e30, dtype=np.float64)\ngate_ok = np.zeros(ndoc, dtype=bool)\nfor i in range(ndoc):\n d = docs[i]; n = d.shape[0]\n if n < MIN_TOK:\n continue\n t = texts[i]\n if len(t) == 0 or len(_NONALPHA.sub(\"\", t)) / len(t) < MIN_ALPHA:\n continue\n di = d.astype(np.int64)\n u = np.unique(di)\n ub = np.unique(bigrams(di)) if n > 1 else np.array([0], dtype=np.int64)\n for ri, r in enumerate(REGISTERS):\n scoremat[i, ri] = uni_w[r][u].mean() + LAMBDA * big_w[r][ub].mean()\n gate_ok[i] = True\n if (i & 0x3FFF) == 0:\n print(f\"[{time.time()-t0:.0f}s] scoring {i}/{ndoc}\", flush=True)\nprint(f\"[{time.time()-t0:.0f}s] scored docs ({gate_ok.sum()} passed gates)\", flush=True)\n\n# ------------- independent per-register ranked lists, interleaved by equal quota\nok = np.where(gate_ok)[0]\nranked = {r: ok[np.argsort(-scoremat[ok, ri])] for ri, r in enumerate(REGISTERS)}\n\nptr = {r: 0 for r in REGISTERS}\ncum = {r: 0 for r in REGISTERS}\nselection, seen, chosen, tot = [], set(), set(), 0\nwhile tot < COVER_TOKENS:\n cand = [r for r in REGISTERS if ptr[r] < len(ranked[r])]\n if not cand:\n break\n r = min(cand, key=lambda r: cum[r] / QUOTA[r])\n # advance past docs already claimed by another register\n while ptr[r] < len(ranked[r]) and ranked[r][ptr[r]] in chosen:\n ptr[r] += 1\n if ptr[r] >= len(ranked[r]):\n continue\n j = int(ranked[r][ptr[r]]); ptr[r] += 1\n chosen.add(j)\n h = hash(texts[j])\n if h in seen: # exact-duplicate guard\n continue\n seen.add(h)\n selection.append(int(ids[j]))\n cum[r] += int(lens[j]); tot += int(lens[j])\n\njson.dump(selection, open(OUT, \"w\"))\nmix = {r: round(cum[r] / max(1, tot), 3) for r in REGISTERS}\nprint(f\"[{time.time()-t0:.0f}s] wrote {len(selection)} ids covering ~{tot/1e6:.1f}M tokens\", flush=True)\nprint(f\"register token mix of emitted list: {mix}\", flush=True)\n", "structuredPatch": [{"oldStart": 43, "oldLines": 7, "newStart": 43, "newLines": 10, "lines": [" LAMBDA = 0.5 # weight of bigram term relative to unigram term", " COVER_TOKENS = 30_000_000 # emit ids covering ~2.5x the 12M budget", " REGISTERS = [\"wiki\", \"web\", \"news\", \"code_qa\"]", "-QUOTA = {r: 0.25 for r in REGISTERS} # equal-parts target mixture", "+# QUOTA (target training mixture) is set below to the target's own per-register", "+# TOKEN proportions: held-out PPL is window-weighted over the concatenated target,", "+# so the eval effectively weights each register by its token share -- matching that", "+# share in training minimizes the token-weighted average loss.", " ", " _NONALPHA = re.compile(r\"[^A-Za-z]\")", " _NEWS = re.compile(r\"\\b(reuters|said|told|according to|\""]}], "userModified": false, "replaceAll": false}