1 files changed, 77 insertions, 0 deletions
diff --git a/experiments/prepare_shakespeare_char.py b/experiments/prepare_shakespeare_char.py
new file mode 100755
index 0000000..f09fc80
--- /dev/null
+++ b/experiments/prepare_shakespeare_char.py
@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+"""Download and deterministically tokenize the frozen Tiny Shakespeare data."""
+import argparse
+import hashlib
+import os
+import pickle
+import urllib.request
+
+import numpy as np
+
+
+URL = (
+    "https://raw.githubusercontent.com/karpathy/char-rnn/master/"
+    "data/tinyshakespeare/input.txt"
+)
+EXPECTED = {
+    "input.txt":
+        "86c4e6aa9db7c042ec79f339dcb96d42b0075e16b8fc2e86bf0ca57e2dc565ed",
+    "train.bin":
+        "6ec305602a99ac2802745a134e1f5e33e2231b4855525b00b9aebb730ac2626f",
+    "val.bin":
+        "d37d30cc0c8327c270d493299c3dca54135f6d5f1c9ef60cda78076e311204b1",
+}
+
+
+def sha256(path):
+    digest = hashlib.sha256()
+    with open(path, "rb") as handle:
+        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def verify(path, expected):
+    actual = sha256(path)
+    if actual != expected:
+        raise RuntimeError(
+            f"hash mismatch for {path}: expected {expected}, found {actual}"
+        )
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--out-dir", required=True)
+    args = parser.parse_args()
+    root = os.path.abspath(args.out_dir)
+    os.makedirs(root, exist_ok=True)
+    input_path = os.path.join(root, "input.txt")
+    if not os.path.isfile(input_path):
+        urllib.request.urlretrieve(URL, input_path)
+    verify(input_path, EXPECTED["input.txt"])
+
+    with open(input_path, encoding="utf-8") as handle:
+        data = handle.read()
+    chars = sorted(set(data))
+    stoi = {char: index for index, char in enumerate(chars)}
+    itos = {index: char for char, index in stoi.items()}
+    encoded = np.asarray([stoi[char] for char in data], dtype=np.uint16)
+    boundary = int(len(encoded) * 0.9)
+    encoded[:boundary].tofile(os.path.join(root, "train.bin"))
+    encoded[boundary:].tofile(os.path.join(root, "val.bin"))
+    with open(os.path.join(root, "meta.pkl"), "wb") as handle:
+        pickle.dump(
+            {"vocab_size": len(chars), "itos": itos, "stoi": stoi},
+            handle,
+        )
+
+    for name, expected in EXPECTED.items():
+        verify(os.path.join(root, name), expected)
+    print(
+        f"prepared Tiny Shakespeare at {root}: "
+        f"{boundary} train, {len(encoded) - boundary} validation tokens"
+    )
+
+
+if __name__ == "__main__":
+    main()