| """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()) |
|
|
