Playbook: run the Tiny Storyteller yourself

This is the short, practical guide. The long, explained version is the notebook tiny_llm_from_scratch.ipynb, which tells the whole thing as a story about Momo the monkey and Ella the elephant. Pick the path that fits you.

What you get

A 5.9-million-parameter GPT-style language model, trained from scratch on simple children's stories (TinyStories) plus a thousand generated stories about Momo the monkey and Ella the elephant. Give it the first few words and it continues with a short, simple story. It is a learning project, not a product: expect charming nonsense now and then, and expect the Momo stories to be recited from memory rather than invented (the notebook explains why).

Path A: run it in Google Colab (no install, free GPU)

Open https://colab.research.google.com and choose Upload notebook.
Pick tiny_llm_from_scratch.ipynb (download it from this repo first).
In Colab go to Runtime > Change runtime type and choose T4 GPU.
Press Run all. Everything above the training section runs in well under a minute. Training takes about 2 minutes on a T4. The whole notebook is about 5 minutes.
Scroll to Part 8 to read the generated stories, and change the prompts to your own.

The publish cell in Part 9 uploads the model to your own Hugging Face account, as a private repo. It needs a Hugging Face access token with write permission. The easiest way on Colab: open the Secrets panel (the key icon on the left), add a secret called HF_TOKEN, switch on notebook access, then run the cell. Without a token, set PUSH_TO_HUB = False in the save cell just above it and the model is only saved locally.

Path B: run it on your own computer

Works on a Mac with Apple Silicon (uses the GPU through MPS), on a PC with an NVIDIA GPU, or on plain CPU (slower, but it works).

# 1. Python environment (any Python 3.10+; uv is fastest but pip works too)
uv venv --python 3.13 && source .venv/bin/activate
uv pip install torch transformers datasets tokenizers huggingface_hub nbformat nbconvert ipykernel jupyterlab

# 2. Run the notebook end to end (about 4 minutes on an M-series Mac)
jupyter nbconvert --to notebook --execute tiny_llm_from_scratch.ipynb --output run.ipynb

# or open it interactively
jupyter lab tiny_llm_from_scratch.ipynb

To publish from your machine, either run hf auth login once, or export HF_TOKEN=hf_... in your shell before starting Jupyter. The notebook never prints the token.

Path C: skip training, just use the model

The trained model lives on the Hugging Face Hub as singhpratech/tiny-storyteller, a public repo, so the lines below work for anyone with no token. If you have published your own copy under your own name, swap in that repo id. The notebook's Part 10, Loading it back later, has the full loader. In short:

from huggingface_hub import hf_hub_download
import json, torch
# GPT class and tell_story() come from Part 6 and Part 8 of the notebook
cfg = json.load(open(hf_hub_download("singhpratech/tiny-storyteller", "config.json")))
state = torch.load(hf_hub_download("singhpratech/tiny-storyteller", "model.pt"), map_location="cpu")

The tokenizer is in the same repo (tokenizer.json), loadable with PreTrainedTokenizerFast.from_pretrained("singhpratech/tiny-storyteller").

Prompts that work well

The model only knows simple children's English. Start the way its stories start:

Once upon a time, a little monkey
One day, Momo went to the river
Ella the elephant was
Momo wanted to play with

Two knobs:

temperature in tell_story(): 0.8 is the default, 1.0 is adventurous, 0.5 is dull but safe.
top_k inside generate(): default 40, which keeps the wildest tokens out. Lower it for safer stories.

Prompts longer than 256 tokens are cropped to the last 256. An empty prompt starts a fresh story from the story marker.

Making it better

Change	Effect
`N_STORIES` 20,000 → 100,000	Noticeably better grammar and memory of characters
`MAX_STEPS` 1,500 → 5,000	Lower loss; about 4 minutes on a T4
`n_layer` 6 → 8, `n_embd` 256 → 384	Bigger model (about 16M parameters), slower, better
More Momo stories, or your own text	The model writes about what it reads

Watch perplexity in the training output. Random guessing is 4,096. Around 10 means real sentences. This recipe lands at about 10 on the hidden quiz stories (200 TinyStories the model never trains on, with no Momo stories in them, so the number is honest).

When something goes wrong

"No module named pip" on the first cell: harmless in some local setups. Install the packages with uv pip install ... or pip install ... and run the next cell.
Login code expired in the publish cell: approve the code within a minute, or use the HF_TOKEN secret described in Path A.
Out of memory on a small GPU: lower BATCH_SIZE in Part 7 to 16.
Stories are all about Lily: the Momo stories are only a few percent of the data. Raise the repeat factor in the Momo data cell from 3 to 6, or prompt with Momo's name.
Momo stories are recited word for word: that is memorising. Add more scenes to the sentence pools in momo_stories.py (more unique stories) and lower the repeat factor.

Rules of the road

The model is trained on public data (TinyStories, CDLA-Sharing-1.0) plus generated stories in this repo. Nothing personal is in it.
It is a toy. Do not use it for anything that matters.
The code is plain PyTorch on purpose, so you can read every line.