File size: 2,371 Bytes
b999890 | 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 63 64 65 66 67 | """Build the two byte streams the search needs: one to search on, one held out.
Default source is text already on this machine (Python's own standard library), so no
download is needed. For a standard benchmark instead, download enwik8 yourself
(http://mattmahoney.net/dc/enwik8.zip, 36 MB), unzip it, and pass --from-file enwik8.
python make_data.py # 1 MB search + 1 MB held out, local text
python make_data.py --from-file enwik8 # same sizes, cut from enwik8
python make_data.py --size-mb 4 # bigger streams (slower evaluations)
"""
from __future__ import annotations
import argparse
import random
import sysconfig
from pathlib import Path
MB = 1024 * 1024
MAX_SIZE_MB = 64
def gather_local_text(root: Path, wanted_bytes: int) -> bytes:
files = sorted(path for path in root.rglob("*.py") if "test" not in path.parts)
random.Random(0).shuffle(files)
chunks, total = [], 0
for path in files:
try:
chunk = path.read_bytes()
except OSError:
continue
chunks.append(chunk)
total += len(chunk)
if total >= wanted_bytes:
break
return b"".join(chunks)
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--from-file", help="cut both streams from this file")
parser.add_argument("--size-mb", type=float, default=1.0, help="size of each stream")
parser.add_argument("--out", default="data")
args = parser.parse_args()
if not 0 < args.size_mb <= MAX_SIZE_MB:
parser.error(f"--size-mb must be between 0 and {MAX_SIZE_MB}")
size = int(args.size_mb * MB)
if args.from_file:
raw = Path(args.from_file).read_bytes()[: 2 * size]
else:
raw = gather_local_text(Path(sysconfig.get_paths()["stdlib"]), 2 * size)
if len(raw) < 2 * size:
parser.error(f"source has only {len(raw)} bytes; need {2 * size}")
out = Path(args.out)
out.mkdir(parents=True, exist_ok=True)
(out / "search.bin").write_bytes(raw[:size])
(out / "holdout.bin").write_bytes(raw[size: 2 * size])
print(f"wrote {out / 'search.bin'} and {out / 'holdout.bin'} ({size} bytes each)")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|