tiny_stories_rl

Uses reinforcement learning to encourage roneneldan/TinyStories-33M to generate stories with alliteration

Docs are here

Installation

If you install uv, it'll get the dependencies.

Backup plan: ./build.sh and ./run.sh will build and run a Docker container that has uv, in case your system is weird (like my NixOS laptop) and doesn't work with uv. Once you're in the container, you can run the commands in the following sections.

Usage

uv run src/tiny_stories_rl/train.py

The KL penalty coefficient is configurable via --kl-coefficient; see here for more.

Running tests

uv run pytest tests