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