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
68
69
70
71
72
73
74
75
76
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()
|
