File size: 2,743 Bytes
d0a770f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | #!/usr/bin/env python3
"""
make_haystack.py -- produce ONE UTF-8 text haystack shared by every peer model.
Why text and not token ids: each peer tokenizes differently, so the only thing
that can be held constant across models is the CHARACTERS. peer_gen.py then cuts
each model's prompt to an exact count of its own tokens from this same story.
Source is the b3_16384_32768 val shard of unparallelled/BluScriptCP-data --
GPT-2 uint16 token ids, decoded back to text here. That shard was held out from
every training arm of ours, and the peers never saw our data at all, so it is
out-of-distribution for everything being compared.
usage: make_haystack.py --out hay.txt --tokens 4000000
"""
import argparse
import os
EOT = "<|endoftext|>"
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--out", required=True)
ap.add_argument("--tokens", type=int, default=4_000_000,
help="GPT-2 tokens to decode; peers need prompt_len*(n+2) of "
"THEIR tokens, and their tokenizers are ~15%% denser than "
"GPT-2 on English, so over-provision")
ap.add_argument("--repo", default="unparallelled/BluScriptCP-data")
ap.add_argument("--file", default="cp_data_by_length/b3_16384_32768/val_shard_00000.bin")
args = ap.parse_args()
import numpy as np, tiktoken
from huggingface_hub import hf_hub_download
p = hf_hub_download(args.repo, args.file, repo_type="dataset")
toks = np.fromfile(p, dtype=np.uint16, count=args.tokens)
print(f" read {toks.size:,} GPT-2 tokens from {os.path.basename(p)}")
enc = tiktoken.get_encoding("gpt2")
# decode in chunks: one 4M-token decode builds a huge intermediate string
parts, CH = [], 200_000
for s in range(0, toks.size, CH):
parts.append(enc.decode(toks[s:s + CH].astype(np.int64).tolist()))
text = "".join(parts)
# STRIP the literal "<|endoftext|>". The corpus was tokenized with
# encode_ordinary, so GPT-2 id 50256 decodes back to that literal string --
# and the peers disagree about it: Qwen3 re-encodes it to ONE special token
# (151643, a document reset that truncates attention), while Spark-X2.5
# treats it as 7 ordinary tokens. Leaving it in would hand the two models
# materially different haystacks from identical text.
n_eot = text.count(EOT)
text = text.replace(EOT, " ")
print(f" stripped {n_eot:,} '{EOT}' markers "
f"(~{n_eot * 16384 / max(1, toks.size):.1f} per 16k-token prompt)")
with open(args.out, "w", encoding="utf-8") as f:
f.write(text)
print(f" wrote {args.out} {len(text):,} chars ({len(text)/toks.size:.2f} chars/gpt2-token)")
if __name__ == "__main__":
main()
|