PocketNarrator: Efficient Story Generation with Small Language Models

PocketNarrator was a research project for the "Efficient Methods in Machine Learning" course (Master Project, WS25/26) at the University of Hamburg. It focuses on building and evaluating small language models for narrative generation using the TinyStories dataset.

Table of Contents

• Overview
• Sample Output
• Features
• Supported Models
• Directory Structure
• Requirements
• Installation
  • 1. Clone the Repository
  • 2. Set Up Environment
  • 3. Install Python Dependencies
• Usage
  • Training
  • Evaluation
  • Generation
  • Demo App
• Resources
• Team

Overview

PocketNarrator is a systematic investigation into the architecture and components of small language models for efficient narrative generation. Our goal was to understand trade-offs between different architectural choices in terms of performance, computational efficiency, and output quality.

The project implements multiple model architectures from scratch using PyTorch, trained on the TinyStories dataset—a clean, restricted-domain collection of children's stories ideal for training efficient models. Our models are designed for next-token prediction and story continuation tasks.

Sample Output

Here is an example story generated by our 28M parameter Transformer model based on the prompt "Once upon a time":

Once upon a time, in a small house, there was a pretty vase. The vase was on a table. The vase was red and shiny. Everyone loved to look at it. One day, a naughty boy named Tim came to the house. Tim saw the pretty vase on the table. He wanted to poke the vase. So, he poked the vase with his finger. The vase fell down and broke. Tim was scared. His mom came into the room and saw the broken vase. She was not happy. Tim said sorry to his mom. Tim's mom hugged him and told him not to poke things. From that day on, Tim learned his lesson. He never poked anything again. And they all lived happily ever after.

Features

• Multiple model architectures supporting N-gram and Transformer models
• Comprehensive evaluation metrics including BLEU, ROUGE, perplexity, distinct-n, LLM-as-a-judge, text quality, and noun carryover analysis
• Flexible tokenization with BPE and character-level tokenizers
• W&B integration for experiment tracking and visualization
• Production-ready evaluation pipeline with model comparison and dataset analysis tools
• Clean, modular codebase with abstract base classes for extensibility

Supported Models

• N-gram Model: Lightweight baseline model for quick experiments
• Transformer Model: Custom decoder-only transformer architecture with configuration options
  • Positional encoding: Sinusoidal & RoPE (Rotary Positional Encoding)
  • Activation function: GELU & SwiGLU
  • Attention mechanism: Softmax & Linear (following Katharopoulos et al.)
  • Optimization & efficiency techniques: Automatic Mixed Precision (AMP), LR scheduling, Gradient Clipping, Weight Decay, Dropout, KV-Caching

Directory Structure

pocket-narrator/
├── README.md                            # Project documentation
├── requirements.txt                     # Python dependencies
├── pyproject.toml                       # Project configuration
├── pytest.ini                           # Pytest configuration
│
├── configs/                             # Configuration files for models and training
│   ├── base_config.yaml
│   ├── evaluation/                      # Evaluation configs
│   ├── large_models_grid_search/        # Grid search configs
│   └── parameter_comparison_runs/       # Parameter comparison configs
│
├── data/                                # Datasets (raw and processed)
│   ├── raw/                             # Original dataset files
│   └── processed/                       # Processed datasets
│
├── models/                              # Trained model checkpoints
│   ├── ngram/
│   ├── transformer/
│   └── cool_models/
│
├── notebooks/                           # Jupyter notebooks for exploration
│
├── pocket_narrator/                     # Main package source code
│   ├── __init__.py
│   ├── models/                          # Model architectures
│   │   ├── base_model.py                # Abstract base class
│   │   ├── ngram_model.py               # N-gram implementation
│   │   ├── components/                  # Reusable components
│   │   │   ├── positional_encoding.py
│   │   │   └── base_pos_encoding.py
│   │   └── transformers/                # Transformer architecture
│   │       ├── model.py
│   │       ├── transformer_block.py
│   │       ├── attention.py
│   │       └── base_attention.py
│   │
│   ├── tokenizers/                      # Tokenization implementations
│   │   ├── base_tokenizer.py
│   │   ├── bpe_tokenizer.py
│   │   └── character_tokenizer.py
│   │
│   ├── trainers/                        # Model trainers
│   │   ├── base_trainer.py
│   │   ├── ngram_trainer.py
│   │   └── transformer_trainer.py
│   │
│   ├── data_loader.py                   # Data loading utilities
│   ├── evaluate.py                      # Evaluation metrics
│   ├── text_quality.py                  # Text quality evaluation
│   ├── noun_carryover.py                # Noun carryover metrics
│   └── gemini_api.py                    # LLM-based evaluation
│
├── scripts/                             # Standalone execution scripts
│   ├── train.py                         # Main training script
│   ├── generate.py                      # Text generation
│   ├── preprocess.py                    # Data preprocessing (filtering, splits)
│   ├── fetch_tinystories.py             # Dataset download
│   ├── evaluate.py                      # Multi-model comparison with W&B logging
│   ├── evaluate_model.py                # Single model evaluation
│   ├── evaluate_dataset.py              # Dataset quality & diversity analysis
│   ├── evaluate_dataset_comprehensive.py # Comprehensive dataset evaluation
│   ├── evaluate_llm_agent_only.py       # seperate LLM-as-a-judge only evaluation
│   ├── inference_benchmark.py           # Inference speed & memory benchmarking (KV-Caching)
│   ├── tokenizer_compare_tokenization_efficiency.py  # Tokenizer efficiency comparison
│   └── demo_app.py                      # Interactive Gradio demo
│
├── tests/                               # Unit and integration tests
│
├── tokenizers/                          # Saved tokenizer artifacts
│
└── wandb/                               # W&B experiment tracking

Requirements

• Python 3.10+
• PyTorch 2.0+
• NumPy, Pandas
• Hugging Face Transformers & Datasets
• gradio (for the demo app)
• wandb (for experiment tracking)
• Optional: spacy, sentence-transformers, google-genai (for advanced evaluation)

Installation

1. Clone the Repository

git clone https://github.com/bschink/pocket-narrator.git
cd pocket-narrator

2. Set Up Environment

# Create conda environment
conda create -n pocket-narrator

# Activate the conda environment
conda activate pocket-narrator

3. Install Python Dependencies

pip install -r requirements.txt

Usage

Training

Before training, you need to download and preprocess the dataset.

Step 1 — Download TinyStories

# Downloads the dataset to data/raw/TinyStories/ (prompts for confirmation)
python scripts/fetch_tinystories.py

# Or skip the prompt and download automatically
python scripts/fetch_tinystories.py --yes

# Custom destination
python scripts/fetch_tinystories.py --yes --dest data/raw/TinyStories

Step 2 — Preprocess the dataset

# Preprocess the training split (removes non-English stories, duplicates, etc.)
# Produces .clean.txt, .clean.half.txt, .clean.quarted.txt, .clean.eighth.txt,
# a .clean.deleted.txt log and a .clean.manifest.json summary
python scripts/preprocess.py \
    --input data/raw/TinyStories/TinyStoriesV2-GPT4-train.txt \
    --output data/processed/TinyStories/

# Key arguments:
#   --min-letters: minimum ASCII letters per story (default: 100)
#   --wordfreq-threshold: zipf frequency cutoff for English detection (default: 2.5)
#   --no-splits: disable half/quarter/eighth subset generation
#   --keep-invalid: keep stories that would otherwise be removed

Step 3 — Train a model

# Train with default config
python scripts/train.py

# Train with specific config
python scripts/train.py --config_path configs/large_models_grid_search/transformer_28M_higher_lr_full_dataset.yaml

# Key arguments:
#   --config_path: Path to training config YAML
#   --model_type: ngram or transformer
#   --dataset_path: Path to dataset
#   --output_dir: Directory to save models
#   --epochs: Number of training epochs
#   --batch_size: Batch size for training

Evaluation

Evaluate a trained model on a dataset:

# Evaluate single model with comprehensive metrics
python scripts/evaluate_model.py \
    --model_path models/transformer/transformer_model.pth \
    --model_type transformer \
    --dataset_path data/test_dataset.txt

# Evaluate dataset without a model (text quality, distinct-n, etc.)
python scripts/evaluate_dataset_comprehensive.py \
    --dataset_path data/validation.txt \
    --dataset_name "TinyStories Validation"

Generation

Generate text continuations with a trained model:

# Generate with default prompt
python scripts/generate.py \
    --model_path models/transformer/transformer_model.pth \
    --model_type transformer

# Generate with custom prompt
python scripts/generate.py \
    --model_path models/transformer/transformer_model.pth \
    --model_type transformer \
    --prompt "A girl went to the" \
    --max_length 100 \
    --temperature 0.7

Demo App

An interactive Gradio demo app is available for story generation with a simple UI:

python scripts/demo_app.py

Before running, configure the parameters at the top of scripts/demo_app.py (line 13 onwards) to point to your model and tokenizer:

MODEL_PATH = "models/transformer/<your_model>.model"  # Path to trained model
GENERATION_STRATEGY = "sample"                        # Generation strategy (e.g. "sample", "greedy")
TOKENIZER_TYPE = "bpe"                                # Tokenizer type ("bpe" or "character")
TOKENIZER_PATH = "tokenizers/<your_tokenizer>"        # Path to tokenizer directory
MAX_NEW_TOKENS = 5000                                 # Maximum number of tokens to generate

Once running, open the local URL printed in the terminal to access the app. Enter a prompt, adjust the temperature slider, and click generate to produce a story continuation.

Resources

Resource	Link
W&B Training Runs	once-upon-a-prompt/pocket-narrator
W&B Evaluation Runs	once-upon-a-prompt/pocket-narrator-eval
Trained Models & Tokenizers	Google Drive

Team

Benedikt Schink & Asiya Yumna