pip install -e ".[all]"

small-llm

A compact, self-contained Python toolkit for text preprocessing, Byte Pair Encoding tokenization, embeddings, and training a tiny transformer language model. Built for teaching — kept clean enough for research.

View on GitHub Quick start

The pipeline

raw text→ clean→ BPE tokenize→ embed→ train LM→ generate

Every stage is an importable module. Use the whole pipeline, or just the piece you need.

What's inside

preprocess

Strip HTML tags & entities, filter non-text Unicode while keeping letters, numbers, and punctuation.

tokenization

Character/word tokenizers plus a full BPE implementation — step-by-step helpers and a reusable BPEWordTokenizer with save/load.

embeddings

Cosine similarity, token lookup, nearest-neighbour search, and t-SNE / dimension plots. Pure NumPy + scikit-learn.

model

A small GPT-style transformer in Keras 3 (JAX/TF/PyTorch), with a live text-generation callback. No course dependency.

Quick start

# Train and reuse a subword tokenizer
from small_llm.tokenization import BPEWordTokenizer

tok = BPEWordTokenizer(["the quick brown fox", "the lazy dog"], num_merges=50)
ids = tok.encode("the fox")      # [12, 7, 33, ...]
tok.decode(ids)                  # "the fox"
tok.save("my_bpe.json")          # reuse across experiments

# Explore embeddings. (May need to get your own .npz file. Still to resolve) 
from small_llm.embeddings import load_embeddings, most_similar
emb, labels = load_embeddings("gemma_embeddings.npz")
most_similar("king", emb, labels, top_k=5)   # [('queen', 0.63), ...]

# Train the tiny language model
from small_llm.model import create_model, prepare_sequences
inputs, targets, max_len = prepare_sequences(encoded, tok.pad_token_id)
model = create_model(max_len, tok.vocabulary_size)
model.fit(inputs, targets, epochs=100, batch_size=2)

From notebooks to a package

The project started as six Google DeepMind AI Foundations lab notebooks. Each is preserved under Notebooks/ and reproduced as a runnable script under examples/:

  1. 01_preprocess.py — clean raw text
  2. 02_tokenize_basic.py — character/word tokenization & frequencies
  3. 03_learn_bpe.py — learn BPE merges step by step
  4. 04_train_tokenizer.py — train & save a BPE tokenizer
  5. 05_embeddings.py — similarity & nearest neighbours
  6. 06_train_slm.py — the full training pipeline

Built for research

Your corpus

Pass any list[str], or load plain text with data.load_text_lines().

Tunable model

embed_dim, num_heads, ff_dim, num_blocks, dropout, learning rate.

Any backend

Set KERAS_BACKEND to jax, tensorflow, or torch.

Tested

A pytest suite covers preprocessing, tokenization, and embeddings.