PLAYBOOK.md

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

1. Open https://colab.research.google.com and choose **Upload notebook**.
2. Pick `tiny_llm_from_scratch.ipynb` (download it from this repo first).
3. In Colab go to **Runtime > Change runtime type** and choose **T4 GPU**.
4. 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.
5. 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).

```bash
# 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:

```python
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.