shalev396/tiny-shakespeare-chat

Tiny Shakespeare Chat

A small GPT written and trained from scratch (no pretrained weights, no tokenizer library). It is a character-level, nanoGPT-style decoder with 10,751,232 parameters. It was first pretrained on the Tiny Shakespeare corpus to predict the next character, then chat-tuned on 7,096 consecutive dialogue-line pairs, so it answers a message with the "next line" of an imaginary play. Its replies are play-style verse, not facts.

Model

Architecture	decoder-only transformer (model.GPT): token + learned position embeddings, 6 pre-LayerNorm blocks (6-head causal self-attention via scaled_dot_product_attention, 4x GELU MLP, dropout 0.2), final LayerNorm, LM head tied to the token embedding. Linear layers have no biases. GPT-2 init with scaled residual projections
Parameters	10,751,232 (the tied embedding / LM-head matrix counted once)
Context	256 characters
Tokenizer	model.CharTokenizer: the 65 characters of the corpus plus <|user|>, <|bot|>, <|end|> as single reserved ids (68 ids). Characters outside the vocabulary are dropped
Prompt	<|user|> {message} <|end|> <|bot|>, preceded by up to 3 earlier (user, bot) turns, cropped to the last 240 tokens
Sampling	temperature 0.8, top-k 40, stops at <|end|> or after max_new_tokens (200) characters
Output	the reply as a plain string (predict) or character by character (stream)
Files	model.safetensors + config.json (vocabulary and sizes, written by PyTorchModelHubMixin), model.py (tokenizer, architecture, sampler, load(), Predictor), handler.py + requirements.txt (Inference Endpoint), metrics.json, assets/ (experiment graphs)

Usage

Python

from huggingface_hub import snapshot_download
import sys
path = snapshot_download("shalev396/tiny-shakespeare-chat")
sys.path.insert(0, path)
import model
predictor = model.load(path, device="cpu")   # or "cuda"
print(predictor.predict("How fares the king?"))
print(predictor.predict("And the queen?", history=[["How fares the king?", "He is well, my lord."]],
                        max_new_tokens=150, temperature=0.7, top_k=40, seed=0))
for piece in predictor.stream("Speak to me of love"):   # streaming, one character at a time
    print(piece, end="", flush=True)

history accepts [[user, bot], ...] pairs or [{"role": ..., "content": ...}, ...] messages. Sampling is random; pass seed for a repeatable reply. Requirements: torch, huggingface_hub, safetensors.

Space API (free): the Space exposes /predict (message, history as a JSON string, temperature, max_new_tokens in; [reply, seconds, device] out). curl and @gradio/client examples are in its README.

Inference Endpoint: deploy this repo from its page (Deploy -> Inference Endpoints, CPU or GPU). handler.py loads the model once (on cuda when the endpoint has a GPU).

curl $ENDPOINT_URL -H "Authorization: Bearer $HF_TOKEN" -H "Content-Type: application/json" \
  -d '{"inputs": {"message": "And the queen?", "history": [["How fares the king?", "He is well, my lord."]]},
       "parameters": {"max_new_tokens": 150, "temperature": 0.7, "top_k": 40, "seed": 1}}'
# [{"generated_text": "..."}]

A plain string ({"inputs": "How fares the king?"}) works too.

Training

• Data: Tiny Shakespeare (karpathy/char-rnn input.txt, 1,115,394 characters of Shakespeare plays).
• Stage A, pretrain: next-character prediction on the raw play, first 90 % train / last 10 % val. 5,000 iterations, batch 64 x 256 characters, AdamW (betas 0.9/0.95, weight decay 0.1 on matrices only), 100 warmup iterations, cosine LR 3e-4 -> 3e-5, gradient clipping 1.0, fp16 autocast on CUDA.
• Stage B, chat-tune: the play is parsed into consecutive (speaker, utterance) turns. Every pair of neighbouring turns becomes <|user|> {line} <|end|> <|bot|> {next line} <|end|>, plus a variant whose reply starts with SPEAKER:. The samples are shuffled (seed 42) and joined, and the last 5 % of that stream is val. 2,000 iterations from the stage A weights, same optimiser, 50 warmup iterations, cosine LR 1e-4 -> 1e-5.
• This checkpoint: converted into model.GPT from the earlier full run (legacy project 13_tiny_shakespeare_chat_pt, one CUDA GPU, 541 s for both stages, 2026-08-01). Every tensor loads with strict=True, the logits match the legacy code exactly (max |diff| 0.0) and seeded generations are identical. All metrics below were re-computed locally on CPU with this repo's training/src/engine.evaluate.

Full code: training/ · Colab.

Experiments

Every variant scored with the same deterministic evaluation (non-overlapping 256-character windows, every validation character scored once) on both validation splits (from metrics.json -> comparison). Cross-entropy in nats per character, lower is better. The deployed model is in bold.

variant	chat val loss	text val loss	notes
bigram baseline (character counts, add-one smoothing)	2.451	2.482	next char from the current char only
stage A only: pretrained GPT, 5,000 iters	1.563	1.495	the clean held-out number: it never saw text val
stage A + B: chat-tuned GPT (deployed)	1.054	1.296	text val is not held out after stage B (see Limitations)

• Pretraining alone takes the loss on unseen text from 2.48 (bigram) to 1.50 nats per character (2.16 bits per character).
• Chat-tuning drops the chat val loss from 1.56 to 1.05. Part of that is the model learning the chat format (markers, speaker names, where a reply ends). Part is the leak described in Limitations, so 1.05 is optimistic.

Training curves of the deployed run (the run's own estimates on 40 random 64 x 256 batches per split). Stage A's val loss flattens around 1.50 after about 3,000 iterations while train keeps falling (1.02 at the end), a mild overfit. Stage B's val loss is still creeping down at 2,000 iterations (1.054).

Evaluation

metric (val)	value
chat_val_loss	1.0540
chat_val_bpc	1.5206
chat_val_ppl	2.8691
text_val_loss	1.2958
text_val_bpc	1.8695

The deployed model on the chat val split (215,040 characters scored) and the text val split (111,360 characters). bpc = bits per character (loss / ln 2), ppl = per-character perplexity (e^loss).

Loss by position in the context window: the first characters of a window are hard to predict (about 2.4 nats with no context), and the loss levels off after roughly 30-50 characters. The model mostly uses about a line of context.

Limitations

• Toy model. 10.75M parameters trained on 1.1 MB of text. It imitates the style (speaker names, verse, archaic words) but invents words, loses the thread within a few lines and does not understand the message. Nothing it says is factual.
• Chat val leaks. The chat samples come from the whole play, and both variants of each pair are shuffled together before the split, so 95.8 % of chat-val samples have their twin (same dialogue pair, other variant) in chat-train. The chat-tuned model has also seen the stage A validation text inside chat pairs. So chat val loss is optimistic, and text val loss after stage B is not a held-out number. The clean held-out figure is stage A's 1.495 on text val. The split was kept so that the code reproduces these weights.
• 68 symbols. Characters outside the corpus alphabet (digits other than 3, brackets, double quotes, accents, emoji) are silently dropped from the message.
• 256-character context. Only the last 3 turns are kept, and long prompts are cropped from the left.
• Random sampling. The same message gives a different reply each time unless seed is set. A reply can be empty (the model may emit <|end|> at once) or cut off at max_new_tokens.
• The plays contain archaic, violent or offensive language, and the model can reproduce it.