File size: 7,015 Bytes
f91d419 | 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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | """Pinned story download, document splits and a tokenizer trained on train only."""
import hashlib
import json
import shutil
from pathlib import Path
import requests
import numpy as np
from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders
from .model import SPECIAL_TOKENS, PAD, BOS, EOS
def digest(path):
h = hashlib.sha256()
with open(path, "rb") as f:
for chunk in iter(lambda: f.read(1024 * 1024), b""):
h.update(chunk)
return h.hexdigest()
def download_stories(revision, filename, count, seen):
url = f"https://huggingface.co/datasets/roneneldan/TinyStories/resolve/{revision}/{filename}"
stories, buffer = [], ""
with requests.get(url, stream=True, timeout=(30, 120)) as response:
response.raise_for_status()
response.encoding = "utf-8"
for chunk in response.iter_content(1024 * 128, decode_unicode=True):
buffer += chunk
while "<|endoftext|>" in buffer:
story, buffer = buffer.split("<|endoftext|>", 1)
story = story.strip()
key = hashlib.sha256(" ".join(story.split()).encode()).hexdigest()
if len(story) >= 100 and key not in seen:
seen.add(key)
stories.append(story)
if len(stories) % 5000 == 0:
print(f"{filename}: {len(stories)} stories", flush=True)
if len(stories) == count:
return stories
raise ValueError(f"Only found {len(stories)} unique stories, requested {count}")
def prepare(args):
if min(args.stories, args.heldout, args.tokenizer_stories) < 1 or args.length < 4 or not 261 <= args.vocab_size <= 65535:
raise ValueError("Positive document counts, length >=4 and vocabulary 261..65535 required")
out = Path(args.output)
if (out / "manifest.json").exists():
raise ValueError("Prepared data already exists; use a new output directory")
out.mkdir(parents=True, exist_ok=True)
metadata = requests.get("https://huggingface.co/api/datasets/roneneldan/TinyStories", timeout=30)
metadata.raise_for_status()
revision = args.revision or metadata.json()["sha"]
seen = set()
# Held-out source first: reject exact document duplicates from training.
heldout = download_stories(revision, "TinyStories-valid.txt", args.heldout * 2, seen)
train = download_stories(revision, "TinyStories-train.txt", args.stories, seen)
splits = {"train": train, "validation": heldout[::2], "test": heldout[1::2]}
existing = getattr(args, "tokenizer", None)
if existing:
tokenizer = Tokenizer.from_file(existing)
if any(tokenizer.token_to_id(t) != i for i, t in enumerate(SPECIAL_TOKENS)):
raise ValueError("Incompatible special-token IDs")
shutil.copy2(existing, out / "tokenizer.json")
else:
tokenizer = Tokenizer(models.BPE(unk_token="[UNK]"))
tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
tokenizer.decoder = decoders.ByteLevel()
tokenizer.train_from_iterator(train[:args.tokenizer_stories], trainers.BpeTrainer(vocab_size=args.vocab_size,
special_tokens=SPECIAL_TOKENS, initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), min_frequency=2))
tokenizer.save(str(out / "tokenizer.json"))
manifest = {"dataset": "roneneldan/TinyStories", "revision": revision, "license": "cdla-sharing-1.0",
"source": "https://huggingface.co/datasets/roneneldan/TinyStories", "synthetic": True,
"selection": "First unique documents in pinned source order; alternating validation/test documents",
"tokenizer_train_documents": min(args.tokenizer_stories, len(train)), "vocab_size": tokenizer.get_vocab_size(),
"max_length": args.length, "tokenizer_sha256": digest(out / "tokenizer.json"), "splits": {}}
if existing:
manifest["tokenizer_source"] = str(existing)
for name, stories in splits.items():
# Nonoverlapping windows within documents, never across split boundaries.
tokens = windows = complete = 0
raw_path = out / f"{name}.rows.tmp"
with raw_path.open("wb") as stream:
for start in range(0, len(stories), 1000):
rows = []
for encoded in tokenizer.encode_batch(stories[start:start+1000]):
ids = encoded.ids
tokens += len(ids)
complete += int(len(ids) <= args.length - 2)
for offset in range(0, len(ids), args.length - 2):
segment = [BOS] + ids[offset:offset + args.length - 2]
if offset + args.length - 2 >= len(ids):
segment += [EOS]
rows.append(segment + [PAD] * (args.length - len(segment)))
np.asarray(rows, dtype=np.uint16).tofile(stream)
windows += len(rows)
path = out / f"{name}.npy"
raw = np.memmap(raw_path, dtype=np.uint16, mode="r", shape=(windows, args.length))
np.save(path, raw)
del raw
raw_path.unlink()
manifest["splits"][name] = {"documents": len(stories), "windows": windows, "text_tokens": tokens,
"complete_story_windows": complete,
"documents_sha256": hashlib.sha256("\n<|endoftext|>\n".join(stories).encode()).hexdigest(), "array_sha256": digest(path)}
print(name, manifest["splits"][name], flush=True)
(out / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
def load_arrays(directory):
directory = Path(directory)
manifest = json.loads((directory / "manifest.json").read_text(encoding="utf-8"))
if digest(directory / "tokenizer.json") != manifest["tokenizer_sha256"]:
raise ValueError("Tokenizer fingerprint mismatch")
arrays = {}
for name in ("train", "validation", "test"):
path = directory / f"{name}.npy"
if digest(path) != manifest["splits"][name]["array_sha256"]:
raise ValueError(f"Corrupted {name} array")
arrays[name] = np.load(path, mmap_mode="r")
array=arrays[name]
error=None
if array.ndim!=2 or array.shape[1]!=manifest["max_length"] or array.dtype!=np.uint16:
error=f"Invalid {name} array shape or dtype"
else:
for start in range(0,len(array),4096):
chunk=array[start:start+4096]
if np.any(chunk[:,0]!=BOS) or np.any(chunk>=manifest["vocab_size"]):
error=f"Invalid {name} token IDs or BOS"
break
if np.any((chunk[:,:-1]==PAD)&(chunk[:,1:]!=PAD)):
error=f"{name} must be right-padded for causal attention"
break
if error:
for opened in arrays.values():opened._mmap.close()
raise ValueError(error)
return arrays, manifest
|