File size: 1,798 Bytes
70a25ae | 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 | import hashlib
import json
import os
from huggingface_hub import hf_hub_download
SEED = b"sprout-v1-split"
VAL_FRAC = 0.01
OUT = "data"
SEP = "<|endoftext|>"
REVISION = "f54c09fd23315a6f9c86f9dc80f725de7d8f9c64"
def is_val(story):
"""Returns True if the story should be held out for validation, False otherwise."""
h = hashlib.blake2b(story.encode("utf-8"), digest_size=8, salt=SEED).digest()
return int.from_bytes(h, "big") % 10000 < VAL_FRAC * 10000
def stories(path):
"""Yield one story at a time. A lone <|endoftext|> line separates them."""
lines = []
with open(path, encoding="utf-8", newline="") as f:
for line in f:
if line.rstrip("\r\n") != SEP:
lines.append(line)
continue
story = "".join(lines).strip()
lines = []
if story:
yield story
story = "".join(lines).strip()
if story:
yield story
def main():
os.makedirs(OUT, exist_ok=True)
raw = hf_hub_download(
"roneneldan/TinyStories",
"TinyStoriesV2-GPT4-train.txt",
repo_type="dataset",
revision=REVISION,
local_dir=OUT,
)
train = open(f"{OUT}/train.jsonl", "w", encoding="utf-8")
val = open(f"{OUT}/val.jsonl", "w", encoding="utf-8")
n = [0, 0]
nbytes = 0
for story in stories(raw):
held_out = is_val(story)
(val if held_out else train).write(
json.dumps({"text": story}, ensure_ascii=False) + "\n")
n[held_out] += 1
nbytes += len(story.encode("utf-8"))
train.close()
val.close()
total = n[0] + n[1]
print(f"{total:,} stories, {nbytes / 1e9:.3f} GB")
print(f"train {n[0]:,} val {n[1]:,} ({n[1] / total:.2%})")
if __name__ == "__main__":
main()
|