BluScriptCP / scripts /peer /make_haystack.py

Upload scripts/peer/make_haystack.py with huggingface_hub

d0a770f verified about 2 months ago

2.74 kB

#!/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()

	#!/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()